From 808f3874b35d4edf3ee5e046af8df2393a9ab890 Mon Sep 17 00:00:00 2001 From: Timothy Zelinsky Date: Sun, 22 Mar 2026 14:36:32 +1100 Subject: [PATCH 1/4] Add V0 memory layer --- .../CodexKit/Memory/InMemoryMemoryStore.swift | 157 +++++ Sources/CodexKit/Memory/MemoryModels.swift | 374 +++++++++++ .../CodexKit/Memory/MemoryQueryEngine.swift | 280 ++++++++ Sources/CodexKit/Memory/MemoryStore.swift | 24 + .../CodexKit/Memory/SQLiteMemoryStore.swift | 626 ++++++++++++++++++ Sources/CodexKit/Runtime/AgentModels.swift | 12 +- Sources/CodexKit/Runtime/AgentRuntime.swift | 223 ++++++- Tests/CodexKitTests/AgentRuntimeTests.swift | 241 +++++++ Tests/CodexKitTests/MemoryStoreTests.swift | 270 ++++++++ 9 files changed, 2200 insertions(+), 7 deletions(-) create mode 100644 Sources/CodexKit/Memory/InMemoryMemoryStore.swift create mode 100644 Sources/CodexKit/Memory/MemoryModels.swift create mode 100644 Sources/CodexKit/Memory/MemoryQueryEngine.swift create mode 100644 Sources/CodexKit/Memory/MemoryStore.swift create mode 100644 Sources/CodexKit/Memory/SQLiteMemoryStore.swift create mode 100644 Tests/CodexKitTests/MemoryStoreTests.swift diff --git a/Sources/CodexKit/Memory/InMemoryMemoryStore.swift b/Sources/CodexKit/Memory/InMemoryMemoryStore.swift new file mode 100644 index 0000000..ad605f6 --- /dev/null +++ b/Sources/CodexKit/Memory/InMemoryMemoryStore.swift @@ -0,0 +1,157 @@ +import Foundation + +public actor InMemoryMemoryStore: MemoryStoring { + private var recordsByNamespace: [String: [String: MemoryRecord]] + + public init(initialRecords: [MemoryRecord] = []) { + recordsByNamespace = Dictionary(grouping: initialRecords, by: \.namespace) + .mapValues { records in + Dictionary(uniqueKeysWithValues: records.map { ($0.id, $0) }) + } + } + + public func put(_ record: MemoryRecord) async throws { + try MemoryQueryEngine.validateNamespace(record.namespace) + var namespaceRecords = recordsByNamespace[record.namespace, default: [:]] + guard namespaceRecords[record.id] == nil else { + throw MemoryStoreError.duplicateRecordID(record.id) + } + + if let dedupeKey = record.dedupeKey, + namespaceRecords.values.contains(where: { $0.dedupeKey == dedupeKey }) { + throw MemoryStoreError.duplicateDedupeKey(dedupeKey) + } + + namespaceRecords[record.id] = record + recordsByNamespace[record.namespace] = namespaceRecords + } + + public func putMany(_ records: [MemoryRecord]) async throws { + var working = recordsByNamespace + + for record in records { + try MemoryQueryEngine.validateNamespace(record.namespace) + var namespaceRecords = working[record.namespace, default: [:]] + guard namespaceRecords[record.id] == nil else { + throw MemoryStoreError.duplicateRecordID(record.id) + } + if let dedupeKey = record.dedupeKey, + namespaceRecords.values.contains(where: { $0.dedupeKey == dedupeKey }) { + throw MemoryStoreError.duplicateDedupeKey(dedupeKey) + } + namespaceRecords[record.id] = record + working[record.namespace] = namespaceRecords + } + + recordsByNamespace = working + } + + public func upsert(_ record: MemoryRecord, dedupeKey: String) async throws { + try MemoryQueryEngine.validateNamespace(record.namespace) + var namespaceRecords = recordsByNamespace[record.namespace, default: [:]] + + if let existing = namespaceRecords.values.first(where: { $0.dedupeKey == dedupeKey }) { + namespaceRecords.removeValue(forKey: existing.id) + } else if let existingByID = namespaceRecords[record.id], + existingByID.dedupeKey != nil, + existingByID.dedupeKey != dedupeKey { + namespaceRecords.removeValue(forKey: existingByID.id) + } + + var updatedRecord = record + updatedRecord.dedupeKey = dedupeKey + namespaceRecords[updatedRecord.id] = updatedRecord + recordsByNamespace[record.namespace] = namespaceRecords + } + + public func query(_ query: MemoryQuery) async throws -> MemoryQueryResult { + try MemoryQueryEngine.validateNamespace(query.namespace) + let namespaceRecords = recordsByNamespace[query.namespace, default: [:]] + let candidates = namespaceRecords.values.map { record in + MemoryQueryEngine.Candidate( + record: record, + rawTextScore: MemoryQueryEngine.defaultTextScore( + for: record, + queryText: query.text + ) + ) + } + + return try MemoryQueryEngine.evaluate( + candidates: candidates, + query: query + ) + } + + public func compact(_ request: MemoryCompactionRequest) async throws { + try MemoryQueryEngine.validateNamespace(request.replacement.namespace) + var working = recordsByNamespace + let namespace = request.replacement.namespace + var namespaceRecords = working[namespace, default: [:]] + + guard namespaceRecords[request.replacement.id] == nil else { + throw MemoryStoreError.duplicateRecordID(request.replacement.id) + } + if let dedupeKey = request.replacement.dedupeKey, + namespaceRecords.values.contains(where: { $0.dedupeKey == dedupeKey }) { + throw MemoryStoreError.duplicateDedupeKey(dedupeKey) + } + + namespaceRecords[request.replacement.id] = request.replacement + for sourceID in request.sourceIDs { + guard var existing = namespaceRecords[sourceID] else { + continue + } + existing.status = .archived + namespaceRecords[sourceID] = existing + } + + working[namespace] = namespaceRecords + recordsByNamespace = working + } + + public func archive(ids: [String], namespace: String) async throws { + try MemoryQueryEngine.validateNamespace(namespace) + var namespaceRecords = recordsByNamespace[namespace, default: [:]] + for id in ids { + guard var record = namespaceRecords[id] else { + continue + } + record.status = .archived + namespaceRecords[id] = record + } + recordsByNamespace[namespace] = namespaceRecords + } + + public func delete(ids: [String], namespace: String) async throws { + try MemoryQueryEngine.validateNamespace(namespace) + var namespaceRecords = recordsByNamespace[namespace, default: [:]] + for id in ids { + namespaceRecords.removeValue(forKey: id) + } + recordsByNamespace[namespace] = namespaceRecords + } + + @discardableResult + public func pruneExpired( + now: Date, + namespace: String + ) async throws -> Int { + try MemoryQueryEngine.validateNamespace(namespace) + var namespaceRecords = recordsByNamespace[namespace, default: [:]] + let expiredIDs = namespaceRecords.values + .filter { record in + !record.isPinned && + record.status == .active && + (record.expiresAt?.compare(now) == .orderedAscending || + record.expiresAt?.compare(now) == .orderedSame) + } + .map(\.id) + + for id in expiredIDs { + namespaceRecords.removeValue(forKey: id) + } + recordsByNamespace[namespace] = namespaceRecords + return expiredIDs.count + } +} diff --git a/Sources/CodexKit/Memory/MemoryModels.swift b/Sources/CodexKit/Memory/MemoryModels.swift new file mode 100644 index 0000000..0bf25ba --- /dev/null +++ b/Sources/CodexKit/Memory/MemoryModels.swift @@ -0,0 +1,374 @@ +import Foundation + +public enum MemoryStoreError: Error, LocalizedError, Equatable, Sendable { + case invalidNamespace + case duplicateRecordID(String) + case duplicateDedupeKey(String) + + public var errorDescription: String? { + switch self { + case .invalidNamespace: + return "Memory namespace must not be empty." + case let .duplicateRecordID(id): + return "A memory record with id \(id) already exists." + case let .duplicateDedupeKey(key): + return "A memory record with dedupe key \(key) already exists." + } + } +} + +public struct MemoryScope: RawRepresentable, Codable, Hashable, Sendable, ExpressibleByStringLiteral { + public var rawValue: String + + public init(rawValue: String) { + self.rawValue = rawValue + } + + public init(stringLiteral value: StringLiteralType) { + self.init(rawValue: value) + } +} + +public enum MemoryRecordStatus: String, Codable, Hashable, Sendable { + case active + case archived +} + +public struct MemoryRecord: Identifiable, Codable, Hashable, Sendable { + public var id: String + public var namespace: String + public var scope: MemoryScope + public var kind: String + public var summary: String + public var evidence: [String] + public var importance: Double + public var createdAt: Date + public var observedAt: Date? + public var expiresAt: Date? + public var tags: [String] + public var relatedIDs: [String] + public var dedupeKey: String? + public var isPinned: Bool + public var attributes: JSONValue? + public var status: MemoryRecordStatus + + public init( + id: String = UUID().uuidString, + namespace: String, + scope: MemoryScope, + kind: String, + summary: String, + evidence: [String] = [], + importance: Double = 0, + createdAt: Date = Date(), + observedAt: Date? = nil, + expiresAt: Date? = nil, + tags: [String] = [], + relatedIDs: [String] = [], + dedupeKey: String? = nil, + isPinned: Bool = false, + attributes: JSONValue? = nil, + status: MemoryRecordStatus = .active + ) { + self.id = id + self.namespace = namespace + self.scope = scope + self.kind = kind + self.summary = summary + self.evidence = evidence + self.importance = importance + self.createdAt = createdAt + self.observedAt = observedAt + self.expiresAt = expiresAt + self.tags = tags + self.relatedIDs = relatedIDs + self.dedupeKey = dedupeKey + self.isPinned = isPinned + self.attributes = attributes + self.status = status + } + + public var effectiveDate: Date { + observedAt ?? createdAt + } +} + +public struct MemoryRankingWeights: Codable, Hashable, Sendable { + public var textWeight: Double + public var importanceWeight: Double + public var recencyWeight: Double + public var kindBoost: Double + public var tagBoost: Double + public var relatedIDBoost: Double + + public init( + textWeight: Double, + importanceWeight: Double, + recencyWeight: Double, + kindBoost: Double, + tagBoost: Double, + relatedIDBoost: Double + ) { + self.textWeight = textWeight + self.importanceWeight = importanceWeight + self.recencyWeight = recencyWeight + self.kindBoost = kindBoost + self.tagBoost = tagBoost + self.relatedIDBoost = relatedIDBoost + } + + public static let `default` = MemoryRankingWeights( + textWeight: 0.50, + importanceWeight: 0.25, + recencyWeight: 0.15, + kindBoost: 0.05, + tagBoost: 0.03, + relatedIDBoost: 0.02 + ) +} + +public struct MemoryReadBudget: Codable, Hashable, Sendable { + public var maxItems: Int + public var maxCharacters: Int + + public init( + maxItems: Int, + maxCharacters: Int + ) { + self.maxItems = maxItems + self.maxCharacters = maxCharacters + } + + public static let runtimeDefault = MemoryReadBudget( + maxItems: 8, + maxCharacters: 1600 + ) +} + +public struct MemoryQuery: Codable, Hashable, Sendable { + public var namespace: String + public var scopes: [MemoryScope] + public var text: String? + public var kinds: [String] + public var tags: [String] + public var relatedIDs: [String] + public var recencyWindow: TimeInterval? + public var minImportance: Double? + public var ranking: MemoryRankingWeights + public var limit: Int + public var maxCharacters: Int + public var includeArchived: Bool + + public init( + namespace: String, + scopes: [MemoryScope] = [], + text: String? = nil, + kinds: [String] = [], + tags: [String] = [], + relatedIDs: [String] = [], + recencyWindow: TimeInterval? = nil, + minImportance: Double? = nil, + ranking: MemoryRankingWeights = .default, + limit: Int = MemoryReadBudget.runtimeDefault.maxItems, + maxCharacters: Int = MemoryReadBudget.runtimeDefault.maxCharacters, + includeArchived: Bool = false + ) { + self.namespace = namespace + self.scopes = scopes + self.text = text + self.kinds = kinds + self.tags = tags + self.relatedIDs = relatedIDs + self.recencyWindow = recencyWindow + self.minImportance = minImportance + self.ranking = ranking + self.limit = limit + self.maxCharacters = maxCharacters + self.includeArchived = includeArchived + } +} + +public struct MemoryMatchExplanation: Codable, Hashable, Sendable { + public var totalScore: Double + public var textScore: Double + public var recencyScore: Double + public var importanceScore: Double + public var kindBoost: Double + public var tagBoost: Double + public var relatedIDBoost: Double + + public init( + totalScore: Double, + textScore: Double, + recencyScore: Double, + importanceScore: Double, + kindBoost: Double, + tagBoost: Double, + relatedIDBoost: Double + ) { + self.totalScore = totalScore + self.textScore = textScore + self.recencyScore = recencyScore + self.importanceScore = importanceScore + self.kindBoost = kindBoost + self.tagBoost = tagBoost + self.relatedIDBoost = relatedIDBoost + } +} + +public struct MemoryQueryMatch: Codable, Hashable, Sendable { + public var record: MemoryRecord + public var explanation: MemoryMatchExplanation + + public init( + record: MemoryRecord, + explanation: MemoryMatchExplanation + ) { + self.record = record + self.explanation = explanation + } +} + +public struct MemoryQueryResult: Codable, Hashable, Sendable { + public var matches: [MemoryQueryMatch] + public var truncated: Bool + + public init( + matches: [MemoryQueryMatch], + truncated: Bool + ) { + self.matches = matches + self.truncated = truncated + } +} + +public struct MemoryCompactionRequest: Codable, Hashable, Sendable { + public var replacement: MemoryRecord + public var sourceIDs: [String] + + public init( + replacement: MemoryRecord, + sourceIDs: [String] + ) { + self.replacement = replacement + self.sourceIDs = sourceIDs + } +} + +public struct AgentMemoryContext: Codable, Hashable, Sendable { + public var namespace: String + public var scopes: [MemoryScope] + public var kinds: [String] + public var tags: [String] + public var relatedIDs: [String] + public var recencyWindow: TimeInterval? + public var minImportance: Double? + public var ranking: MemoryRankingWeights? + public var readBudget: MemoryReadBudget? + + public init( + namespace: String, + scopes: [MemoryScope] = [], + kinds: [String] = [], + tags: [String] = [], + relatedIDs: [String] = [], + recencyWindow: TimeInterval? = nil, + minImportance: Double? = nil, + ranking: MemoryRankingWeights? = nil, + readBudget: MemoryReadBudget? = nil + ) { + self.namespace = namespace + self.scopes = scopes + self.kinds = kinds + self.tags = tags + self.relatedIDs = relatedIDs + self.recencyWindow = recencyWindow + self.minImportance = minImportance + self.ranking = ranking + self.readBudget = readBudget + } +} + +public enum MemorySelectionMode: String, Codable, Hashable, Sendable { + case inherit + case append + case replace + case disable +} + +public struct MemorySelection: Codable, Hashable, Sendable { + public var mode: MemorySelectionMode + public var namespace: String? + public var scopes: [MemoryScope] + public var kinds: [String] + public var tags: [String] + public var relatedIDs: [String] + public var recencyWindow: TimeInterval? + public var minImportance: Double? + public var ranking: MemoryRankingWeights? + public var readBudget: MemoryReadBudget? + public var text: String? + + public init( + mode: MemorySelectionMode = .inherit, + namespace: String? = nil, + scopes: [MemoryScope] = [], + kinds: [String] = [], + tags: [String] = [], + relatedIDs: [String] = [], + recencyWindow: TimeInterval? = nil, + minImportance: Double? = nil, + ranking: MemoryRankingWeights? = nil, + readBudget: MemoryReadBudget? = nil, + text: String? = nil + ) { + self.mode = mode + self.namespace = namespace + self.scopes = scopes + self.kinds = kinds + self.tags = tags + self.relatedIDs = relatedIDs + self.recencyWindow = recencyWindow + self.minImportance = minImportance + self.ranking = ranking + self.readBudget = readBudget + self.text = text + } +} + +public protocol MemoryPromptRendering: Sendable { + func render(result: MemoryQueryResult, budget: MemoryReadBudget) -> String +} + +public struct DefaultMemoryPromptRenderer: MemoryPromptRendering, Sendable { + public init() {} + + public func render( + result: MemoryQueryResult, + budget: MemoryReadBudget + ) -> String { + MemoryQueryEngine.renderPrompt( + matches: result.matches, + budget: budget + ) + } +} + +public struct AgentMemoryConfiguration: Sendable { + public let store: any MemoryStoring + public let defaultRanking: MemoryRankingWeights + public let defaultReadBudget: MemoryReadBudget + public let promptRenderer: any MemoryPromptRendering + + public init( + store: any MemoryStoring, + defaultRanking: MemoryRankingWeights = .default, + defaultReadBudget: MemoryReadBudget = .runtimeDefault, + promptRenderer: any MemoryPromptRendering = DefaultMemoryPromptRenderer() + ) { + self.store = store + self.defaultRanking = defaultRanking + self.defaultReadBudget = defaultReadBudget + self.promptRenderer = promptRenderer + } +} diff --git a/Sources/CodexKit/Memory/MemoryQueryEngine.swift b/Sources/CodexKit/Memory/MemoryQueryEngine.swift new file mode 100644 index 0000000..47a0509 --- /dev/null +++ b/Sources/CodexKit/Memory/MemoryQueryEngine.swift @@ -0,0 +1,280 @@ +import Foundation + +internal enum MemoryQueryEngine { + internal struct Candidate { + let record: MemoryRecord + let rawTextScore: Double? + } + + private struct ScoredCandidate { + let match: MemoryQueryMatch + let characterCost: Int + } + + static func evaluate( + candidates: [Candidate], + query: MemoryQuery, + now: Date = Date() + ) throws -> MemoryQueryResult { + try validateNamespace(query.namespace) + + let activeCandidates = candidates.filter { candidate in + matchesFilters(candidate.record, query: query, now: now) + } + + let textScores = normalizedTextScores(from: activeCandidates) + + let scored = activeCandidates.map { candidate -> ScoredCandidate in + let textScore = textScores[candidate.record.id] ?? 0 + let recencyScore = recencyScore( + for: candidate.record, + query: query, + now: now + ) + let importanceScore = clamp(candidate.record.importance) + let kindBoost = query.kinds.contains(candidate.record.kind) ? query.ranking.kindBoost : 0 + let tagBoost = candidate.record.tags.contains(where: query.tags.contains) ? query.ranking.tagBoost : 0 + let relatedBoost = candidate.record.relatedIDs.contains(where: query.relatedIDs.contains) ? query.ranking.relatedIDBoost : 0 + let totalScore = + (textScore * query.ranking.textWeight) + + (importanceScore * query.ranking.importanceWeight) + + (recencyScore * query.ranking.recencyWeight) + + kindBoost + + tagBoost + + relatedBoost + + let explanation = MemoryMatchExplanation( + totalScore: totalScore, + textScore: textScore, + recencyScore: recencyScore, + importanceScore: importanceScore, + kindBoost: kindBoost, + tagBoost: tagBoost, + relatedIDBoost: relatedBoost + ) + let match = MemoryQueryMatch( + record: candidate.record, + explanation: explanation + ) + + return ScoredCandidate( + match: match, + characterCost: renderMatch(match).count + ) + } + .sorted { + if $0.match.explanation.totalScore == $1.match.explanation.totalScore { + if $0.match.record.effectiveDate == $1.match.record.effectiveDate { + return $0.match.record.id < $1.match.record.id + } + return $0.match.record.effectiveDate > $1.match.record.effectiveDate + } + return $0.match.explanation.totalScore > $1.match.explanation.totalScore + } + + var selected: [MemoryQueryMatch] = [] + var characterCount = 0 + var truncated = false + + for candidate in scored { + if selected.count >= query.limit { + truncated = true + break + } + + let nextCount = characterCount + candidate.characterCost + if !selected.isEmpty, nextCount > query.maxCharacters { + truncated = true + break + } + + if selected.isEmpty, candidate.characterCost > query.maxCharacters { + truncated = true + break + } + + selected.append(candidate.match) + characterCount = nextCount + } + + if !truncated { + truncated = selected.count < scored.count + } + + return MemoryQueryResult( + matches: selected, + truncated: truncated + ) + } + + static func renderPrompt( + matches: [MemoryQueryMatch], + budget: MemoryReadBudget + ) -> String { + var lines: [String] = [] + var characterCount = 0 + + for match in matches.prefix(budget.maxItems) { + let rendered = renderMatch(match) + let nextCount = characterCount + rendered.count + (lines.isEmpty ? 0 : 1) + if !lines.isEmpty, nextCount > budget.maxCharacters { + break + } + if lines.isEmpty, rendered.count > budget.maxCharacters { + break + } + lines.append(rendered) + characterCount = nextCount + } + + guard !lines.isEmpty else { + return "" + } + + return """ + Relevant Memory: + \(lines.joined(separator: "\n")) + """ + } + + static func defaultTextScore( + for record: MemoryRecord, + queryText: String? + ) -> Double { + let queryTokens = tokenize(queryText) + guard !queryTokens.isEmpty else { + return 0 + } + + let haystack = tokenize( + ([record.summary] + record.evidence + record.tags + [record.kind]).joined(separator: " ") + ) + guard !haystack.isEmpty else { + return 0 + } + + let overlap = Set(queryTokens).intersection(Set(haystack)) + return Double(overlap.count) / Double(Set(queryTokens).count) + } + + static func validateNamespace(_ namespace: String) throws { + guard !namespace.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw MemoryStoreError.invalidNamespace + } + } + + static func tokenize(_ value: String?) -> [String] { + guard let value else { + return [] + } + + return value + .lowercased() + .split { !$0.isLetter && !$0.isNumber } + .map(String.init) + .filter { !$0.isEmpty } + } + + private static func matchesFilters( + _ record: MemoryRecord, + query: MemoryQuery, + now: Date + ) -> Bool { + guard record.namespace == query.namespace else { + return false + } + + if !query.includeArchived, record.status == .archived { + return false + } + + if !record.isPinned, + let expiresAt = record.expiresAt, + expiresAt <= now { + return false + } + + if !query.scopes.isEmpty, !query.scopes.contains(record.scope) { + return false + } + + if !query.kinds.isEmpty, !query.kinds.contains(record.kind) { + return false + } + + if !query.tags.isEmpty, !record.tags.contains(where: query.tags.contains) { + return false + } + + if !query.relatedIDs.isEmpty, !record.relatedIDs.contains(where: query.relatedIDs.contains) { + return false + } + + if let minImportance = query.minImportance, + clamp(record.importance) < minImportance { + return false + } + + if let recencyWindow = query.recencyWindow, + now.timeIntervalSince(record.effectiveDate) > recencyWindow { + return false + } + + return true + } + + private static func normalizedTextScores( + from candidates: [Candidate] + ) -> [String: Double] { + let rawScores = candidates.compactMap(\.rawTextScore) + guard let maxScore = rawScores.max(), + let minScore = rawScores.min() + else { + return [:] + } + + return candidates.reduce(into: [String: Double]()) { partial, candidate in + guard let rawScore = candidate.rawTextScore else { + partial[candidate.record.id] = 0 + return + } + + if maxScore == minScore { + partial[candidate.record.id] = 1 + } else { + partial[candidate.record.id] = clamp((maxScore - rawScore) / (maxScore - minScore)) + } + } + } + + private static func recencyScore( + for record: MemoryRecord, + query: MemoryQuery, + now: Date + ) -> Double { + let halfLife = max(query.recencyWindow ?? (30 * 24 * 60 * 60), 1) + let age = max(now.timeIntervalSince(record.effectiveDate), 0) + return clamp(pow(0.5, age / halfLife)) + } + + private static func renderMatch(_ match: MemoryQueryMatch) -> String { + var components: [String] = [ + "- [\(match.record.scope.rawValue)] [\(match.record.kind)] \(match.record.summary)" + ] + + if let evidence = match.record.evidence.first, + !evidence.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { + components.append(" Evidence: \(evidence)") + } + + if !match.record.tags.isEmpty { + components.append(" Tags: \(match.record.tags.joined(separator: ", "))") + } + + return components.joined(separator: "\n") + } + + private static func clamp(_ value: Double) -> Double { + min(1, max(0, value)) + } +} diff --git a/Sources/CodexKit/Memory/MemoryStore.swift b/Sources/CodexKit/Memory/MemoryStore.swift new file mode 100644 index 0000000..0e791cd --- /dev/null +++ b/Sources/CodexKit/Memory/MemoryStore.swift @@ -0,0 +1,24 @@ +import Foundation + +public protocol MemoryStoring: Sendable { + func put(_ record: MemoryRecord) async throws + func putMany(_ records: [MemoryRecord]) async throws + func upsert(_ record: MemoryRecord, dedupeKey: String) async throws + func query(_ query: MemoryQuery) async throws -> MemoryQueryResult + func compact(_ request: MemoryCompactionRequest) async throws + func archive(ids: [String], namespace: String) async throws + func delete(ids: [String], namespace: String) async throws + + @discardableResult + func pruneExpired( + now: Date, + namespace: String + ) async throws -> Int +} + +public extension MemoryStoring { + @discardableResult + func pruneExpired(namespace: String) async throws -> Int { + try await pruneExpired(now: Date(), namespace: namespace) + } +} diff --git a/Sources/CodexKit/Memory/SQLiteMemoryStore.swift b/Sources/CodexKit/Memory/SQLiteMemoryStore.swift new file mode 100644 index 0000000..2bff1f2 --- /dev/null +++ b/Sources/CodexKit/Memory/SQLiteMemoryStore.swift @@ -0,0 +1,626 @@ +import Foundation +import SQLite3 + +public actor SQLiteMemoryStore: MemoryStoring { + private let url: URL + private nonisolated(unsafe) var database: OpaquePointer? + private let encoder = JSONEncoder() + private let decoder = JSONDecoder() + + public init(url: URL) throws { + self.url = url + self.database = try Self.openDatabase(at: url) + try Self.createSchemaIfNeeded(in: database) + } + + deinit { + if let database { + sqlite3_close(database) + } + } + + public func put(_ record: MemoryRecord) async throws { + try MemoryQueryEngine.validateNamespace(record.namespace) + try transaction { + try ensureRecordIDAvailable(record.id, namespace: record.namespace) + if let dedupeKey = record.dedupeKey { + try ensureDedupeKeyAvailable(dedupeKey, namespace: record.namespace) + } + try upsertRecord(record) + } + } + + public func putMany(_ records: [MemoryRecord]) async throws { + try transaction { + for record in records { + try MemoryQueryEngine.validateNamespace(record.namespace) + try ensureRecordIDAvailable(record.id, namespace: record.namespace) + if let dedupeKey = record.dedupeKey { + try ensureDedupeKeyAvailable(dedupeKey, namespace: record.namespace) + } + try upsertRecord(record) + } + } + } + + public func upsert(_ record: MemoryRecord, dedupeKey: String) async throws { + try MemoryQueryEngine.validateNamespace(record.namespace) + try transaction { + try deleteRecord(withDedupeKey: dedupeKey, namespace: record.namespace) + try deleteRecord(id: record.id, namespace: record.namespace) + var updatedRecord = record + updatedRecord.dedupeKey = dedupeKey + try upsertRecord(updatedRecord) + } + } + + public func query(_ query: MemoryQuery) async throws -> MemoryQueryResult { + try MemoryQueryEngine.validateNamespace(query.namespace) + let records = try loadRecords(namespace: query.namespace) + let rawScores = try loadFTSRawScores( + namespace: query.namespace, + queryText: query.text + ) + + let candidates = records.map { record in + MemoryQueryEngine.Candidate( + record: record, + rawTextScore: rawScores[record.id] + ) + } + + return try MemoryQueryEngine.evaluate( + candidates: candidates, + query: query + ) + } + + public func compact(_ request: MemoryCompactionRequest) async throws { + try MemoryQueryEngine.validateNamespace(request.replacement.namespace) + try transaction { + try ensureRecordIDAvailable(request.replacement.id, namespace: request.replacement.namespace) + if let dedupeKey = request.replacement.dedupeKey { + try ensureDedupeKeyAvailable(dedupeKey, namespace: request.replacement.namespace) + } + try upsertRecord(request.replacement) + for sourceID in request.sourceIDs { + try archiveRecord(id: sourceID, namespace: request.replacement.namespace) + } + } + } + + public func archive(ids: [String], namespace: String) async throws { + try MemoryQueryEngine.validateNamespace(namespace) + try transaction { + for id in ids { + try archiveRecord(id: id, namespace: namespace) + } + } + } + + public func delete(ids: [String], namespace: String) async throws { + try MemoryQueryEngine.validateNamespace(namespace) + try transaction { + for id in ids { + try deleteRecord(id: id, namespace: namespace) + } + } + } + + @discardableResult + public func pruneExpired( + now: Date, + namespace: String + ) async throws -> Int { + try MemoryQueryEngine.validateNamespace(namespace) + let expiredIDs = try loadRecords(namespace: namespace) + .filter { record in + !record.isPinned && + record.status == .active && + (record.expiresAt?.compare(now) == .orderedAscending || + record.expiresAt?.compare(now) == .orderedSame) + } + .map(\.id) + + try transaction { + for id in expiredIDs { + try deleteRecord(id: id, namespace: namespace) + } + } + return expiredIDs.count + } + + private static func openDatabase(at url: URL) throws -> OpaquePointer { + let directory = url.deletingLastPathComponent() + if !directory.path.isEmpty { + try FileManager.default.createDirectory( + at: directory, + withIntermediateDirectories: true + ) + } + + var database: OpaquePointer? + let result = sqlite3_open_v2( + url.path, + &database, + SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX, + nil + ) + guard result == SQLITE_OK, let database else { + throw sqliteError( + database, + message: "Failed to open SQLite memory store." + ) + } + sqlite3_exec(database, "PRAGMA foreign_keys = ON;", nil, nil, nil) + return database + } + + private static func createSchemaIfNeeded(in database: OpaquePointer?) throws { + let schema = """ + CREATE TABLE IF NOT EXISTS memory_records ( + namespace TEXT NOT NULL, + id TEXT NOT NULL, + scope TEXT NOT NULL, + kind TEXT NOT NULL, + summary TEXT NOT NULL, + evidence_json TEXT NOT NULL, + importance REAL NOT NULL, + created_at REAL NOT NULL, + observed_at REAL, + expires_at REAL, + tags_json TEXT NOT NULL, + related_ids_json TEXT NOT NULL, + dedupe_key TEXT, + is_pinned INTEGER NOT NULL, + attributes_json TEXT, + status TEXT NOT NULL, + PRIMARY KEY(namespace, id) + ); + CREATE UNIQUE INDEX IF NOT EXISTS memory_records_namespace_dedupe + ON memory_records(namespace, dedupe_key) + WHERE dedupe_key IS NOT NULL; + CREATE INDEX IF NOT EXISTS memory_records_namespace_scope + ON memory_records(namespace, scope); + CREATE INDEX IF NOT EXISTS memory_records_namespace_kind + ON memory_records(namespace, kind); + CREATE INDEX IF NOT EXISTS memory_records_namespace_status + ON memory_records(namespace, status); + CREATE TABLE IF NOT EXISTS memory_tags ( + namespace TEXT NOT NULL, + record_id TEXT NOT NULL, + tag TEXT NOT NULL, + FOREIGN KEY(namespace, record_id) + REFERENCES memory_records(namespace, id) + ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS memory_tags_lookup + ON memory_tags(namespace, tag, record_id); + CREATE TABLE IF NOT EXISTS memory_related_ids ( + namespace TEXT NOT NULL, + record_id TEXT NOT NULL, + related_id TEXT NOT NULL, + FOREIGN KEY(namespace, record_id) + REFERENCES memory_records(namespace, id) + ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS memory_related_lookup + ON memory_related_ids(namespace, related_id, record_id); + CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts + USING fts5(namespace UNINDEXED, record_id UNINDEXED, content); + """ + try execSQL(database, schema) + } + + private func ensureRecordIDAvailable( + _ id: String, + namespace: String + ) throws { + if try recordExists(id: id, namespace: namespace) { + throw MemoryStoreError.duplicateRecordID(id) + } + } + + private func ensureDedupeKeyAvailable( + _ dedupeKey: String, + namespace: String + ) throws { + if try recordExists(dedupeKey: dedupeKey, namespace: namespace) { + throw MemoryStoreError.duplicateDedupeKey(dedupeKey) + } + } + + private func loadRecords(namespace: String) throws -> [MemoryRecord] { + let sql = """ + SELECT + id, scope, kind, summary, evidence_json, importance, + created_at, observed_at, expires_at, tags_json, + related_ids_json, dedupe_key, is_pinned, attributes_json, status + FROM memory_records + WHERE namespace = ?; + """ + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + try bindText(namespace, to: statement, index: 1) + + var records: [MemoryRecord] = [] + while sqlite3_step(statement) == SQLITE_ROW { + let id = try columnText(statement, index: 0) + let scope = MemoryScope(rawValue: try columnText(statement, index: 1)) + let kind = try columnText(statement, index: 2) + let summary = try columnText(statement, index: 3) + let evidence = try decodeJSON([String].self, from: try columnText(statement, index: 4)) + let importance = sqlite3_column_double(statement, 5) + let createdAt = Date(timeIntervalSince1970: sqlite3_column_double(statement, 6)) + let observedAt = sqlite3_column_type(statement, 7) == SQLITE_NULL + ? nil + : Date(timeIntervalSince1970: sqlite3_column_double(statement, 7)) + let expiresAt = sqlite3_column_type(statement, 8) == SQLITE_NULL + ? nil + : Date(timeIntervalSince1970: sqlite3_column_double(statement, 8)) + let tags = try decodeJSON([String].self, from: try columnText(statement, index: 9)) + let relatedIDs = try decodeJSON([String].self, from: try columnText(statement, index: 10)) + let dedupeKey = sqlite3_column_type(statement, 11) == SQLITE_NULL ? nil : try columnText(statement, index: 11) + let isPinned = sqlite3_column_int(statement, 12) == 1 + let attributes = sqlite3_column_type(statement, 13) == SQLITE_NULL + ? nil + : try decodeJSON(JSONValue.self, from: try columnText(statement, index: 13)) + let status = MemoryRecordStatus(rawValue: try columnText(statement, index: 14)) ?? .active + + records.append( + MemoryRecord( + id: id, + namespace: namespace, + scope: scope, + kind: kind, + summary: summary, + evidence: evidence, + importance: importance, + createdAt: createdAt, + observedAt: observedAt, + expiresAt: expiresAt, + tags: tags, + relatedIDs: relatedIDs, + dedupeKey: dedupeKey, + isPinned: isPinned, + attributes: attributes, + status: status + ) + ) + } + + return records + } + + private func loadFTSRawScores( + namespace: String, + queryText: String? + ) throws -> [String: Double] { + let matchQuery = ftsQuery(from: queryText) + guard !matchQuery.isEmpty else { + return [:] + } + + let sql = """ + SELECT record_id, bm25(memory_fts) + FROM memory_fts + WHERE namespace = ? AND memory_fts MATCH ?; + """ + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + try bindText(namespace, to: statement, index: 1) + try bindText(matchQuery, to: statement, index: 2) + + var scores: [String: Double] = [:] + while sqlite3_step(statement) == SQLITE_ROW { + let recordID = try columnText(statement, index: 0) + let score = sqlite3_column_double(statement, 1) + scores[recordID] = score + } + return scores + } + + private func recordExists( + id: String, + namespace: String + ) throws -> Bool { + let sql = "SELECT 1 FROM memory_records WHERE namespace = ? AND id = ? LIMIT 1;" + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + try bindText(namespace, to: statement, index: 1) + try bindText(id, to: statement, index: 2) + return sqlite3_step(statement) == SQLITE_ROW + } + + private func recordExists( + dedupeKey: String, + namespace: String + ) throws -> Bool { + let sql = "SELECT 1 FROM memory_records WHERE namespace = ? AND dedupe_key = ? LIMIT 1;" + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + try bindText(namespace, to: statement, index: 1) + try bindText(dedupeKey, to: statement, index: 2) + return sqlite3_step(statement) == SQLITE_ROW + } + + private func archiveRecord( + id: String, + namespace: String + ) throws { + let sql = """ + UPDATE memory_records + SET status = ? + WHERE namespace = ? AND id = ?; + """ + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + try bindText(MemoryRecordStatus.archived.rawValue, to: statement, index: 1) + try bindText(namespace, to: statement, index: 2) + try bindText(id, to: statement, index: 3) + try step(statement) + } + + private func deleteRecord( + id: String, + namespace: String + ) throws { + try exec( + "DELETE FROM memory_fts WHERE namespace = ? AND record_id = ?;", + bindings: [.text(namespace), .text(id)] + ) + try exec( + "DELETE FROM memory_related_ids WHERE namespace = ? AND record_id = ?;", + bindings: [.text(namespace), .text(id)] + ) + try exec( + "DELETE FROM memory_tags WHERE namespace = ? AND record_id = ?;", + bindings: [.text(namespace), .text(id)] + ) + try exec( + "DELETE FROM memory_records WHERE namespace = ? AND id = ?;", + bindings: [.text(namespace), .text(id)] + ) + } + + private func deleteRecord( + withDedupeKey dedupeKey: String, + namespace: String + ) throws { + let statement = try prepare( + "SELECT id FROM memory_records WHERE namespace = ? AND dedupe_key = ? LIMIT 1;" + ) + defer { sqlite3_finalize(statement) } + try bindText(namespace, to: statement, index: 1) + try bindText(dedupeKey, to: statement, index: 2) + if sqlite3_step(statement) == SQLITE_ROW { + let id = try columnText(statement, index: 0) + try deleteRecord(id: id, namespace: namespace) + } + } + + private func upsertRecord(_ record: MemoryRecord) throws { + let sql = """ + INSERT OR REPLACE INTO memory_records ( + namespace, id, scope, kind, summary, evidence_json, importance, + created_at, observed_at, expires_at, tags_json, related_ids_json, + dedupe_key, is_pinned, attributes_json, status + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?); + """ + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + + try bindText(record.namespace, to: statement, index: 1) + try bindText(record.id, to: statement, index: 2) + try bindText(record.scope.rawValue, to: statement, index: 3) + try bindText(record.kind, to: statement, index: 4) + try bindText(record.summary, to: statement, index: 5) + try bindText(try encodeJSON(record.evidence), to: statement, index: 6) + sqlite3_bind_double(statement, 7, record.importance) + sqlite3_bind_double(statement, 8, record.createdAt.timeIntervalSince1970) + if let observedAt = record.observedAt { + sqlite3_bind_double(statement, 9, observedAt.timeIntervalSince1970) + } else { + sqlite3_bind_null(statement, 9) + } + if let expiresAt = record.expiresAt { + sqlite3_bind_double(statement, 10, expiresAt.timeIntervalSince1970) + } else { + sqlite3_bind_null(statement, 10) + } + try bindText(try encodeJSON(record.tags), to: statement, index: 11) + try bindText(try encodeJSON(record.relatedIDs), to: statement, index: 12) + if let dedupeKey = record.dedupeKey { + try bindText(dedupeKey, to: statement, index: 13) + } else { + sqlite3_bind_null(statement, 13) + } + sqlite3_bind_int(statement, 14, record.isPinned ? 1 : 0) + if let attributes = record.attributes { + try bindText(try encodeJSON(attributes), to: statement, index: 15) + } else { + sqlite3_bind_null(statement, 15) + } + try bindText(record.status.rawValue, to: statement, index: 16) + try step(statement) + + try exec( + "DELETE FROM memory_tags WHERE namespace = ? AND record_id = ?;", + bindings: [.text(record.namespace), .text(record.id)] + ) + try exec( + "DELETE FROM memory_related_ids WHERE namespace = ? AND record_id = ?;", + bindings: [.text(record.namespace), .text(record.id)] + ) + try exec( + "DELETE FROM memory_fts WHERE namespace = ? AND record_id = ?;", + bindings: [.text(record.namespace), .text(record.id)] + ) + + for tag in record.tags { + try exec( + "INSERT INTO memory_tags(namespace, record_id, tag) VALUES (?, ?, ?);", + bindings: [.text(record.namespace), .text(record.id), .text(tag)] + ) + } + + for relatedID in record.relatedIDs { + try exec( + "INSERT INTO memory_related_ids(namespace, record_id, related_id) VALUES (?, ?, ?);", + bindings: [.text(record.namespace), .text(record.id), .text(relatedID)] + ) + } + + let ftsContent = ([record.summary] + record.evidence + record.tags + [record.kind]).joined(separator: " ") + try exec( + "INSERT INTO memory_fts(namespace, record_id, content) VALUES (?, ?, ?);", + bindings: [.text(record.namespace), .text(record.id), .text(ftsContent)] + ) + } + + private func transaction(_ operation: () throws -> Void) throws { + try exec("BEGIN IMMEDIATE;") + do { + try operation() + try exec("COMMIT;") + } catch { + try? exec("ROLLBACK;") + throw error + } + } + + private func prepare(_ sql: String) throws -> OpaquePointer? { + guard let database else { + throw sqliteError(message: "SQLite database is unavailable.") + } + var statement: OpaquePointer? + let result = sqlite3_prepare_v2(database, sql, -1, &statement, nil) + guard result == SQLITE_OK else { + throw sqliteError(message: "Failed to prepare SQLite statement.") + } + return statement + } + + private func exec( + _ sql: String, + bindings: [SQLiteBinding] = [] + ) throws { + let statement = try prepare(sql) + defer { sqlite3_finalize(statement) } + + for (index, binding) in bindings.enumerated() { + try bind(binding, to: statement, index: Int32(index + 1)) + } + + try step(statement) + } + + private func step(_ statement: OpaquePointer?) throws { + let result = sqlite3_step(statement) + guard result == SQLITE_DONE || result == SQLITE_ROW else { + throw sqliteError(message: "SQLite step failed.") + } + } + + private func bind( + _ binding: SQLiteBinding, + to statement: OpaquePointer?, + index: Int32 + ) throws { + switch binding { + case let .text(value): + try bindText(value, to: statement, index: index) + case let .double(value): + sqlite3_bind_double(statement, index, value) + case .null: + sqlite3_bind_null(statement, index) + } + } + + private func bindText( + _ value: String, + to statement: OpaquePointer?, + index: Int32 + ) throws { + let result = sqlite3_bind_text(statement, index, value, -1, SQLITE_TRANSIENT) + guard result == SQLITE_OK else { + throw sqliteError(message: "Failed to bind SQLite text value.") + } + } + + private func columnText( + _ statement: OpaquePointer?, + index: Int32 + ) throws -> String { + guard let cString = sqlite3_column_text(statement, index) else { + throw sqliteError(message: "SQLite column was unexpectedly null.") + } + return String(cString: cString) + } + + private func encodeJSON(_ value: T) throws -> String { + let data = try encoder.encode(value) + return String(decoding: data, as: UTF8.self) + } + + private func decodeJSON( + _ type: T.Type, + from string: String + ) throws -> T { + try decoder.decode(type, from: Data(string.utf8)) + } + + private func ftsQuery(from value: String?) -> String { + let tokens = MemoryQueryEngine.tokenize(value) + guard !tokens.isEmpty else { + return "" + } + return tokens.joined(separator: " OR ") + } + + private static func sqliteError( + _ database: OpaquePointer?, + message: String + ) -> NSError { + let detail = if let database, let messagePointer = sqlite3_errmsg(database) { + String(cString: messagePointer) + } else { + "Unknown SQLite error" + } + + return NSError( + domain: "CodexKit.SQLiteMemoryStore", + code: Int(sqlite3_errcode(database)), + userInfo: [NSLocalizedDescriptionKey: "\(message) \(detail)"] + ) + } + + private func sqliteError(message: String) -> NSError { + Self.sqliteError(database, message: message) + } +} + +private enum SQLiteBinding { + case text(String) + case double(Double) + case null +} + +private let SQLITE_TRANSIENT = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + +private func execSQL( + _ database: OpaquePointer?, + _ sql: String +) throws { + var errorPointer: UnsafeMutablePointer? + let result = sqlite3_exec(database, sql, nil, nil, &errorPointer) + guard result == SQLITE_OK else { + let detail = errorPointer.map { String(cString: $0) } ?? "Unknown SQLite error" + sqlite3_free(errorPointer) + throw NSError( + domain: "CodexKit.SQLiteMemoryStore", + code: Int(result), + userInfo: [NSLocalizedDescriptionKey: detail] + ) + } +} diff --git a/Sources/CodexKit/Runtime/AgentModels.swift b/Sources/CodexKit/Runtime/AgentModels.swift index c436c7a..13c48c1 100644 --- a/Sources/CodexKit/Runtime/AgentModels.swift +++ b/Sources/CodexKit/Runtime/AgentModels.swift @@ -212,17 +212,20 @@ public struct UserMessageRequest: Codable, Hashable, Sendable { public var images: [AgentImageAttachment] public var personaOverride: AgentPersonaStack? public var skillOverrideIDs: [String]? + public var memorySelection: MemorySelection? public init( text: String, images: [AgentImageAttachment] = [], personaOverride: AgentPersonaStack? = nil, - skillOverrideIDs: [String]? = nil + skillOverrideIDs: [String]? = nil, + memorySelection: MemorySelection? = nil ) { self.text = text self.images = images self.personaOverride = personaOverride self.skillOverrideIDs = skillOverrideIDs + self.memorySelection = memorySelection } public var hasContent: Bool { @@ -234,6 +237,7 @@ public struct UserMessageRequest: Codable, Hashable, Sendable { case images case personaOverride case skillOverrideIDs + case memorySelection } public init(from decoder: Decoder) throws { @@ -242,6 +246,7 @@ public struct UserMessageRequest: Codable, Hashable, Sendable { images = try container.decodeIfPresent([AgentImageAttachment].self, forKey: .images) ?? [] personaOverride = try container.decodeIfPresent(AgentPersonaStack.self, forKey: .personaOverride) skillOverrideIDs = try container.decodeIfPresent([String].self, forKey: .skillOverrideIDs) + memorySelection = try container.decodeIfPresent(MemorySelection.self, forKey: .memorySelection) } } @@ -250,6 +255,7 @@ public struct AgentThread: Identifiable, Codable, Hashable, Sendable { public var title: String? public var personaStack: AgentPersonaStack? public var skillIDs: [String] + public var memoryContext: AgentMemoryContext? public var createdAt: Date public var updatedAt: Date public var status: AgentThreadStatus @@ -259,6 +265,7 @@ public struct AgentThread: Identifiable, Codable, Hashable, Sendable { title: String? = nil, personaStack: AgentPersonaStack? = nil, skillIDs: [String] = [], + memoryContext: AgentMemoryContext? = nil, createdAt: Date = Date(), updatedAt: Date = Date(), status: AgentThreadStatus = .idle @@ -267,6 +274,7 @@ public struct AgentThread: Identifiable, Codable, Hashable, Sendable { self.title = title self.personaStack = personaStack self.skillIDs = skillIDs + self.memoryContext = memoryContext self.createdAt = createdAt self.updatedAt = updatedAt self.status = status @@ -277,6 +285,7 @@ public struct AgentThread: Identifiable, Codable, Hashable, Sendable { case title case personaStack case skillIDs + case memoryContext case createdAt case updatedAt case status @@ -288,6 +297,7 @@ public struct AgentThread: Identifiable, Codable, Hashable, Sendable { title = try container.decodeIfPresent(String.self, forKey: .title) personaStack = try container.decodeIfPresent(AgentPersonaStack.self, forKey: .personaStack) skillIDs = try container.decodeIfPresent([String].self, forKey: .skillIDs) ?? [] + memoryContext = try container.decodeIfPresent(AgentMemoryContext.self, forKey: .memoryContext) createdAt = try container.decodeIfPresent(Date.self, forKey: .createdAt) ?? Date() updatedAt = try container.decodeIfPresent(Date.self, forKey: .updatedAt) ?? createdAt status = try container.decodeIfPresent(AgentThreadStatus.self, forKey: .status) ?? .idle diff --git a/Sources/CodexKit/Runtime/AgentRuntime.swift b/Sources/CodexKit/Runtime/AgentRuntime.swift index dea445d..7b4decd 100644 --- a/Sources/CodexKit/Runtime/AgentRuntime.swift +++ b/Sources/CodexKit/Runtime/AgentRuntime.swift @@ -20,6 +20,7 @@ public actor AgentRuntime { public let backend: any AgentBackend public let approvalPresenter: any ApprovalPresenting public let stateStore: any RuntimeStateStoring + public let memory: AgentMemoryConfiguration? public let baseInstructions: String? public let tools: [ToolRegistration] public let skills: [AgentSkill] @@ -31,6 +32,7 @@ public actor AgentRuntime { backend: any AgentBackend, approvalPresenter: any ApprovalPresenting, stateStore: any RuntimeStateStoring, + memory: AgentMemoryConfiguration? = nil, baseInstructions: String? = nil, tools: [ToolRegistration] = [], skills: [AgentSkill] = [], @@ -41,6 +43,7 @@ public actor AgentRuntime { self.backend = backend self.approvalPresenter = approvalPresenter self.stateStore = stateStore + self.memory = memory self.baseInstructions = baseInstructions self.tools = tools self.skills = skills @@ -53,6 +56,7 @@ public actor AgentRuntime { private let sessionManager: ChatGPTSessionManager private let toolRegistry: ToolRegistry private let approvalCoordinator: ApprovalCoordinator + private let memoryConfiguration: AgentMemoryConfiguration? private let baseInstructions: String? private let definitionSourceLoader: AgentDefinitionSourceLoader private var skillsByID: [String: AgentSkill] @@ -153,6 +157,7 @@ public actor AgentRuntime { self.approvalCoordinator = ApprovalCoordinator( presenter: configuration.approvalPresenter ) + self.memoryConfiguration = configuration.memory self.baseInstructions = configuration.baseInstructions ?? configuration.backend.baseInstructions self.definitionSourceLoader = configuration.definitionSourceLoader self.skillsByID = try Self.validatedSkills(from: configuration.skills) @@ -264,7 +269,8 @@ public actor AgentRuntime { title: String? = nil, personaStack: AgentPersonaStack? = nil, personaSource: AgentDefinitionSource? = nil, - skillIDs: [String] = [] + skillIDs: [String] = [], + memoryContext: AgentMemoryContext? = nil ) async throws -> AgentThread { try assertSkillsExist(skillIDs) let resolvedPersonaStack: AgentPersonaStack? @@ -288,6 +294,7 @@ public actor AgentRuntime { } thread.personaStack = resolvedPersonaStack thread.skillIDs = skillIDs + thread.memoryContext = memoryContext try await upsertThread(thread) return thread } @@ -329,7 +336,7 @@ public actor AgentRuntime { thread: thread, message: request ) - let resolvedInstructions = resolveInstructions( + let resolvedInstructions = await resolveInstructions( thread: thread, message: request, resolvedTurnSkills: resolvedTurnSkills @@ -395,7 +402,7 @@ public actor AgentRuntime { public func resolvedInstructionsPreview( for threadID: String, request: UserMessageRequest - ) throws -> String { + ) async throws -> String { guard let thread = thread(for: threadID) else { throw AgentRuntimeError.threadNotFound(threadID) } @@ -405,13 +412,27 @@ public actor AgentRuntime { message: request ) - return resolveInstructions( + return await resolveInstructions( thread: thread, message: request, resolvedTurnSkills: resolvedTurnSkills ) } + public func memoryQueryPreview( + for threadID: String, + request: UserMessageRequest + ) async throws -> MemoryQueryResult? { + guard let thread = thread(for: threadID) else { + throw AgentRuntimeError.threadNotFound(threadID) + } + + return await resolvedMemoryQuery( + thread: thread, + message: request + ) + } + private func consumeTurnStream( _ turnStream: any AgentTurnStreaming, for threadID: String, @@ -598,6 +619,14 @@ public actor AgentRuntime { return thread.skillIDs } + public func memoryContext(for threadID: String) throws -> AgentMemoryContext? { + guard let thread = thread(for: threadID) else { + throw AgentRuntimeError.threadNotFound(threadID) + } + + return thread.memoryContext + } + public func setSkillIDs( _ skillIDs: [String], for threadID: String @@ -612,6 +641,19 @@ public actor AgentRuntime { try await persistState() } + public func setMemoryContext( + _ memoryContext: AgentMemoryContext?, + for threadID: String + ) async throws { + guard let index = state.threads.firstIndex(where: { $0.id == threadID }) else { + throw AgentRuntimeError.threadNotFound(threadID) + } + + state.threads[index].memoryContext = memoryContext + state.threads[index].updatedAt = Date() + try await persistState() + } + private func upsertThread(_ thread: AgentThread) async throws { if let index = state.threads.firstIndex(where: { $0.id == thread.id }) { var mergedThread = thread @@ -624,6 +666,9 @@ public actor AgentRuntime { if mergedThread.skillIDs.isEmpty { mergedThread.skillIDs = state.threads[index].skillIDs } + if mergedThread.memoryContext == nil { + mergedThread.memoryContext = state.threads[index].memoryContext + } state.threads[index] = mergedThread } else { state.threads.append(thread) @@ -671,14 +716,46 @@ public actor AgentRuntime { thread: AgentThread, message: UserMessageRequest, resolvedTurnSkills: ResolvedTurnSkills - ) -> String { - AgentInstructionCompiler.compile( + ) async -> String { + let compiled = AgentInstructionCompiler.compile( baseInstructions: baseInstructions, threadPersonaStack: thread.personaStack, threadSkills: resolvedTurnSkills.threadSkills, turnPersonaOverride: message.personaOverride, turnSkills: resolvedTurnSkills.turnSkills ) + + guard let queryResult = await resolvedMemoryQuery( + thread: thread, + message: message + ), + let memoryConfiguration + else { + return compiled + } + + let budget = resolvedMemoryBudget( + thread: thread, + message: message, + fallback: memoryConfiguration.defaultReadBudget + ) + let renderedMemory = memoryConfiguration.promptRenderer.render( + result: queryResult, + budget: budget + ) + guard !renderedMemory.isEmpty else { + return compiled + } + + if compiled.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { + return renderedMemory + } + + return """ + \(compiled) + + \(renderedMemory) + """ } private static func isUnauthorizedError(_ error: Error) -> Bool { @@ -725,6 +802,140 @@ public actor AgentRuntime { ) } + private func resolvedMemoryQuery( + thread: AgentThread, + message: UserMessageRequest + ) async -> MemoryQueryResult? { + guard let memoryConfiguration else { + return nil + } + + guard let query = resolvedMemoryQuery( + thread: thread, + message: message, + fallbackRanking: memoryConfiguration.defaultRanking, + fallbackBudget: memoryConfiguration.defaultReadBudget + ) else { + return nil + } + + return try? await memoryConfiguration.store.query(query) + } + + private func resolvedMemoryQuery( + thread: AgentThread, + message: UserMessageRequest, + fallbackRanking: MemoryRankingWeights, + fallbackBudget: MemoryReadBudget + ) -> MemoryQuery? { + let selection = message.memorySelection + if selection?.mode == .disable { + return nil + } + + let threadContext = thread.memoryContext + let namespace = selection?.namespace ?? + threadContext?.namespace + + guard let namespace else { + return nil + } + + let scopes: [MemoryScope] + switch selection?.mode ?? .inherit { + case .append: + scopes = uniqueScopes((threadContext?.scopes ?? []) + (selection?.scopes ?? [])) + case .replace: + scopes = selection?.scopes ?? [] + case .disable: + return nil + case .inherit: + if let selection, + !selection.scopes.isEmpty { + scopes = selection.scopes + } else { + scopes = threadContext?.scopes ?? [] + } + } + + let kinds = resolvedValues( + mode: selection?.mode ?? .inherit, + threadValues: threadContext?.kinds ?? [], + selectionValues: selection?.kinds ?? [] + ) + let tags = resolvedValues( + mode: selection?.mode ?? .inherit, + threadValues: threadContext?.tags ?? [], + selectionValues: selection?.tags ?? [] + ) + let relatedIDs = resolvedValues( + mode: selection?.mode ?? .inherit, + threadValues: threadContext?.relatedIDs ?? [], + selectionValues: selection?.relatedIDs ?? [] + ) + + let recencyWindow = selection?.recencyWindow + ?? threadContext?.recencyWindow + let minImportance = selection?.minImportance + ?? threadContext?.minImportance + let ranking = selection?.ranking + ?? threadContext?.ranking + ?? fallbackRanking + let budget = resolvedMemoryBudget( + thread: thread, + message: message, + fallback: fallbackBudget + ) + let text = selection?.text ?? message.text + + return MemoryQuery( + namespace: namespace, + scopes: scopes, + text: text, + kinds: kinds, + tags: tags, + relatedIDs: relatedIDs, + recencyWindow: recencyWindow, + minImportance: minImportance, + ranking: ranking, + limit: budget.maxItems, + maxCharacters: budget.maxCharacters, + includeArchived: false + ) + } + + private func resolvedMemoryBudget( + thread: AgentThread, + message: UserMessageRequest, + fallback: MemoryReadBudget + ) -> MemoryReadBudget { + message.memorySelection?.readBudget + ?? thread.memoryContext?.readBudget + ?? fallback + } + + private func uniqueScopes(_ scopes: [MemoryScope]) -> [MemoryScope] { + var seen: Set = [] + return scopes.filter { seen.insert($0).inserted } + } + + private func resolvedValues( + mode: MemorySelectionMode, + threadValues: [String], + selectionValues: [String] + ) -> [String] { + switch mode { + case .append: + return Array(Set(threadValues + selectionValues)).sorted() + case .replace: + return selectionValues + case .disable: + return [] + case .inherit: + return selectionValues.isEmpty ? threadValues : selectionValues + } + } + private func compileToolPolicy(from skills: [AgentSkill]) -> CompiledSkillToolPolicy { var allowedToolNames: Set? var requiredToolNames: Set = [] diff --git a/Tests/CodexKitTests/AgentRuntimeTests.swift b/Tests/CodexKitTests/AgentRuntimeTests.swift index 19ae1a1..630129e 100644 --- a/Tests/CodexKitTests/AgentRuntimeTests.swift +++ b/Tests/CodexKitTests/AgentRuntimeTests.swift @@ -587,6 +587,177 @@ final class AgentRuntimeTests: XCTestCase { XCTAssertTrue(preview.contains("[health_coach: Health Coach]")) } + func testRuntimeInjectsRelevantMemoryIntoInstructionsAndPreviewMatches() async throws { + let backend = InMemoryAgentBackend( + baseInstructions: "Base host instructions." + ) + let store = InMemoryMemoryStore(initialRecords: [ + MemoryRecord( + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Eleanor remembers being overruled on the trade bill.", + evidence: ["She warned the player twice before being ignored."], + importance: 0.9, + tags: ["trade"] + ), + MemoryRecord( + namespace: "oval-office", + scope: "actor:sophia_ramirez", + kind: "grievance", + summary: "Sophia is focused on education messaging.", + importance: 0.8 + ), + ]) + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: backend, + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init(store: store) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Memory", + memoryContext: AgentMemoryContext( + namespace: "oval-office", + scopes: ["actor:eleanor_price"] + ) + ) + + let preview = try await runtime.memoryQueryPreview( + for: thread.id, + request: UserMessageRequest(text: "What does Eleanor still remember about trade?") + ) + XCTAssertEqual(preview?.matches.map(\.record.scope.rawValue), ["actor:eleanor_price"]) + + let stream = try await runtime.sendMessage( + UserMessageRequest(text: "What does Eleanor still remember about trade?"), + in: thread.id + ) + for try await _ in stream {} + + let instructions = await backend.receivedInstructions() + let resolved = try XCTUnwrap(instructions.last) + XCTAssertTrue(resolved.contains("Relevant Memory:")) + XCTAssertTrue(resolved.contains("Eleanor remembers being overruled on the trade bill.")) + XCTAssertFalse(resolved.contains("Sophia is focused on education messaging.")) + } + + func testRuntimeMemorySelectionCanReplaceOrDisableThreadDefaults() async throws { + let backend = InMemoryAgentBackend( + baseInstructions: "Base host instructions." + ) + let store = InMemoryMemoryStore(initialRecords: [ + MemoryRecord( + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Eleanor grievance." + ), + MemoryRecord( + namespace: "oval-office", + scope: "actor:sophia_ramirez", + kind: "grievance", + summary: "Sophia grievance." + ), + ]) + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: backend, + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init(store: store) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Scoped Memory", + memoryContext: AgentMemoryContext( + namespace: "oval-office", + scopes: ["actor:eleanor_price"] + ) + ) + + let replaceStream = try await runtime.sendMessage( + UserMessageRequest( + text: "Use Sophia memory instead.", + memorySelection: MemorySelection( + mode: .replace, + scopes: ["actor:sophia_ramirez"] + ) + ), + in: thread.id + ) + for try await _ in replaceStream {} + + let disableStream = try await runtime.sendMessage( + UserMessageRequest( + text: "Now disable memory.", + memorySelection: MemorySelection(mode: .disable) + ), + in: thread.id + ) + for try await _ in disableStream {} + + let instructions = await backend.receivedInstructions() + XCTAssertEqual(instructions.count, 2) + XCTAssertTrue(instructions[0].contains("Sophia grievance.")) + XCTAssertFalse(instructions[0].contains("Eleanor grievance.")) + XCTAssertFalse(instructions[1].contains("Relevant Memory:")) + } + + func testRuntimeGracefullyDegradesWhenMemoryStoreFails() async throws { + let backend = InMemoryAgentBackend( + baseInstructions: "Base host instructions." + ) + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: backend, + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init(store: ThrowingMemoryStore()) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Graceful", + memoryContext: AgentMemoryContext( + namespace: "oval-office", + scopes: ["actor:eleanor_price"] + ) + ) + + let stream = try await runtime.sendMessage( + UserMessageRequest(text: "This should still work."), + in: thread.id + ) + for try await _ in stream {} + + let instructions = await backend.receivedInstructions() + let resolved = try XCTUnwrap(instructions.last) + XCTAssertFalse(resolved.contains("Relevant Memory:")) + } + func testResolvedInstructionsPreviewThrowsForMissingThread() async throws { let runtime = try AgentRuntime(configuration: .init( authProvider: DemoChatGPTAuthProvider(), @@ -805,10 +976,54 @@ final class AgentRuntimeTests: XCTestCase { XCTAssertEqual(state.threads.count, 1) XCTAssertEqual(state.threads.first?.personaStack, nil) + XCTAssertEqual(state.threads.first?.memoryContext, nil) XCTAssertEqual(state.messagesByThread["thread-1"]?.first?.images, []) XCTAssertEqual(state.messagesByThread["thread-1"]?.first?.text, "Hello from legacy state") } + func testThreadMemoryContextPersistsAcrossRestore() async throws { + let stateStore = InMemoryRuntimeStateStore() + let secureStore = KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ) + let memoryContext = AgentMemoryContext( + namespace: "oval-office", + scopes: ["actor:eleanor_price", "world:public"], + readBudget: .init(maxItems: 4, maxCharacters: 800) + ) + + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: secureStore, + backend: InMemoryAgentBackend(), + approvalPresenter: AutoApprovalPresenter(), + stateStore: stateStore + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + let thread = try await runtime.createThread( + title: "Memory Restore", + memoryContext: memoryContext + ) + + let restoredRuntime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: secureStore, + backend: InMemoryAgentBackend(), + approvalPresenter: AutoApprovalPresenter(), + stateStore: stateStore + )) + + _ = try await restoredRuntime.restore() + + let restoredContext = try await restoredRuntime.memoryContext(for: thread.id) + let restoredThreads = await restoredRuntime.threads() + XCTAssertEqual(restoredContext, memoryContext) + XCTAssertEqual(restoredThreads.first?.memoryContext, memoryContext) + } + func testRuntimeStreamsToolApprovalAndCompletion() async throws { let runtime = try AgentRuntime(configuration: .init( authProvider: DemoChatGPTAuthProvider(), @@ -1162,3 +1377,29 @@ private final class ImageReplyTurn: AgentTurnStreaming, @unchecked Sendable { for _: String ) async throws {} } + +private actor ThrowingMemoryStore: MemoryStoring { + func put(_ record: MemoryRecord) async throws {} + + func putMany(_ records: [MemoryRecord]) async throws {} + + func upsert(_ record: MemoryRecord, dedupeKey: String) async throws {} + + func query(_ query: MemoryQuery) async throws -> MemoryQueryResult { + throw NSError( + domain: "CodexKitTests.ThrowingMemoryStore", + code: 1, + userInfo: [NSLocalizedDescriptionKey: "Simulated memory failure"] + ) + } + + func compact(_ request: MemoryCompactionRequest) async throws {} + + func archive(ids: [String], namespace: String) async throws {} + + func delete(ids: [String], namespace: String) async throws {} + + func pruneExpired(now: Date, namespace: String) async throws -> Int { + 0 + } +} diff --git a/Tests/CodexKitTests/MemoryStoreTests.swift b/Tests/CodexKitTests/MemoryStoreTests.swift new file mode 100644 index 0000000..6e09ef2 --- /dev/null +++ b/Tests/CodexKitTests/MemoryStoreTests.swift @@ -0,0 +1,270 @@ +import CodexKit +import Foundation +import XCTest + +final class MemoryStoreTests: XCTestCase { + func testSQLiteStorePersistsAndReloadsRecords() async throws { + let url = temporarySQLiteURL() + defer { try? FileManager.default.removeItem(at: url) } + + let store = try SQLiteMemoryStore(url: url) + let record = MemoryRecord( + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Eleanor remembers being overruled on the trade bill.", + evidence: ["The player dismissed her warning on day 12."], + importance: 0.9, + tags: ["trade", "advisors"], + relatedIDs: ["bill-12"], + dedupeKey: "eleanor-trade-day-12" + ) + + try await store.put(record) + + let reloaded = try SQLiteMemoryStore(url: url) + let result = try await reloaded.query( + MemoryQuery( + namespace: "oval-office", + scopes: ["actor:eleanor_price"], + text: "overruled trade warning", + limit: 5, + maxCharacters: 600 + ) + ) + + XCTAssertEqual(result.matches.map(\.record.id), [record.id]) + XCTAssertGreaterThan(result.matches[0].explanation.textScore, 0) + } + + func testPutManyIsAtomicWhenDuplicateIDIsPresent() async throws { + let url = temporarySQLiteURL() + defer { try? FileManager.default.removeItem(at: url) } + + let store = try SQLiteMemoryStore(url: url) + let existing = MemoryRecord( + id: "memory-1", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "fact", + summary: "Existing memory." + ) + try await store.put(existing) + + await XCTAssertThrowsErrorAsync( + try await store.putMany([ + MemoryRecord( + id: "memory-2", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "fact", + summary: "Should roll back." + ), + MemoryRecord( + id: "memory-1", + namespace: "oval-office", + scope: "actor:sophia_ramirez", + kind: "fact", + summary: "Duplicate id." + ), + ]) + ) { error in + XCTAssertEqual( + error as? MemoryStoreError, + .duplicateRecordID("memory-1") + ) + } + + let result = try await store.query( + MemoryQuery( + namespace: "oval-office", + scopes: [], + limit: 10, + maxCharacters: 1000 + ) + ) + XCTAssertEqual(result.matches.map(\.record.id), ["memory-1"]) + } + + func testUpsertIsRetrySafeByDedupeKey() async throws { + let store = InMemoryMemoryStore() + + try await store.upsert( + MemoryRecord( + id: "memory-1", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "quote_record", + summary: "Initial quote." + ), + dedupeKey: "press-quote-17" + ) + + try await store.upsert( + MemoryRecord( + id: "memory-2", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "quote_record", + summary: "Updated quote after retry." + ), + dedupeKey: "press-quote-17" + ) + + let result = try await store.query( + MemoryQuery( + namespace: "oval-office", + scopes: ["actor:eleanor_price"], + limit: 10, + maxCharacters: 1000 + ) + ) + + XCTAssertEqual(result.matches.count, 1) + XCTAssertEqual(result.matches[0].record.id, "memory-2") + XCTAssertEqual(result.matches[0].record.dedupeKey, "press-quote-17") + } + + func testQueryFiltersRankingAndCharacterBudget() async throws { + let store = InMemoryMemoryStore() + try await store.putMany([ + MemoryRecord( + id: "match-1", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Eleanor is still angry about the farm subsidy reversal.", + evidence: ["She warned the player twice before being ignored."], + importance: 0.95, + tags: ["farm", "economy"], + relatedIDs: ["policy-farm"] + ), + MemoryRecord( + id: "match-2", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "A second long grievance that should be trimmed by the budget.", + evidence: ["This extra line makes the rendered memory longer than the cap."], + importance: 0.80, + tags: ["farm"], + relatedIDs: ["policy-farm"] + ), + MemoryRecord( + id: "other-scope", + namespace: "oval-office", + scope: "actor:sophia_ramirez", + kind: "grievance", + summary: "Sophia has a separate grievance.", + importance: 0.99, + tags: ["farm"], + relatedIDs: ["policy-farm"] + ), + ]) + + let result = try await store.query( + MemoryQuery( + namespace: "oval-office", + scopes: ["actor:eleanor_price"], + text: "farm grievance warning", + kinds: ["grievance"], + tags: ["farm"], + relatedIDs: ["policy-farm"], + minImportance: 0.5, + limit: 10, + maxCharacters: 220 + ) + ) + + XCTAssertEqual(result.matches.count, 1) + XCTAssertEqual(result.matches.map(\.record.id), ["match-1"]) + XCTAssertTrue(result.truncated) + XCTAssertGreaterThan(result.matches[0].explanation.kindBoost, 0) + XCTAssertGreaterThan(result.matches[0].explanation.tagBoost, 0) + XCTAssertGreaterThan(result.matches[0].explanation.relatedIDBoost, 0) + } + + func testCompactArchivesSourcesAndPruneExpiredSkipsPinned() async throws { + let store = InMemoryMemoryStore() + let now = Date() + + try await store.putMany([ + MemoryRecord( + id: "source-1", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "fact", + summary: "Source memory one." + ), + MemoryRecord( + id: "source-2", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "fact", + summary: "Source memory two." + ), + MemoryRecord( + id: "expired-pinned", + namespace: "oval-office", + scope: "world:press", + kind: "summary", + summary: "Pinned memory should survive pruning.", + expiresAt: now.addingTimeInterval(-60), + isPinned: true + ), + MemoryRecord( + id: "expired-unpinned", + namespace: "oval-office", + scope: "world:press", + kind: "summary", + summary: "Unpinned memory should be removed.", + expiresAt: now.addingTimeInterval(-60) + ), + ]) + + try await store.compact( + MemoryCompactionRequest( + replacement: MemoryRecord( + id: "replacement", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "summary", + summary: "Compacted grievance summary." + ), + sourceIDs: ["source-1", "source-2"] + ) + ) + + let active = try await store.query( + MemoryQuery( + namespace: "oval-office", + scopes: ["actor:eleanor_price"], + limit: 10, + maxCharacters: 1000 + ) + ) + XCTAssertEqual(active.matches.map(\.record.id), ["replacement"]) + + let prunedCount = try await store.pruneExpired( + now: now, + namespace: "oval-office" + ) + XCTAssertEqual(prunedCount, 1) + + let remaining = try await store.query( + MemoryQuery( + namespace: "oval-office", + scopes: ["world:press"], + limit: 10, + maxCharacters: 1000 + ) + ) + XCTAssertEqual(remaining.matches.map(\.record.id), ["expired-pinned"]) + } + + private func temporarySQLiteURL() -> URL { + FileManager.default.temporaryDirectory + .appendingPathComponent(UUID().uuidString) + .appendingPathExtension("sqlite") + } +} From d96cca61479d6643a455460bbaa68c156ba5c686 Mon Sep 17 00:00:00 2001 From: Timothy Zelinsky Date: Sun, 22 Mar 2026 16:25:47 +1100 Subject: [PATCH 2/4] Extend memory tooling and docs --- README.md | 90 ++++++++++++++ .../CodexKit/Memory/InMemoryMemoryStore.swift | 74 ++++++++++++ Sources/CodexKit/Memory/MemoryModels.swift | 71 ++++++++++- Sources/CodexKit/Memory/MemoryStore.swift | 21 ++++ .../CodexKit/Memory/SQLiteMemoryStore.swift | 112 ++++++++++++++++-- Sources/CodexKit/Runtime/AgentRuntime.swift | 22 +++- Tests/CodexKitTests/AgentRuntimeTests.swift | 90 ++++++++++++++ Tests/CodexKitTests/MemoryStoreTests.swift | 73 ++++++++++++ 8 files changed, 541 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index e467cc7..3929294 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ Use `CodexKit` if you are building a SwiftUI/iOS app and want: - ChatGPT sign-in (device code or OAuth) - secure session persistence - resumable threaded conversations +- structured local memory with optional prompt injection - streamed assistant output - host-defined tools with approval gates - persona-aware agent behavior @@ -77,6 +78,7 @@ let stream = try await runtime.sendMessage( | Configurable thinking level | Yes | | Web search toggle (`enableWebSearch`) | Yes | | Built-in request retry/backoff | Yes (configurable) | +| Structured local memory layer | Yes | | Text + image input | Yes | | Assistant image attachment rendering | Yes | | Video/audio input attachments | Not yet | @@ -177,6 +179,94 @@ let stream = try await runtime.sendMessage( Custom tools can also return image URLs via `ToolResultContent.image(URL)`, and `CodexKit` attempts to hydrate those into assistant image attachments for chat rendering. +## Memory Layer + +`CodexKit` includes a generic memory layer for app-authored records. The SDK owns storage, retrieval, ranking, and optional prompt injection. Your app still decides what to remember and when to write it. + +```swift +let memoryURL = FileManager.default.urls( + for: .applicationSupportDirectory, + in: .userDomainMask +).first! + .appendingPathComponent("CodexKit/memory.sqlite") + +let memoryStore = try SQLiteMemoryStore(url: memoryURL) + +try await memoryStore.upsert( + MemoryRecord( + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Eleanor remembers being overruled on the trade bill.", + evidence: ["She warned the player twice before being ignored."], + importance: 0.9, + tags: ["trade", "advisor"] + ), + dedupeKey: "trade-bill-overruled-eleanor" +) + +let runtime = try AgentRuntime(configuration: .init( + authProvider: try ChatGPTAuthProvider( + method: .deviceCode, + deviceCodePresenter: deviceCodeCoordinator + ), + secureStore: KeychainSessionSecureStore( + service: "CodexKit.ChatGPTSession", + account: "main" + ), + backend: CodexResponsesBackend(), + approvalPresenter: approvalInbox, + stateStore: FileRuntimeStateStore( + url: FileManager.default.urls( + for: .applicationSupportDirectory, + in: .userDomainMask + ).first! + .appendingPathComponent("CodexKit/runtime-state.json") + ), + memory: .init(store: memoryStore) +)) + +let thread = try await runtime.createThread( + title: "Press Chat", + memoryContext: AgentMemoryContext( + namespace: "oval-office", + scopes: ["actor:eleanor_price", "thread:press"] + ) +) +``` + +Per-turn memory can be narrowed, expanded, replaced, or disabled with `MemorySelection`: + +```swift +let stream = try await runtime.sendMessage( + UserMessageRequest( + text: "How should Eleanor frame this rebuttal?", + memorySelection: MemorySelection( + mode: .append, + scopes: ["world:public"], + tags: ["trade"] + ) + ), + in: thread.id +) +``` + +For debugging and tooling, memory stores also support direct inspection: + +```swift +let stored = try await memoryStore.record( + id: "some-memory-id", + namespace: "oval-office" +) +let records = try await memoryStore.list( + namespace: "oval-office", + scopes: ["actor:eleanor_price"], + includeArchived: true, + limit: 20 +) +let diagnostics = try await memoryStore.diagnostics(namespace: "oval-office") +``` + ## Pinned And Dynamic Personas `CodexKit` supports layered persona precedence: diff --git a/Sources/CodexKit/Memory/InMemoryMemoryStore.swift b/Sources/CodexKit/Memory/InMemoryMemoryStore.swift index ad605f6..4b3ea0f 100644 --- a/Sources/CodexKit/Memory/InMemoryMemoryStore.swift +++ b/Sources/CodexKit/Memory/InMemoryMemoryStore.swift @@ -83,6 +83,51 @@ public actor InMemoryMemoryStore: MemoryStoring { ) } + public func record( + id: String, + namespace: String + ) async throws -> MemoryRecord? { + try MemoryQueryEngine.validateNamespace(namespace) + return recordsByNamespace[namespace, default: [:]][id] + } + + public func list(_ query: MemoryRecordListQuery) async throws -> [MemoryRecord] { + try MemoryQueryEngine.validateNamespace(query.namespace) + return recordsByNamespace[query.namespace, default: [:]] + .values + .filter { record in + if !query.includeArchived, record.status == .archived { + return false + } + if !query.scopes.isEmpty, !query.scopes.contains(record.scope) { + return false + } + if !query.kinds.isEmpty, !query.kinds.contains(record.kind) { + return false + } + return true + } + .sorted { + if $0.effectiveDate == $1.effectiveDate { + return $0.id < $1.id + } + return $0.effectiveDate > $1.effectiveDate + } + .prefix(query.limit ?? .max) + .map { $0 } + } + + public func diagnostics(namespace: String) async throws -> MemoryStoreDiagnostics { + try MemoryQueryEngine.validateNamespace(namespace) + let records = Array(recordsByNamespace[namespace, default: [:]].values) + return diagnostics( + namespace: namespace, + implementation: "in_memory", + schemaVersion: nil, + records: records + ) + } + public func compact(_ request: MemoryCompactionRequest) async throws { try MemoryQueryEngine.validateNamespace(request.replacement.namespace) var working = recordsByNamespace @@ -154,4 +199,33 @@ public actor InMemoryMemoryStore: MemoryStoring { recordsByNamespace[namespace] = namespaceRecords return expiredIDs.count } + + private func diagnostics( + namespace: String, + implementation: String, + schemaVersion: Int?, + records: [MemoryRecord] + ) -> MemoryStoreDiagnostics { + var countsByScope: [MemoryScope: Int] = [:] + var countsByKind: [String: Int] = [:] + + for record in records { + countsByScope[record.scope, default: 0] += 1 + countsByKind[record.kind, default: 0] += 1 + } + + let activeRecords = records.filter { $0.status == .active }.count + let archivedRecords = records.count - activeRecords + + return MemoryStoreDiagnostics( + namespace: namespace, + implementation: implementation, + schemaVersion: schemaVersion, + totalRecords: records.count, + activeRecords: activeRecords, + archivedRecords: archivedRecords, + countsByScope: countsByScope, + countsByKind: countsByKind + ) + } } diff --git a/Sources/CodexKit/Memory/MemoryModels.swift b/Sources/CodexKit/Memory/MemoryModels.swift index 0bf25ba..c79584d 100644 --- a/Sources/CodexKit/Memory/MemoryModels.swift +++ b/Sources/CodexKit/Memory/MemoryModels.swift @@ -4,6 +4,7 @@ public enum MemoryStoreError: Error, LocalizedError, Equatable, Sendable { case invalidNamespace case duplicateRecordID(String) case duplicateDedupeKey(String) + case unsupportedSchemaVersion(Int) public var errorDescription: String? { switch self { @@ -13,6 +14,8 @@ public enum MemoryStoreError: Error, LocalizedError, Equatable, Sendable { return "A memory record with id \(id) already exists." case let .duplicateDedupeKey(key): return "A memory record with dedupe key \(key) already exists." + case let .unsupportedSchemaVersion(version): + return "The memory store schema version \(version) is newer than this SDK supports." } } } @@ -255,6 +258,59 @@ public struct MemoryCompactionRequest: Codable, Hashable, Sendable { } } +public struct MemoryRecordListQuery: Codable, Hashable, Sendable { + public var namespace: String + public var scopes: [MemoryScope] + public var kinds: [String] + public var includeArchived: Bool + public var limit: Int? + + public init( + namespace: String, + scopes: [MemoryScope] = [], + kinds: [String] = [], + includeArchived: Bool = false, + limit: Int? = nil + ) { + self.namespace = namespace + self.scopes = scopes + self.kinds = kinds + self.includeArchived = includeArchived + self.limit = limit + } +} + +public struct MemoryStoreDiagnostics: Codable, Hashable, Sendable { + public var namespace: String + public var implementation: String + public var schemaVersion: Int? + public var totalRecords: Int + public var activeRecords: Int + public var archivedRecords: Int + public var countsByScope: [MemoryScope: Int] + public var countsByKind: [String: Int] + + public init( + namespace: String, + implementation: String, + schemaVersion: Int?, + totalRecords: Int, + activeRecords: Int, + archivedRecords: Int, + countsByScope: [MemoryScope: Int], + countsByKind: [String: Int] + ) { + self.namespace = namespace + self.implementation = implementation + self.schemaVersion = schemaVersion + self.totalRecords = totalRecords + self.activeRecords = activeRecords + self.archivedRecords = archivedRecords + self.countsByScope = countsByScope + self.countsByKind = countsByKind + } +} + public struct AgentMemoryContext: Codable, Hashable, Sendable { public var namespace: String public var scopes: [MemoryScope] @@ -340,6 +396,16 @@ public protocol MemoryPromptRendering: Sendable { func render(result: MemoryQueryResult, budget: MemoryReadBudget) -> String } +public enum MemoryObservationEvent: Sendable { + case queryStarted(MemoryQuery) + case querySucceeded(query: MemoryQuery, result: MemoryQueryResult) + case queryFailed(query: MemoryQuery, message: String) +} + +public protocol MemoryObserving: Sendable { + func handle(event: MemoryObservationEvent) async +} + public struct DefaultMemoryPromptRenderer: MemoryPromptRendering, Sendable { public init() {} @@ -359,16 +425,19 @@ public struct AgentMemoryConfiguration: Sendable { public let defaultRanking: MemoryRankingWeights public let defaultReadBudget: MemoryReadBudget public let promptRenderer: any MemoryPromptRendering + public let observer: (any MemoryObserving)? public init( store: any MemoryStoring, defaultRanking: MemoryRankingWeights = .default, defaultReadBudget: MemoryReadBudget = .runtimeDefault, - promptRenderer: any MemoryPromptRendering = DefaultMemoryPromptRenderer() + promptRenderer: any MemoryPromptRendering = DefaultMemoryPromptRenderer(), + observer: (any MemoryObserving)? = nil ) { self.store = store self.defaultRanking = defaultRanking self.defaultReadBudget = defaultReadBudget self.promptRenderer = promptRenderer + self.observer = observer } } diff --git a/Sources/CodexKit/Memory/MemoryStore.swift b/Sources/CodexKit/Memory/MemoryStore.swift index 0e791cd..97bf0ca 100644 --- a/Sources/CodexKit/Memory/MemoryStore.swift +++ b/Sources/CodexKit/Memory/MemoryStore.swift @@ -5,6 +5,9 @@ public protocol MemoryStoring: Sendable { func putMany(_ records: [MemoryRecord]) async throws func upsert(_ record: MemoryRecord, dedupeKey: String) async throws func query(_ query: MemoryQuery) async throws -> MemoryQueryResult + func record(id: String, namespace: String) async throws -> MemoryRecord? + func list(_ query: MemoryRecordListQuery) async throws -> [MemoryRecord] + func diagnostics(namespace: String) async throws -> MemoryStoreDiagnostics func compact(_ request: MemoryCompactionRequest) async throws func archive(ids: [String], namespace: String) async throws func delete(ids: [String], namespace: String) async throws @@ -17,6 +20,24 @@ public protocol MemoryStoring: Sendable { } public extension MemoryStoring { + func list( + namespace: String, + scopes: [MemoryScope] = [], + kinds: [String] = [], + includeArchived: Bool = false, + limit: Int? = nil + ) async throws -> [MemoryRecord] { + try await list( + MemoryRecordListQuery( + namespace: namespace, + scopes: scopes, + kinds: kinds, + includeArchived: includeArchived, + limit: limit + ) + ) + } + @discardableResult func pruneExpired(namespace: String) async throws -> Int { try await pruneExpired(now: Date(), namespace: namespace) diff --git a/Sources/CodexKit/Memory/SQLiteMemoryStore.swift b/Sources/CodexKit/Memory/SQLiteMemoryStore.swift index 2bff1f2..ee4fa15 100644 --- a/Sources/CodexKit/Memory/SQLiteMemoryStore.swift +++ b/Sources/CodexKit/Memory/SQLiteMemoryStore.swift @@ -2,6 +2,7 @@ import Foundation import SQLite3 public actor SQLiteMemoryStore: MemoryStoring { + private static let currentSchemaVersion = 1 private let url: URL private nonisolated(unsafe) var database: OpaquePointer? private let encoder = JSONEncoder() @@ -10,7 +11,7 @@ public actor SQLiteMemoryStore: MemoryStoring { public init(url: URL) throws { self.url = url self.database = try Self.openDatabase(at: url) - try Self.createSchemaIfNeeded(in: database) + try Self.migrateIfNeeded(in: database) } deinit { @@ -75,6 +76,54 @@ public actor SQLiteMemoryStore: MemoryStoring { ) } + public func record( + id: String, + namespace: String + ) async throws -> MemoryRecord? { + try MemoryQueryEngine.validateNamespace(namespace) + return try loadRecords(namespace: namespace).first { $0.id == id } + } + + public func list(_ query: MemoryRecordListQuery) async throws -> [MemoryRecord] { + try MemoryQueryEngine.validateNamespace(query.namespace) + return try loadRecords(namespace: query.namespace) + .filter { record in + if !query.includeArchived, record.status == .archived { + return false + } + if !query.scopes.isEmpty, !query.scopes.contains(record.scope) { + return false + } + if !query.kinds.isEmpty, !query.kinds.contains(record.kind) { + return false + } + return true + } + .sorted { + if $0.effectiveDate == $1.effectiveDate { + return $0.id < $1.id + } + return $0.effectiveDate > $1.effectiveDate + } + .prefix(query.limit ?? .max) + .map { $0 } + } + + public func diagnostics(namespace: String) async throws -> MemoryStoreDiagnostics { + try MemoryQueryEngine.validateNamespace(namespace) + let records = try loadRecords(namespace: namespace) + return MemoryStoreDiagnostics( + namespace: namespace, + implementation: "sqlite", + schemaVersion: try Self.schemaVersion(in: database), + totalRecords: records.count, + activeRecords: records.filter { $0.status == .active }.count, + archivedRecords: records.filter { $0.status == .archived }.count, + countsByScope: Dictionary(grouping: records, by: \.scope).mapValues(\.count), + countsByKind: Dictionary(grouping: records, by: \.kind).mapValues(\.count) + ) + } + public func compact(_ request: MemoryCompactionRequest) async throws { try MemoryQueryEngine.validateNamespace(request.replacement.namespace) try transaction { @@ -156,6 +205,18 @@ public actor SQLiteMemoryStore: MemoryStoring { return database } + private static func migrateIfNeeded(in database: OpaquePointer?) throws { + let existingVersion = try schemaVersion(in: database) + if existingVersion > currentSchemaVersion { + throw MemoryStoreError.unsupportedSchemaVersion(existingVersion) + } + + try createSchemaIfNeeded(in: database) + if existingVersion < currentSchemaVersion { + try setSchemaVersion(currentSchemaVersion, in: database) + } + } + private static func createSchemaIfNeeded(in database: OpaquePointer?) throws { let schema = """ CREATE TABLE IF NOT EXISTS memory_records ( @@ -212,6 +273,22 @@ public actor SQLiteMemoryStore: MemoryStoring { try execSQL(database, schema) } + private static func schemaVersion(in database: OpaquePointer?) throws -> Int { + let statement = try prepareSQL(database, "PRAGMA user_version;") + defer { sqlite3_finalize(statement) } + guard sqlite3_step(statement) == SQLITE_ROW else { + throw sqliteError(database, message: "Failed to read SQLite schema version.") + } + return Int(sqlite3_column_int(statement, 0)) + } + + private static func setSchemaVersion( + _ version: Int, + in database: OpaquePointer? + ) throws { + try execSQL(database, "PRAGMA user_version = \(version);") + } + private func ensureRecordIDAvailable( _ id: String, namespace: String @@ -490,15 +567,7 @@ public actor SQLiteMemoryStore: MemoryStoring { } private func prepare(_ sql: String) throws -> OpaquePointer? { - guard let database else { - throw sqliteError(message: "SQLite database is unavailable.") - } - var statement: OpaquePointer? - let result = sqlite3_prepare_v2(database, sql, -1, &statement, nil) - guard result == SQLITE_OK else { - throw sqliteError(message: "Failed to prepare SQLite statement.") - } - return statement + try prepareSQL(database, sql) } private func exec( @@ -624,3 +693,26 @@ private func execSQL( ) } } + +private func prepareSQL( + _ database: OpaquePointer?, + _ sql: String +) throws -> OpaquePointer? { + guard let database else { + throw NSError( + domain: "CodexKit.SQLiteMemoryStore", + code: 0, + userInfo: [NSLocalizedDescriptionKey: "SQLite database is unavailable."] + ) + } + var statement: OpaquePointer? + let result = sqlite3_prepare_v2(database, sql, -1, &statement, nil) + guard result == SQLITE_OK else { + throw NSError( + domain: "CodexKit.SQLiteMemoryStore", + code: Int(result), + userInfo: [NSLocalizedDescriptionKey: "Failed to prepare SQLite statement."] + ) + } + return statement +} diff --git a/Sources/CodexKit/Runtime/AgentRuntime.swift b/Sources/CodexKit/Runtime/AgentRuntime.swift index 7b4decd..a5b034c 100644 --- a/Sources/CodexKit/Runtime/AgentRuntime.swift +++ b/Sources/CodexKit/Runtime/AgentRuntime.swift @@ -819,7 +819,27 @@ public actor AgentRuntime { return nil } - return try? await memoryConfiguration.store.query(query) + if let observer = memoryConfiguration.observer { + await observer.handle(event: .queryStarted(query)) + } + + do { + let result = try await memoryConfiguration.store.query(query) + if let observer = memoryConfiguration.observer { + await observer.handle(event: .querySucceeded(query: query, result: result)) + } + return result + } catch { + if let observer = memoryConfiguration.observer { + await observer.handle( + event: .queryFailed( + query: query, + message: error.localizedDescription + ) + ) + } + return nil + } } private func resolvedMemoryQuery( diff --git a/Tests/CodexKitTests/AgentRuntimeTests.swift b/Tests/CodexKitTests/AgentRuntimeTests.swift index 630129e..a0b3c4f 100644 --- a/Tests/CodexKitTests/AgentRuntimeTests.swift +++ b/Tests/CodexKitTests/AgentRuntimeTests.swift @@ -758,6 +758,63 @@ final class AgentRuntimeTests: XCTestCase { XCTAssertFalse(resolved.contains("Relevant Memory:")) } + func testRuntimeReportsMemoryObservationEvents() async throws { + let backend = InMemoryAgentBackend( + baseInstructions: "Base host instructions." + ) + let store = InMemoryMemoryStore(initialRecords: [ + MemoryRecord( + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Observed memory." + ), + ]) + let observer = RecordingMemoryObserver() + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: backend, + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init( + store: store, + observer: observer + ) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Observed", + memoryContext: AgentMemoryContext( + namespace: "oval-office", + scopes: ["actor:eleanor_price"] + ) + ) + + let stream = try await runtime.sendMessage( + UserMessageRequest(text: "Use memory."), + in: thread.id + ) + for try await _ in stream {} + + let events = await observer.events() + XCTAssertEqual(events.count, 2) + guard case let .queryStarted(startedQuery) = events[0] else { + return XCTFail("Expected queryStarted event.") + } + XCTAssertEqual(startedQuery.namespace, "oval-office") + guard case let .querySucceeded(_, result) = events[1] else { + return XCTFail("Expected querySucceeded event.") + } + XCTAssertEqual(result.matches.count, 1) + } + func testResolvedInstructionsPreviewThrowsForMissingThread() async throws { let runtime = try AgentRuntime(configuration: .init( authProvider: DemoChatGPTAuthProvider(), @@ -1393,6 +1450,27 @@ private actor ThrowingMemoryStore: MemoryStoring { ) } + func record(id: String, namespace: String) async throws -> MemoryRecord? { + nil + } + + func list(_ query: MemoryRecordListQuery) async throws -> [MemoryRecord] { + [] + } + + func diagnostics(namespace: String) async throws -> MemoryStoreDiagnostics { + .init( + namespace: namespace, + implementation: "throwing", + schemaVersion: nil, + totalRecords: 0, + activeRecords: 0, + archivedRecords: 0, + countsByScope: [:], + countsByKind: [:] + ) + } + func compact(_ request: MemoryCompactionRequest) async throws {} func archive(ids: [String], namespace: String) async throws {} @@ -1403,3 +1481,15 @@ private actor ThrowingMemoryStore: MemoryStoring { 0 } } + +private actor RecordingMemoryObserver: MemoryObserving { + private var observedEvents: [MemoryObservationEvent] = [] + + func handle(event: MemoryObservationEvent) async { + observedEvents.append(event) + } + + func events() -> [MemoryObservationEvent] { + observedEvents + } +} diff --git a/Tests/CodexKitTests/MemoryStoreTests.swift b/Tests/CodexKitTests/MemoryStoreTests.swift index 6e09ef2..b66655a 100644 --- a/Tests/CodexKitTests/MemoryStoreTests.swift +++ b/Tests/CodexKitTests/MemoryStoreTests.swift @@ -1,5 +1,6 @@ import CodexKit import Foundation +import SQLite3 import XCTest final class MemoryStoreTests: XCTestCase { @@ -35,6 +36,35 @@ final class MemoryStoreTests: XCTestCase { XCTAssertEqual(result.matches.map(\.record.id), [record.id]) XCTAssertGreaterThan(result.matches[0].explanation.textScore, 0) + + let diagnostics = try await reloaded.diagnostics(namespace: "oval-office") + XCTAssertEqual(diagnostics.implementation, "sqlite") + XCTAssertEqual(diagnostics.schemaVersion, 1) + } + + func testSQLiteStoreRejectsUnsupportedFutureSchemaVersion() async throws { + let url = temporarySQLiteURL() + defer { try? FileManager.default.removeItem(at: url) } + + var database: OpaquePointer? + XCTAssertEqual( + sqlite3_open_v2( + url.path, + &database, + SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX, + nil + ), + SQLITE_OK + ) + XCTAssertEqual(sqlite3_exec(database, "PRAGMA user_version = 999;", nil, nil, nil), SQLITE_OK) + sqlite3_close(database) + + XCTAssertThrowsError(try SQLiteMemoryStore(url: url)) { error in + XCTAssertEqual( + error as? MemoryStoreError, + .unsupportedSchemaVersion(999) + ) + } } func testPutManyIsAtomicWhenDuplicateIDIsPresent() async throws { @@ -125,6 +155,49 @@ final class MemoryStoreTests: XCTestCase { XCTAssertEqual(result.matches[0].record.dedupeKey, "press-quote-17") } + func testStoreInspectionAPIsReturnRecordsAndDiagnostics() async throws { + let store = InMemoryMemoryStore() + try await store.putMany([ + MemoryRecord( + id: "active-memory", + namespace: "oval-office", + scope: "actor:eleanor_price", + kind: "grievance", + summary: "Active memory." + ), + MemoryRecord( + id: "archived-memory", + namespace: "oval-office", + scope: "world:press", + kind: "summary", + summary: "Archived memory.", + status: .archived + ), + ]) + + let fetched = try await store.record( + id: "active-memory", + namespace: "oval-office" + ) + XCTAssertEqual(fetched?.kind, "grievance") + + let listed = try await store.list( + namespace: "oval-office", + includeArchived: true, + limit: 10 + ) + XCTAssertEqual(listed.map(\.id).sorted(), ["active-memory", "archived-memory"]) + + let diagnostics = try await store.diagnostics(namespace: "oval-office") + XCTAssertEqual(diagnostics.implementation, "in_memory") + XCTAssertNil(diagnostics.schemaVersion) + XCTAssertEqual(diagnostics.totalRecords, 2) + XCTAssertEqual(diagnostics.activeRecords, 1) + XCTAssertEqual(diagnostics.archivedRecords, 1) + XCTAssertEqual(diagnostics.countsByScope["actor:eleanor_price"], 1) + XCTAssertEqual(diagnostics.countsByKind["summary"], 1) + } + func testQueryFiltersRankingAndCharacterBudget() async throws { let store = InMemoryMemoryStore() try await store.putMany([ From 682589fb31dc6e3837fec76a461c1a9acdc7d85e Mon Sep 17 00:00:00 2001 From: Timothy Zelinsky Date: Sun, 22 Mar 2026 18:03:04 +1100 Subject: [PATCH 3/4] Add layered memory capture APIs and demos --- .../project.pbxproj | 12 + .../AssistantRuntimeDemoApp.swift | 14 + .../Shared/AgentDemoRuntimeFactory.swift | 43 +- .../Shared/AgentDemoViewModel+Memory.swift | 205 +++++++++ .../Shared/AgentDemoViewModel.swift | 46 ++ .../Shared/DemoMemoryExamples.swift | 53 +++ .../Shared/MemoryDemoView.swift | 393 ++++++++++++++++++ README.md | 147 ++++++- .../CodexKit/Memory/InMemoryMemoryStore.swift | 5 +- Sources/CodexKit/Memory/MemoryCapture.swift | 128 ++++++ Sources/CodexKit/Memory/MemoryModels.swift | 29 +- .../CodexKit/Memory/MemoryQueryEngine.swift | 37 +- Sources/CodexKit/Memory/MemoryWriter.swift | 311 ++++++++++++++ .../CodexKit/Memory/SQLiteMemoryStore.swift | 3 +- Sources/CodexKit/Runtime/AgentModels.swift | 7 + Sources/CodexKit/Runtime/AgentRuntime.swift | 257 ++++++++++++ Tests/CodexKitTests/AgentRuntimeTests.swift | 295 ++++++++++--- Tests/CodexKitTests/MemoryStoreTests.swift | 303 ++++++++++---- 18 files changed, 2119 insertions(+), 169 deletions(-) create mode 100644 DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel+Memory.swift create mode 100644 DemoApp/AssistantRuntimeDemoApp/Shared/DemoMemoryExamples.swift create mode 100644 DemoApp/AssistantRuntimeDemoApp/Shared/MemoryDemoView.swift create mode 100644 Sources/CodexKit/Memory/MemoryCapture.swift create mode 100644 Sources/CodexKit/Memory/MemoryWriter.swift diff --git a/DemoApp/AssistantRuntimeDemoApp.xcodeproj/project.pbxproj b/DemoApp/AssistantRuntimeDemoApp.xcodeproj/project.pbxproj index f9a976a..39fd65e 100644 --- a/DemoApp/AssistantRuntimeDemoApp.xcodeproj/project.pbxproj +++ b/DemoApp/AssistantRuntimeDemoApp.xcodeproj/project.pbxproj @@ -22,6 +22,9 @@ 1A2B3C4D5E6F70000000000D /* DemoUIComponents.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2A2B3C4D5E6F70000000000D /* DemoUIComponents.swift */; }; 1A2B3C4D5E6F70000000000E /* StructuredOutputDemoView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2A2B3C4D5E6F70000000000E /* StructuredOutputDemoView.swift */; }; 1A2B3C4D5E6F70000000000F /* ThreadDetailView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2A2B3C4D5E6F70000000000F /* ThreadDetailView.swift */; }; + 1A2B3C4D5E6F700000000010 /* DemoMemoryExamples.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2A2B3C4D5E6F700000000010 /* DemoMemoryExamples.swift */; }; + 1A2B3C4D5E6F700000000011 /* AgentDemoViewModel+Memory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2A2B3C4D5E6F700000000011 /* AgentDemoViewModel+Memory.swift */; }; + 1A2B3C4D5E6F700000000012 /* MemoryDemoView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2A2B3C4D5E6F700000000012 /* MemoryDemoView.swift */; }; 7482123BC63AC10F104DE092 /* AssistantRuntimeDemoApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5A6999E6475919476E726E8C /* AssistantRuntimeDemoApp.swift */; }; 84726927B752451499D9257F /* Foundation.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 906A95007C8ECB92CFC2CE15 /* Foundation.framework */; }; B060448C6464C41789B56EED /* AgentDemoView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 3CA22585116A120BA97F76B8 /* AgentDemoView.swift */; }; @@ -47,6 +50,9 @@ 2A2B3C4D5E6F70000000000D /* DemoUIComponents.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = DemoUIComponents.swift; sourceTree = ""; }; 2A2B3C4D5E6F70000000000E /* StructuredOutputDemoView.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = StructuredOutputDemoView.swift; sourceTree = ""; }; 2A2B3C4D5E6F70000000000F /* ThreadDetailView.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = ThreadDetailView.swift; sourceTree = ""; }; + 2A2B3C4D5E6F700000000010 /* DemoMemoryExamples.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = DemoMemoryExamples.swift; sourceTree = ""; }; + 2A2B3C4D5E6F700000000011 /* AgentDemoViewModel+Memory.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = "AgentDemoViewModel+Memory.swift"; sourceTree = ""; }; + 2A2B3C4D5E6F700000000012 /* MemoryDemoView.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = MemoryDemoView.swift; sourceTree = ""; }; 2481147A958D00EB4A70C928 /* AgentDemoViewModel.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = AgentDemoViewModel.swift; sourceTree = ""; }; 3CA22585116A120BA97F76B8 /* AgentDemoView.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = AgentDemoView.swift; sourceTree = ""; }; 5A6999E6475919476E726E8C /* AssistantRuntimeDemoApp.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = AssistantRuntimeDemoApp.swift; sourceTree = ""; }; @@ -137,6 +143,9 @@ 2A2B3C4D5E6F70000000000D /* DemoUIComponents.swift */, 2A2B3C4D5E6F70000000000E /* StructuredOutputDemoView.swift */, 2A2B3C4D5E6F70000000000F /* ThreadDetailView.swift */, + 2A2B3C4D5E6F700000000010 /* DemoMemoryExamples.swift */, + 2A2B3C4D5E6F700000000011 /* AgentDemoViewModel+Memory.swift */, + 2A2B3C4D5E6F700000000012 /* MemoryDemoView.swift */, ); name = Shared; path = Shared; @@ -228,6 +237,9 @@ 1A2B3C4D5E6F70000000000D /* DemoUIComponents.swift in Sources */, 1A2B3C4D5E6F70000000000E /* StructuredOutputDemoView.swift in Sources */, 1A2B3C4D5E6F70000000000F /* ThreadDetailView.swift in Sources */, + 1A2B3C4D5E6F700000000010 /* DemoMemoryExamples.swift in Sources */, + 1A2B3C4D5E6F700000000011 /* AgentDemoViewModel+Memory.swift in Sources */, + 1A2B3C4D5E6F700000000012 /* MemoryDemoView.swift in Sources */, 7482123BC63AC10F104DE092 /* AssistantRuntimeDemoApp.swift in Sources */, BB4F38E64D1EBBB3821AC4E3 /* AgentDemoRuntimeFactory.swift in Sources */, B060448C6464C41789B56EED /* AgentDemoView.swift in Sources */, diff --git a/DemoApp/AssistantRuntimeDemoApp/AssistantRuntimeDemoApp.swift b/DemoApp/AssistantRuntimeDemoApp/AssistantRuntimeDemoApp.swift index 1e9b9a4..ef42ed6 100644 --- a/DemoApp/AssistantRuntimeDemoApp/AssistantRuntimeDemoApp.swift +++ b/DemoApp/AssistantRuntimeDemoApp/AssistantRuntimeDemoApp.swift @@ -3,6 +3,7 @@ import SwiftUI enum DemoTab: Hashable { case assistant case structuredOutput + case memory case healthCoach } @@ -45,6 +46,19 @@ struct AssistantRuntimeDemoApp: App { Label("Structured", systemImage: "square.stack.3d.up") } + NavigationStack { + MemoryDemoView( + viewModel: viewModel, + selectedTab: $selectedTab + ) + .navigationTitle("Memory") + .navigationBarTitleDisplayMode(.inline) + } + .tag(DemoTab.memory) + .tabItem { + Label("Memory", systemImage: "brain") + } + NavigationStack { HealthCoachView(viewModel: viewModel) .navigationTitle("Health Coach") diff --git a/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoRuntimeFactory.swift b/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoRuntimeFactory.swift index b41a659..f2ab5fc 100644 --- a/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoRuntimeFactory.swift +++ b/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoRuntimeFactory.swift @@ -102,7 +102,21 @@ enum AgentDemoRuntimeFactory { ) ), approvalPresenter: approvalInbox, - stateStore: FileRuntimeStateStore(url: stateURL ?? defaultStateURL()) + stateStore: FileRuntimeStateStore(url: stateURL ?? defaultStateURL()), + memory: .init( + store: try! SQLiteMemoryStore(url: defaultMemoryURL()), + automaticCapturePolicy: .init( + source: .lastTurn, + options: .init( + defaults: .init( + namespace: DemoMemoryExamples.namespace, + kind: "preference", + tags: ["demo", "auto-capture"] + ), + maxMemories: 2 + ) + ) + ) )) } #endif @@ -129,7 +143,21 @@ enum AgentDemoRuntimeFactory { ) ), approvalPresenter: NonInteractiveApprovalPresenter(), - stateStore: FileRuntimeStateStore(url: defaultStateURL()) + stateStore: FileRuntimeStateStore(url: defaultStateURL()), + memory: .init( + store: try! SQLiteMemoryStore(url: defaultMemoryURL()), + automaticCapturePolicy: .init( + source: .lastTurn, + options: .init( + defaults: .init( + namespace: DemoMemoryExamples.namespace, + kind: "preference", + tags: ["demo", "auto-capture"] + ), + maxMemories: 2 + ) + ) + ) )) } @@ -143,6 +171,17 @@ enum AgentDemoRuntimeFactory { .appendingPathComponent("AssistantRuntimeDemoApp", isDirectory: true) .appendingPathComponent("runtime-state.json") } + + static func defaultMemoryURL() -> URL { + let baseDirectory = FileManager.default.urls( + for: .applicationSupportDirectory, + in: .userDomainMask + ).first ?? URL(fileURLWithPath: NSTemporaryDirectory()) + + return baseDirectory + .appendingPathComponent("AssistantRuntimeDemoApp", isDirectory: true) + .appendingPathComponent("memory.sqlite") + } } private struct NonInteractiveApprovalPresenter: ApprovalPresenting { diff --git a/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel+Memory.swift b/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel+Memory.swift new file mode 100644 index 0000000..6e3d4bf --- /dev/null +++ b/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel+Memory.swift @@ -0,0 +1,205 @@ +import CodexKit +import Foundation + +@MainActor +extension AgentDemoViewModel { + func runAutomaticPolicyMemoryDemo() async { + guard session != nil else { + lastError = "Sign in before running policy-based memory capture." + return + } + guard !isRunningMemoryDemo else { + return + } + + isRunningMemoryDemo = true + lastError = nil + defer { + isRunningMemoryDemo = false + } + + do { + let thread = try await runtime.createThread( + title: "Memory Demo: Automatic Policy", + memoryContext: AgentMemoryContext( + namespace: DemoMemoryExamples.namespace, + scopes: [DemoMemoryExamples.healthCoachScope], + kinds: ["preference"] + ) + ) + _ = try await runtime.sendMessage( + UserMessageRequest(text: DemoMemoryExamples.automaticPolicyPrompt), + in: thread.id + ) + + let store = try SQLiteMemoryStore(url: AgentDemoRuntimeFactory.defaultMemoryURL()) + let result = try await store.query( + MemoryQuery( + namespace: DemoMemoryExamples.namespace, + scopes: [DemoMemoryExamples.healthCoachScope], + text: "direct blunt steps", + limit: 4, + maxCharacters: 800 + ) + ) + + automaticPolicyMemoryResult = AutomaticPolicyMemoryDemoResult( + threadID: thread.id, + threadTitle: thread.title ?? "Memory Demo: Automatic Policy", + prompt: DemoMemoryExamples.automaticPolicyPrompt, + records: result.matches.map(\.record) + ) + threads = await runtime.threads() + } catch { + lastError = error.localizedDescription + } + } + + func runAutomaticMemoryDemo() async { + guard session != nil else { + lastError = "Sign in before running automatic memory capture." + return + } + guard !isRunningMemoryDemo else { + return + } + + isRunningMemoryDemo = true + lastError = nil + defer { + isRunningMemoryDemo = false + } + + do { + let thread = try await runtime.createThread( + title: "Memory Demo: Automatic Capture", + memoryContext: DemoMemoryExamples.previewContext + ) + let capture = try await runtime.captureMemories( + from: .text(DemoMemoryExamples.automaticCaptureTranscript), + for: thread.id, + options: .init( + defaults: DemoMemoryExamples.guidedDefaults, + maxMemories: 3 + ) + ) + automaticMemoryResult = AutomaticMemoryDemoResult( + threadID: thread.id, + threadTitle: thread.title ?? "Memory Demo: Automatic Capture", + capture: capture + ) + threads = await runtime.threads() + } catch { + lastError = error.localizedDescription + } + } + + func runGuidedMemoryDemo() async { + guard !isRunningMemoryDemo else { + return + } + + isRunningMemoryDemo = true + lastError = nil + defer { + isRunningMemoryDemo = false + } + + do { + let writer = try await runtime.memoryWriter(defaults: DemoMemoryExamples.guidedDefaults) + let record = try await writer.upsert(DemoMemoryExamples.guidedDraft) + let diagnostics = try await writer.diagnostics() + + guidedMemoryResult = GuidedMemoryDemoResult( + record: record, + diagnostics: diagnostics + ) + } catch { + lastError = error.localizedDescription + } + } + + func runRawMemoryDemo() async { + guard !isRunningMemoryDemo else { + return + } + + isRunningMemoryDemo = true + lastError = nil + defer { + isRunningMemoryDemo = false + } + + do { + let store = try SQLiteMemoryStore(url: AgentDemoRuntimeFactory.defaultMemoryURL()) + try await store.upsert( + DemoMemoryExamples.rawRecord, + dedupeKey: DemoMemoryExamples.rawRecord.dedupeKey ?? DemoMemoryExamples.rawRecord.id + ) + let diagnostics = try await store.diagnostics(namespace: DemoMemoryExamples.namespace) + + rawMemoryResult = RawMemoryDemoResult( + record: DemoMemoryExamples.rawRecord, + diagnostics: diagnostics + ) + } catch { + lastError = error.localizedDescription + } + } + + func runMemoryPreviewDemo() async { + guard !isRunningMemoryDemo else { + return + } + + isRunningMemoryDemo = true + lastError = nil + defer { + isRunningMemoryDemo = false + } + + do { + let result: MemoryQueryResult + var previewThreadID: String? + var previewThreadTitle: String? + + if session != nil { + let thread = try await runtime.createThread( + title: "Memory Demo: Prompt Injection", + memoryContext: DemoMemoryExamples.previewContext + ) + previewThreadID = thread.id + previewThreadTitle = thread.title + result = try await runtime.memoryQueryPreview( + for: thread.id, + request: UserMessageRequest(text: DemoMemoryExamples.previewRequestText) + ) ?? MemoryQueryResult(matches: [], truncated: false) + threads = await runtime.threads() + } else { + let store = try SQLiteMemoryStore(url: AgentDemoRuntimeFactory.defaultMemoryURL()) + result = try await store.query( + MemoryQuery( + namespace: DemoMemoryExamples.namespace, + scopes: DemoMemoryExamples.previewContext.scopes, + text: DemoMemoryExamples.previewRequestText, + limit: DemoMemoryExamples.previewBudget.maxItems, + maxCharacters: DemoMemoryExamples.previewBudget.maxCharacters + ) + ) + } + + memoryPreviewResult = MemoryPreviewDemoResult( + threadID: previewThreadID, + threadTitle: previewThreadTitle, + requestText: DemoMemoryExamples.previewRequestText, + result: result, + renderedPrompt: DefaultMemoryPromptRenderer().render( + result: result, + budget: DemoMemoryExamples.previewBudget + ) + ) + } catch { + lastError = error.localizedDescription + } + } +} diff --git a/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel.swift b/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel.swift index 2adfe17..4a6094d 100644 --- a/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel.swift +++ b/DemoApp/AssistantRuntimeDemoApp/Shared/AgentDemoViewModel.swift @@ -39,6 +39,37 @@ struct StructuredOutputDemoImportResult: Sendable { let summary: StructuredImportedContentSummary } +struct GuidedMemoryDemoResult: Sendable { + let record: MemoryRecord + let diagnostics: MemoryStoreDiagnostics +} + +struct RawMemoryDemoResult: Sendable { + let record: MemoryRecord + let diagnostics: MemoryStoreDiagnostics +} + +struct MemoryPreviewDemoResult: Sendable { + let threadID: String? + let threadTitle: String? + let requestText: String + let result: MemoryQueryResult + let renderedPrompt: String +} + +struct AutomaticMemoryDemoResult: Sendable { + let threadID: String + let threadTitle: String + let capture: MemoryCaptureResult +} + +struct AutomaticPolicyMemoryDemoResult: Sendable { + let threadID: String + let threadTitle: String + let prompt: String + let records: [MemoryRecord] +} + @MainActor @Observable final class AgentDemoViewModel: @unchecked Sendable { @@ -106,6 +137,12 @@ final class AgentDemoViewModel: @unchecked Sendable { var isRunningStructuredOutputDemo = false var structuredShippingReplyResult: StructuredOutputDemoDraftResult? var structuredImportedSummaryResult: StructuredOutputDemoImportResult? + var isRunningMemoryDemo = false + var automaticMemoryResult: AutomaticMemoryDemoResult? + var automaticPolicyMemoryResult: AutomaticPolicyMemoryDemoResult? + var guidedMemoryResult: GuidedMemoryDemoResult? + var rawMemoryResult: RawMemoryDemoResult? + var memoryPreviewResult: MemoryPreviewDemoResult? var healthKitAuthorized = false var notificationAuthorized = false var isRefreshingHealthCoach = false @@ -433,6 +470,15 @@ final class AgentDemoViewModel: @unchecked Sendable { lastResolvedInstructionsThreadTitle = nil isRunningSkillPolicyProbe = false skillPolicyProbeResult = nil + isRunningStructuredOutputDemo = false + structuredShippingReplyResult = nil + structuredImportedSummaryResult = nil + isRunningMemoryDemo = false + automaticMemoryResult = nil + automaticPolicyMemoryResult = nil + guidedMemoryResult = nil + rawMemoryResult = nil + memoryPreviewResult = nil activeThreadID = nil healthCoachThreadID = nil healthCoachFeedback = "Set a step goal, then start moving." diff --git a/DemoApp/AssistantRuntimeDemoApp/Shared/DemoMemoryExamples.swift b/DemoApp/AssistantRuntimeDemoApp/Shared/DemoMemoryExamples.swift new file mode 100644 index 0000000..1b65ec5 --- /dev/null +++ b/DemoApp/AssistantRuntimeDemoApp/Shared/DemoMemoryExamples.swift @@ -0,0 +1,53 @@ +import CodexKit +import Foundation + +enum DemoMemoryExamples { + static let namespace = "demo-assistant" + static let healthCoachScope: MemoryScope = "feature:health-coach" + static let travelPlannerScope: MemoryScope = "feature:travel-planner" + + static let guidedDraft = MemoryDraft( + summary: "Health Coach should use blunt accountability when the user is behind on steps.", + evidence: ["The user responds better to direct push language than soft encouragement."], + importance: 0.95, + tags: ["tone", "steps"], + relatedIDs: ["health-goal-10000"], + dedupeKey: "health-coach-direct-accountability" + ) + + static let rawRecord = MemoryRecord( + namespace: namespace, + scope: travelPlannerScope, + kind: "preference", + summary: "Travel Planner should keep itineraries compact, walkable, and transit-aware.", + evidence: ["The demo works best when travel plans stay practical for mobile users."], + importance: 0.82, + tags: ["travel", "logistics"], + relatedIDs: ["travel-style-compact"], + dedupeKey: "travel-planner-compact-itinerary" + ) + + static let previewRequestText = "How should the assistant behave for the health coach and travel planner demos?" + static let previewBudget = MemoryReadBudget(maxItems: 4, maxCharacters: 500) + static let automaticCaptureTranscript = """ + User: I respond better when the health coach is direct and blunt if I am behind on steps. + Assistant: Understood. I will stop soft-pedaling reminders and push harder when you are off pace. + User: For travel planning, keep itineraries compact, walkable, and transit-aware. I hate sprawling plans. + """ + static let automaticPolicyPrompt = "When I fall behind on steps, drop the sugar-coating and be direct with me." + + static let guidedDefaults = MemoryWriterDefaults( + namespace: namespace, + scope: healthCoachScope, + kind: "preference", + tags: ["demo", "guided"], + relatedIDs: ["demo-memory"], + status: .active + ) + + static let previewContext = AgentMemoryContext( + namespace: namespace, + scopes: [healthCoachScope, travelPlannerScope], + readBudget: previewBudget + ) +} diff --git a/DemoApp/AssistantRuntimeDemoApp/Shared/MemoryDemoView.swift b/DemoApp/AssistantRuntimeDemoApp/Shared/MemoryDemoView.swift new file mode 100644 index 0000000..769511f --- /dev/null +++ b/DemoApp/AssistantRuntimeDemoApp/Shared/MemoryDemoView.swift @@ -0,0 +1,393 @@ +import CodexKit +import Foundation +import SwiftUI + +@available(iOS 17.0, macOS 14.0, *) +struct MemoryDemoView: View { + @State var viewModel: AgentDemoViewModel + @Binding var selectedTab: DemoTab + + init(viewModel: AgentDemoViewModel, selectedTab: Binding) { + _viewModel = State(initialValue: viewModel) + _selectedTab = selectedTab + } + + var body: some View { + ScrollView { + VStack(alignment: .leading, spacing: 18) { + overviewCard + automaticPolicyCard + automaticCaptureCard + guidedAuthoringCard + lowLevelAuthoringCard + retrievalCard + } + .padding(20) + } + } +} + +@available(iOS 17.0, macOS 14.0, *) +private extension MemoryDemoView { + var overviewCard: some View { + DemoSectionCard { + Text("Memory Layer") + .font(.title2.weight(.semibold)) + + Text("This demo shows the full stack: high-level automatic capture policies, mid-level guided authoring with `MemoryWriter`, and low-level raw `MemoryRecord` control.") + .font(.subheadline) + .foregroundStyle(.secondary) + + if viewModel.session != nil { + Label("Signed in: prompt-injection preview can also create a live thread.", systemImage: "checkmark.seal.fill") + .font(.subheadline) + .foregroundStyle(.secondary) + } else { + Label("You can save and query memory without signing in. Sign in if you want the preview to open a real thread.", systemImage: "info.circle") + .font(.subheadline) + .foregroundStyle(.secondary) + } + } + } + + var automaticPolicyCard: some View { + DemoSectionCard { + Text("Automatic Policy") + .font(.headline) + + Text("This is the high-level option. The runtime is configured to capture memory automatically after successful turns for threads with memory context.") + .font(.subheadline) + .foregroundStyle(.secondary) + + sampleCard( + title: "Normal chat prompt", + body: DemoMemoryExamples.automaticPolicyPrompt + ) + + DemoActionTile( + title: viewModel.isRunningMemoryDemo ? "Running Automatic Policy..." : "Run Auto Memory Policy Demo", + subtitle: "Sends a normal assistant message, then lets the runtime extract and save durable memory on its own.", + systemImage: "bolt.badge.automatic", + isProminent: true, + isDisabled: viewModel.session == nil || viewModel.isRunningMemoryDemo + ) { + Task { + await viewModel.runAutomaticPolicyMemoryDemo() + } + } + + if let result = viewModel.automaticPolicyMemoryResult { + VStack(alignment: .leading, spacing: 10) { + Text("Captured after a regular turn") + .font(.subheadline.weight(.semibold)) + + ForEach(Array(result.records.enumerated()), id: \.offset) { _, record in + memoryRecordCard( + title: record.summary, + record: record, + diagnostics: MemoryStoreDiagnostics( + namespace: record.namespace, + implementation: "automatic_policy", + schemaVersion: nil, + totalRecords: result.records.count, + activeRecords: result.records.count, + archivedRecords: 0, + countsByScope: [:], + countsByKind: [:] + ) + ) + } + + Button("Open Policy Thread In Assistant") { + Task { + await viewModel.activateThread(id: result.threadID) + selectedTab = .assistant + } + } + .buttonStyle(.bordered) + } + .padding(14) + .background( + RoundedRectangle(cornerRadius: 16, style: .continuous) + .fill(Color.accentColor.opacity(0.10)) + ) + } + } + } + + var guidedAuthoringCard: some View { + DemoSectionCard { + Text("Guided Authoring") + .font(.headline) + + Text("Uses `MemoryWriter` defaults so the app only supplies the memory payload. Namespace, scope, kind, and tags are resolved automatically.") + .font(.subheadline) + .foregroundStyle(.secondary) + + sampleCard( + title: "MemoryDraft", + body: """ + summary: \(DemoMemoryExamples.guidedDraft.summary) + evidence: \(DemoMemoryExamples.guidedDraft.evidence.first ?? "") + dedupeKey: \(DemoMemoryExamples.guidedDraft.dedupeKey ?? "none") + """ + ) + + DemoActionTile( + title: viewModel.isRunningMemoryDemo ? "Saving Guided Memory..." : "Save Guided Health Coach Memory", + subtitle: "Calls `runtime.memoryWriter(defaults:)` and resolves a real `MemoryRecord`.", + systemImage: "wand.and.stars", + isProminent: true, + isDisabled: viewModel.isRunningMemoryDemo + ) { + Task { + await viewModel.runGuidedMemoryDemo() + } + } + + if let result = viewModel.guidedMemoryResult { + memoryRecordCard( + title: "Resolved record", + record: result.record, + diagnostics: result.diagnostics + ) + } + } + } + + var automaticCaptureCard: some View { + DemoSectionCard { + Text("Automatic Capture") + .font(.headline) + + Text("Lets the runtime extract durable memory candidates from a transcript, then writes them through the same memory store automatically.") + .font(.subheadline) + .foregroundStyle(.secondary) + + sampleCard( + title: "Sample transcript", + body: DemoMemoryExamples.automaticCaptureTranscript + ) + + DemoActionTile( + title: viewModel.isRunningMemoryDemo ? "Capturing Memory..." : "Auto-Capture Memory From Transcript", + subtitle: "Uses structured output under the hood, then saves the extracted drafts with `MemoryWriter`.", + systemImage: "sparkles.rectangle.stack", + isProminent: true, + isDisabled: viewModel.session == nil || viewModel.isRunningMemoryDemo + ) { + Task { + await viewModel.runAutomaticMemoryDemo() + } + } + + if let result = viewModel.automaticMemoryResult { + VStack(alignment: .leading, spacing: 10) { + Text("Captured \(result.capture.records.count) memory record\(result.capture.records.count == 1 ? "" : "s")") + .font(.subheadline.weight(.semibold)) + + ForEach(Array(result.capture.records.enumerated()), id: \.offset) { _, record in + memoryRecordCard( + title: record.summary, + record: record, + diagnostics: MemoryStoreDiagnostics( + namespace: record.namespace, + implementation: "runtime_capture", + schemaVersion: nil, + totalRecords: result.capture.records.count, + activeRecords: result.capture.records.count, + archivedRecords: 0, + countsByScope: [:], + countsByKind: [:] + ) + ) + } + + Button("Open Capture Thread In Assistant") { + Task { + await viewModel.activateThread(id: result.threadID) + selectedTab = .assistant + } + } + .buttonStyle(.bordered) + } + .padding(14) + .background( + RoundedRectangle(cornerRadius: 16, style: .continuous) + .fill(Color.accentColor.opacity(0.10)) + ) + } + } + } + + var lowLevelAuthoringCard: some View { + DemoSectionCard { + Text("Raw Store Control") + .font(.headline) + + Text("Writes a full `MemoryRecord` directly into the SQLite store. This is the low-level escape hatch for apps that want exact IDs, scopes, compaction flows, or custom pipelines.") + .font(.subheadline) + .foregroundStyle(.secondary) + + sampleCard( + title: "MemoryRecord", + body: """ + namespace: \(DemoMemoryExamples.rawRecord.namespace) + scope: \(DemoMemoryExamples.rawRecord.scope.rawValue) + kind: \(DemoMemoryExamples.rawRecord.kind) + summary: \(DemoMemoryExamples.rawRecord.summary) + """ + ) + + DemoActionTile( + title: viewModel.isRunningMemoryDemo ? "Saving Raw Memory..." : "Save Raw Travel Planner Memory", + subtitle: "Calls `SQLiteMemoryStore.upsert(...)` directly with a fully specified record.", + systemImage: "shippingbox.circle", + isDisabled: viewModel.isRunningMemoryDemo + ) { + Task { + await viewModel.runRawMemoryDemo() + } + } + + if let result = viewModel.rawMemoryResult { + memoryRecordCard( + title: "Stored raw record", + record: result.record, + diagnostics: result.diagnostics + ) + } + } + } + + var retrievalCard: some View { + DemoSectionCard { + Text("Retrieval Preview") + .font(.headline) + + Text("Queries the stored memories and renders the exact prompt block that would be injected into a turn. If you are signed in, it also creates a live thread with matching memory context.") + .font(.subheadline) + .foregroundStyle(.secondary) + + sampleCard( + title: "Preview request", + body: DemoMemoryExamples.previewRequestText + ) + + DemoActionTile( + title: viewModel.isRunningMemoryDemo ? "Building Preview..." : "Preview Memory Injection", + subtitle: viewModel.session == nil + ? "Runs a local memory query and renders the prompt block." + : "Runs `memoryQueryPreview` and prepares a real thread you can open in Assistant.", + systemImage: "brain.head.profile", + isDisabled: viewModel.isRunningMemoryDemo + ) { + Task { + await viewModel.runMemoryPreviewDemo() + } + } + + if let result = viewModel.memoryPreviewResult { + VStack(alignment: .leading, spacing: 12) { + Text("Matched \(result.result.matches.count) record\(result.result.matches.count == 1 ? "" : "s")") + .font(.subheadline.weight(.semibold)) + + if result.result.matches.isEmpty { + Text("No memory matched yet. Save one of the demo records above and run the preview again.") + .font(.subheadline) + .foregroundStyle(.secondary) + } else { + ForEach(Array(result.result.matches.enumerated()), id: \.offset) { _, match in + VStack(alignment: .leading, spacing: 6) { + Text("\(match.record.scope.rawValue) • \(match.record.kind)") + .font(.caption.weight(.semibold)) + .foregroundStyle(.secondary) + Text(match.record.summary) + .font(.body) + Text("score \(match.explanation.totalScore.formatted(.number.precision(.fractionLength(2))))") + .font(.caption) + .foregroundStyle(.secondary) + } + .frame(maxWidth: .infinity, alignment: .leading) + } + } + + sampleCard( + title: "Rendered prompt block", + body: result.renderedPrompt.isEmpty ? "(empty)" : result.renderedPrompt + ) + + if let threadID = result.threadID { + Button("Open Preview Thread In Assistant") { + Task { + await viewModel.activateThread(id: threadID) + selectedTab = .assistant + } + } + .buttonStyle(.bordered) + } + } + .padding(14) + .background( + RoundedRectangle(cornerRadius: 16, style: .continuous) + .fill(Color.accentColor.opacity(0.10)) + ) + } + } + } + + func sampleCard(title: String, body: String) -> some View { + VStack(alignment: .leading, spacing: 8) { + Text(title) + .font(.caption.weight(.semibold)) + .foregroundStyle(.secondary) + + Text(body) + .font(.callout.monospaced()) + .frame(maxWidth: .infinity, alignment: .leading) + .textSelection(.enabled) + } + .padding(14) + .background( + RoundedRectangle(cornerRadius: 16, style: .continuous) + .fill(Color.primary.opacity(0.04)) + ) + } + + func memoryRecordCard( + title: String, + record: MemoryRecord, + diagnostics: MemoryStoreDiagnostics + ) -> some View { + VStack(alignment: .leading, spacing: 10) { + Text(title) + .font(.subheadline.weight(.semibold)) + + Group { + detailRow(label: "Namespace", value: record.namespace) + detailRow(label: "Scope", value: record.scope.rawValue) + detailRow(label: "Kind", value: record.kind) + detailRow(label: "Summary", value: record.summary) + detailRow(label: "Tags", value: record.tags.joined(separator: ", ")) + detailRow(label: "Dedupe Key", value: record.dedupeKey ?? "none") + detailRow(label: "Store Count", value: "\(diagnostics.totalRecords)") + } + } + .padding(14) + .background( + RoundedRectangle(cornerRadius: 16, style: .continuous) + .fill(Color.primary.opacity(0.04)) + ) + } + + func detailRow(label: String, value: String) -> some View { + VStack(alignment: .leading, spacing: 4) { + Text(label) + .font(.caption.weight(.semibold)) + .foregroundStyle(.secondary) + Text(value.isEmpty ? "none" : value) + .font(.body) + .frame(maxWidth: .infinity, alignment: .leading) + } + } +} diff --git a/README.md b/README.md index 7dbd972..a69ff05 100644 --- a/README.md +++ b/README.md @@ -231,7 +231,110 @@ Custom tools can also return image URLs via `ToolResultContent.image(URL)`, and ## Memory Layer -`CodexKit` includes a generic memory layer for app-authored records. The SDK owns storage, retrieval, ranking, and optional prompt injection. Your app still decides what to remember and when to write it. +`CodexKit` now supports three memory layers: + +- high-level automatic capture policies for apps that want the runtime to extract memory after successful turns +- a guided `MemoryWriter` layer that resolves defaults into concrete records +- the raw `MemoryRecord` / `MemoryStoring` APIs for apps that want exact control + +The SDK owns storage, retrieval, ranking, and optional prompt injection. Your app can choose how automatic or explicit memory authoring should be. + +High-level automatic capture looks like this: + +```swift +let runtime = try AgentRuntime(configuration: .init( + authProvider: try ChatGPTAuthProvider( + method: .deviceCode, + deviceCodePresenter: deviceCodeCoordinator + ), + secureStore: KeychainSessionSecureStore( + service: "CodexKit.ChatGPTSession", + account: "demo" + ), + backend: CodexResponsesBackend( + configuration: .init(model: "gpt-5.4") + ), + approvalPresenter: approvalPresenter, + stateStore: FileRuntimeStateStore(url: stateURL), + memory: .init( + store: try SQLiteMemoryStore(url: memoryURL), + automaticCapturePolicy: .init( + source: .lastTurn, + options: .init( + defaults: .init( + namespace: "demo-assistant", + kind: "preference" + ), + maxMemories: 2 + ) + ) + ) +)) + +let thread = try await runtime.createThread( + title: "Health Coach", + memoryContext: .init( + namespace: "demo-assistant", + scopes: ["feature:health-coach"] + ) +) + +_ = try await runtime.sendMessage( + UserMessageRequest(text: "Be direct with me when I fall behind on steps."), + in: thread.id +) +``` + +Mid-level guided authoring looks like this: + +```swift +let writer = try await runtime.memoryWriter( + defaults: .init( + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + tags: ["steps", "tone"] + ) +) + +let record = try await writer.upsert( + MemoryDraft( + summary: "Health Coach should use direct accountability when the user is behind on steps.", + evidence: ["The user responds better to blunt reminders than soft encouragement."], + importance: 0.9, + dedupeKey: "health-coach-direct-accountability" + ) +) +``` + +If you want the SDK to capture memory for you, `AgentRuntime` can extract durable memory candidates from a thread or transcript and write them automatically: + +```swift +let thread = try await runtime.createThread( + title: "Health Coach", + memoryContext: .init( + namespace: "demo-assistant", + scopes: ["feature:health-coach"] + ) +) + +let result = try await runtime.captureMemories( + from: .threadHistory(maxMessages: 6), + for: thread.id, + options: .init( + defaults: .init( + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference" + ), + maxMemories: 3 + ) +) + +print(result.records.count) +``` + +If you want full control, the low-level store API is still there: ```swift let memoryURL = FileManager.default.urls( @@ -244,15 +347,15 @@ let memoryStore = try SQLiteMemoryStore(url: memoryURL) try await memoryStore.upsert( MemoryRecord( - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", - summary: "Eleanor remembers being overruled on the trade bill.", - evidence: ["She warned the player twice before being ignored."], + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + summary: "Health Coach should use direct accountability when the user is behind on steps.", + evidence: ["The user responds better to blunt coaching than soft encouragement."], importance: 0.9, - tags: ["trade", "advisor"] + tags: ["steps", "tone"] ), - dedupeKey: "trade-bill-overruled-eleanor" + dedupeKey: "health-coach-direct-accountability" ) let runtime = try AgentRuntime(configuration: .init( @@ -279,8 +382,8 @@ let runtime = try AgentRuntime(configuration: .init( let thread = try await runtime.createThread( title: "Press Chat", memoryContext: AgentMemoryContext( - namespace: "oval-office", - scopes: ["actor:eleanor_price", "thread:press"] + namespace: "demo-assistant", + scopes: ["feature:health-coach", "thread:daily-checkin"] ) ) ``` @@ -288,13 +391,13 @@ let thread = try await runtime.createThread( Per-turn memory can be narrowed, expanded, replaced, or disabled with `MemorySelection`: ```swift -let stream = try await runtime.sendMessage( +let reply = try await runtime.sendMessage( UserMessageRequest( - text: "How should Eleanor frame this rebuttal?", + text: "How should the health coach respond when the user is behind on steps?", memorySelection: MemorySelection( mode: .append, - scopes: ["world:public"], - tags: ["trade"] + scopes: ["feature:travel-planner"], + tags: ["steps"] ) ), in: thread.id @@ -306,17 +409,25 @@ For debugging and tooling, memory stores also support direct inspection: ```swift let stored = try await memoryStore.record( id: "some-memory-id", - namespace: "oval-office" + namespace: "demo-assistant" ) let records = try await memoryStore.list( - namespace: "oval-office", - scopes: ["actor:eleanor_price"], + namespace: "demo-assistant", + scopes: ["feature:health-coach"], includeArchived: true, limit: 20 ) -let diagnostics = try await memoryStore.diagnostics(namespace: "oval-office") +let diagnostics = try await memoryStore.diagnostics(namespace: "demo-assistant") ``` +The demo app now includes a dedicated `Memory` tab that shows: + +- high-level automatic capture after a normal turn +- mid-level automatic capture from transcript +- guided authoring with `MemoryWriter` +- raw record writes against the underlying store +- preview of the exact prompt block injected into a turn + ## Pinned And Dynamic Personas `CodexKit` supports layered persona precedence: diff --git a/Sources/CodexKit/Memory/InMemoryMemoryStore.swift b/Sources/CodexKit/Memory/InMemoryMemoryStore.swift index 4b3ea0f..ded560e 100644 --- a/Sources/CodexKit/Memory/InMemoryMemoryStore.swift +++ b/Sources/CodexKit/Memory/InMemoryMemoryStore.swift @@ -70,10 +70,11 @@ public actor InMemoryMemoryStore: MemoryStoring { let candidates = namespaceRecords.values.map { record in MemoryQueryEngine.Candidate( record: record, - rawTextScore: MemoryQueryEngine.defaultTextScore( + textScore: MemoryQueryEngine.defaultTextScore( for: record, queryText: query.text - ) + ), + textScoreOrdering: .higherIsBetter ) } diff --git a/Sources/CodexKit/Memory/MemoryCapture.swift b/Sources/CodexKit/Memory/MemoryCapture.swift new file mode 100644 index 0000000..f748b63 --- /dev/null +++ b/Sources/CodexKit/Memory/MemoryCapture.swift @@ -0,0 +1,128 @@ +import Foundation + +public enum MemoryCaptureSource: Sendable { + case threadHistory(maxMessages: Int = 8) + case messages([AgentMessage]) + case text(String) +} + +public struct MemoryCaptureOptions: Sendable { + public var defaults: MemoryWriterDefaults + public var maxMemories: Int + public var instructions: String? + + public init( + defaults: MemoryWriterDefaults = .init(), + maxMemories: Int = 3, + instructions: String? = nil + ) { + self.defaults = defaults + self.maxMemories = maxMemories + self.instructions = instructions + } +} + +public struct MemoryCaptureResult: Sendable { + public var sourceText: String + public var drafts: [MemoryDraft] + public var records: [MemoryRecord] + + public init( + sourceText: String, + drafts: [MemoryDraft], + records: [MemoryRecord] + ) { + self.sourceText = sourceText + self.drafts = drafts + self.records = records + } +} + +struct MemoryExtractionDraftResponse: Decodable, Sendable { + let memories: [MemoryExtractionDraft] + + static func responseFormat(maxMemories: Int) -> AgentStructuredOutputFormat { + AgentStructuredOutputFormat( + name: "memory_extraction_batch", + description: "Durable memory candidates extracted from app conversation history.", + schema: .object( + properties: [ + "memories": .array( + items: .object( + properties: [ + "summary": .string(), + "scope": .nullable(.string()), + "kind": .nullable(.string()), + "evidence": .array(items: .string()), + "importance": .number, + "tags": .array(items: .string()), + "relatedIDs": .array(items: .string()), + "dedupeKey": .nullable(.string()), + ], + required: ["summary", "evidence", "importance", "tags", "relatedIDs", "dedupeKey"], + additionalProperties: false + ) + ), + ], + required: ["memories"], + additionalProperties: false + ), + strict: true + ) + } + + static func prompt( + sourceText: String, + maxMemories: Int + ) -> String { + """ + Extract up to \(maxMemories) durable memory records from the source conversation below. + + Save only information worth carrying into future turns: + - stable user preferences + - durable project constraints + - recurring behavioral guidance + - persistent goals or important facts + + Do not save: + - greetings or filler + - one-off requests + - temporary status updates + - details that are too vague to be useful later + + If nothing is worth saving, return an empty `memories` array. + + Source conversation: + \(sourceText) + """ + } + + static let instructions = """ + You extract durable app memory from conversations for future turns. + Return compact memory candidates only. Prefer specific, reusable facts over temporary requests. + """ +} + +struct MemoryExtractionDraft: Decodable, Sendable { + let summary: String + let scope: String? + let kind: String? + let evidence: [String] + let importance: Double + let tags: [String] + let relatedIDs: [String] + let dedupeKey: String? + + var memoryDraft: MemoryDraft { + MemoryDraft( + scope: scope.map(MemoryScope.init(rawValue:)), + kind: kind, + summary: summary, + evidence: evidence, + importance: importance, + tags: tags, + relatedIDs: relatedIDs, + dedupeKey: dedupeKey + ) + } +} diff --git a/Sources/CodexKit/Memory/MemoryModels.swift b/Sources/CodexKit/Memory/MemoryModels.swift index c79584d..c4e6212 100644 --- a/Sources/CodexKit/Memory/MemoryModels.swift +++ b/Sources/CodexKit/Memory/MemoryModels.swift @@ -400,12 +400,36 @@ public enum MemoryObservationEvent: Sendable { case queryStarted(MemoryQuery) case querySucceeded(query: MemoryQuery, result: MemoryQueryResult) case queryFailed(query: MemoryQuery, message: String) + case captureStarted(threadID: String, sourceDescription: String) + case captureSucceeded(threadID: String, result: MemoryCaptureResult) + case captureFailed(threadID: String, message: String) } public protocol MemoryObserving: Sendable { func handle(event: MemoryObservationEvent) async } +public enum MemoryAutomaticCaptureSource: Hashable, Sendable { + case lastTurn + case threadHistory(maxMessages: Int = 8) +} + +public struct MemoryAutomaticCapturePolicy: Sendable { + public var source: MemoryAutomaticCaptureSource + public var options: MemoryCaptureOptions + public var requiresThreadMemoryContext: Bool + + public init( + source: MemoryAutomaticCaptureSource = .lastTurn, + options: MemoryCaptureOptions = .init(), + requiresThreadMemoryContext: Bool = true + ) { + self.source = source + self.options = options + self.requiresThreadMemoryContext = requiresThreadMemoryContext + } +} + public struct DefaultMemoryPromptRenderer: MemoryPromptRendering, Sendable { public init() {} @@ -426,18 +450,21 @@ public struct AgentMemoryConfiguration: Sendable { public let defaultReadBudget: MemoryReadBudget public let promptRenderer: any MemoryPromptRendering public let observer: (any MemoryObserving)? + public let automaticCapturePolicy: MemoryAutomaticCapturePolicy? public init( store: any MemoryStoring, defaultRanking: MemoryRankingWeights = .default, defaultReadBudget: MemoryReadBudget = .runtimeDefault, promptRenderer: any MemoryPromptRendering = DefaultMemoryPromptRenderer(), - observer: (any MemoryObserving)? = nil + observer: (any MemoryObserving)? = nil, + automaticCapturePolicy: MemoryAutomaticCapturePolicy? = nil ) { self.store = store self.defaultRanking = defaultRanking self.defaultReadBudget = defaultReadBudget self.promptRenderer = promptRenderer self.observer = observer + self.automaticCapturePolicy = automaticCapturePolicy } } diff --git a/Sources/CodexKit/Memory/MemoryQueryEngine.swift b/Sources/CodexKit/Memory/MemoryQueryEngine.swift index 47a0509..c21dd4c 100644 --- a/Sources/CodexKit/Memory/MemoryQueryEngine.swift +++ b/Sources/CodexKit/Memory/MemoryQueryEngine.swift @@ -1,9 +1,15 @@ import Foundation internal enum MemoryQueryEngine { + internal enum TextScoreOrdering { + case higherIsBetter + case lowerIsBetter + } + internal struct Candidate { let record: MemoryRecord - let rawTextScore: Double? + let textScore: Double? + let textScoreOrdering: TextScoreOrdering } private struct ScoredCandidate { @@ -82,15 +88,10 @@ internal enum MemoryQueryEngine { break } - let nextCount = characterCount + candidate.characterCost - if !selected.isEmpty, nextCount > query.maxCharacters { + let nextCount = characterCount + candidate.characterCost + (selected.isEmpty ? 0 : 1) + if nextCount > query.maxCharacters { truncated = true - break - } - - if selected.isEmpty, candidate.characterCost > query.maxCharacters { - truncated = true - break + continue } selected.append(candidate.match) @@ -226,7 +227,7 @@ internal enum MemoryQueryEngine { private static func normalizedTextScores( from candidates: [Candidate] ) -> [String: Double] { - let rawScores = candidates.compactMap(\.rawTextScore) + let rawScores = candidates.compactMap(\.textScore) guard let maxScore = rawScores.max(), let minScore = rawScores.min() else { @@ -234,15 +235,25 @@ internal enum MemoryQueryEngine { } return candidates.reduce(into: [String: Double]()) { partial, candidate in - guard let rawScore = candidate.rawTextScore else { + guard let rawScore = candidate.textScore else { partial[candidate.record.id] = 0 return } if maxScore == minScore { - partial[candidate.record.id] = 1 + switch candidate.textScoreOrdering { + case .higherIsBetter: + partial[candidate.record.id] = rawScore > 0 ? 1 : 0 + case .lowerIsBetter: + partial[candidate.record.id] = 1 + } } else { - partial[candidate.record.id] = clamp((maxScore - rawScore) / (maxScore - minScore)) + switch candidate.textScoreOrdering { + case .higherIsBetter: + partial[candidate.record.id] = clamp((rawScore - minScore) / (maxScore - minScore)) + case .lowerIsBetter: + partial[candidate.record.id] = clamp((maxScore - rawScore) / (maxScore - minScore)) + } } } } diff --git a/Sources/CodexKit/Memory/MemoryWriter.swift b/Sources/CodexKit/Memory/MemoryWriter.swift new file mode 100644 index 0000000..3fdff01 --- /dev/null +++ b/Sources/CodexKit/Memory/MemoryWriter.swift @@ -0,0 +1,311 @@ +import Foundation + +public enum MemoryAuthoringError: Error, LocalizedError, Equatable, Sendable { + case missingNamespace + case missingScope + case missingKind + case missingDedupeKey + + public var errorDescription: String? { + switch self { + case .missingNamespace: + return "A memory namespace is required before writing memory." + case .missingScope: + return "A memory scope is required before writing memory." + case .missingKind: + return "A memory kind is required before writing memory." + case .missingDedupeKey: + return "A dedupe key is required for memory upserts." + } + } +} + +public struct MemoryDraft: Codable, Hashable, Sendable { + public var id: String? + public var namespace: String? + public var scope: MemoryScope? + public var kind: String? + public var summary: String + public var evidence: [String] + public var importance: Double? + public var createdAt: Date? + public var observedAt: Date? + public var expiresAt: Date? + public var expiresIn: TimeInterval? + public var tags: [String] + public var relatedIDs: [String] + public var dedupeKey: String? + public var isPinned: Bool? + public var attributes: JSONValue? + public var status: MemoryRecordStatus? + + public init( + id: String? = nil, + namespace: String? = nil, + scope: MemoryScope? = nil, + kind: String? = nil, + summary: String, + evidence: [String] = [], + importance: Double? = nil, + createdAt: Date? = nil, + observedAt: Date? = nil, + expiresAt: Date? = nil, + expiresIn: TimeInterval? = nil, + tags: [String] = [], + relatedIDs: [String] = [], + dedupeKey: String? = nil, + isPinned: Bool? = nil, + attributes: JSONValue? = nil, + status: MemoryRecordStatus? = nil + ) { + self.id = id + self.namespace = namespace + self.scope = scope + self.kind = kind + self.summary = summary + self.evidence = evidence + self.importance = importance + self.createdAt = createdAt + self.observedAt = observedAt + self.expiresAt = expiresAt + self.expiresIn = expiresIn + self.tags = tags + self.relatedIDs = relatedIDs + self.dedupeKey = dedupeKey + self.isPinned = isPinned + self.attributes = attributes + self.status = status + } +} + +public struct MemoryWriterDefaults: Codable, Hashable, Sendable { + public var namespace: String? + public var scope: MemoryScope? + public var kind: String? + public var importance: Double + public var tags: [String] + public var relatedIDs: [String] + public var isPinned: Bool + public var status: MemoryRecordStatus + + public init( + namespace: String? = nil, + scope: MemoryScope? = nil, + kind: String? = nil, + importance: Double = 0, + tags: [String] = [], + relatedIDs: [String] = [], + isPinned: Bool = false, + status: MemoryRecordStatus = .active + ) { + self.namespace = namespace + self.scope = scope + self.kind = kind + self.importance = importance + self.tags = tags + self.relatedIDs = relatedIDs + self.isPinned = isPinned + self.status = status + } + + public func fillingMissingValues(from inherited: MemoryWriterDefaults) -> MemoryWriterDefaults { + MemoryWriterDefaults( + namespace: namespace ?? inherited.namespace, + scope: scope ?? inherited.scope, + kind: kind ?? inherited.kind, + importance: importance, + tags: Self.uniqueMerged(values: inherited.tags, with: tags), + relatedIDs: Self.uniqueMerged(values: inherited.relatedIDs, with: relatedIDs), + isPinned: isPinned, + status: status + ) + } + + fileprivate static func uniqueMerged(values base: [String], with override: [String]) -> [String] { + var merged: [String] = [] + var seen = Set() + + for value in base + override { + guard seen.insert(value).inserted else { + continue + } + merged.append(value) + } + + return merged + } +} + +public actor MemoryWriter { + private let store: any MemoryStoring + public let defaults: MemoryWriterDefaults + + public init( + store: any MemoryStoring, + defaults: MemoryWriterDefaults = .init() + ) { + self.store = store + self.defaults = defaults + } + + public nonisolated func resolve( + _ draft: MemoryDraft, + now: Date = Date() + ) throws -> MemoryRecord { + let namespace = try resolvedNamespace(for: draft) + let scope = try resolvedScope(for: draft) + let kind = try resolvedKind(for: draft) + let createdAt = draft.createdAt ?? now + let baseExpiryDate = draft.observedAt ?? createdAt + let expiresAt = draft.expiresAt ?? draft.expiresIn.map { baseExpiryDate.addingTimeInterval($0) } + + return MemoryRecord( + id: draft.id ?? UUID().uuidString, + namespace: namespace, + scope: scope, + kind: kind, + summary: draft.summary, + evidence: draft.evidence, + importance: draft.importance ?? defaults.importance, + createdAt: createdAt, + observedAt: draft.observedAt, + expiresAt: expiresAt, + tags: MemoryWriterDefaults.uniqueMerged(values: defaults.tags, with: draft.tags), + relatedIDs: MemoryWriterDefaults.uniqueMerged(values: defaults.relatedIDs, with: draft.relatedIDs), + dedupeKey: draft.dedupeKey, + isPinned: draft.isPinned ?? defaults.isPinned, + attributes: draft.attributes, + status: draft.status ?? defaults.status + ) + } + + @discardableResult + public func put( + _ draft: MemoryDraft, + now: Date = Date() + ) async throws -> MemoryRecord { + let record = try resolve(draft, now: now) + try await store.put(record) + return record + } + + @discardableResult + public func putMany( + _ drafts: [MemoryDraft], + now: Date = Date() + ) async throws -> [MemoryRecord] { + let records = try drafts.map { try resolve($0, now: now) } + try await store.putMany(records) + return records + } + + @discardableResult + public func upsert( + _ draft: MemoryDraft, + now: Date = Date() + ) async throws -> MemoryRecord { + guard let dedupeKey = draft.dedupeKey?.trimmingCharacters(in: .whitespacesAndNewlines), + !dedupeKey.isEmpty + else { + throw MemoryAuthoringError.missingDedupeKey + } + + var draft = draft + draft.dedupeKey = dedupeKey + let record = try resolve(draft, now: now) + try await store.upsert(record, dedupeKey: dedupeKey) + return record + } + + @discardableResult + public func compact( + replacement draft: MemoryDraft, + sourceIDs: [String], + now: Date = Date() + ) async throws -> MemoryRecord { + let record = try resolve(draft, now: now) + try await store.compact( + MemoryCompactionRequest( + replacement: record, + sourceIDs: sourceIDs + ) + ) + return record + } + + public func diagnostics(namespace: String? = nil) async throws -> MemoryStoreDiagnostics { + try await store.diagnostics(namespace: try resolvedNamespace(namespace)) + } + + public func list( + namespace: String? = nil, + scopes: [MemoryScope] = [], + kinds: [String] = [], + includeArchived: Bool = false, + limit: Int? = nil + ) async throws -> [MemoryRecord] { + try await store.list( + namespace: try resolvedNamespace(namespace), + scopes: scopes, + kinds: kinds, + includeArchived: includeArchived, + limit: limit + ) + } + + public func archive( + ids: [String], + namespace: String? = nil + ) async throws { + try await store.archive(ids: ids, namespace: try resolvedNamespace(namespace)) + } + + public func delete( + ids: [String], + namespace: String? = nil + ) async throws { + try await store.delete(ids: ids, namespace: try resolvedNamespace(namespace)) + } + + @discardableResult + public func pruneExpired( + now: Date = Date(), + namespace: String? = nil + ) async throws -> Int { + try await store.pruneExpired( + now: now, + namespace: try resolvedNamespace(namespace) + ) + } + + private nonisolated func resolvedNamespace(for draft: MemoryDraft) throws -> String { + try resolvedNamespace(draft.namespace) + } + + private nonisolated func resolvedNamespace(_ namespace: String?) throws -> String { + guard let namespace = (namespace ?? defaults.namespace)?.trimmingCharacters(in: .whitespacesAndNewlines), + !namespace.isEmpty + else { + throw MemoryAuthoringError.missingNamespace + } + + return namespace + } + + private nonisolated func resolvedScope(for draft: MemoryDraft) throws -> MemoryScope { + guard let scope = draft.scope ?? defaults.scope else { + throw MemoryAuthoringError.missingScope + } + return scope + } + + private nonisolated func resolvedKind(for draft: MemoryDraft) throws -> String { + guard let kind = (draft.kind ?? defaults.kind)?.trimmingCharacters(in: .whitespacesAndNewlines), + !kind.isEmpty + else { + throw MemoryAuthoringError.missingKind + } + + return kind + } +} diff --git a/Sources/CodexKit/Memory/SQLiteMemoryStore.swift b/Sources/CodexKit/Memory/SQLiteMemoryStore.swift index ee4fa15..b7b6c10 100644 --- a/Sources/CodexKit/Memory/SQLiteMemoryStore.swift +++ b/Sources/CodexKit/Memory/SQLiteMemoryStore.swift @@ -66,7 +66,8 @@ public actor SQLiteMemoryStore: MemoryStoring { let candidates = records.map { record in MemoryQueryEngine.Candidate( record: record, - rawTextScore: rawScores[record.id] + textScore: rawScores[record.id], + textScoreOrdering: .lowerIsBetter ) } diff --git a/Sources/CodexKit/Runtime/AgentModels.swift b/Sources/CodexKit/Runtime/AgentModels.swift index 569f9a0..a7dcea3 100644 --- a/Sources/CodexKit/Runtime/AgentModels.swift +++ b/Sources/CodexKit/Runtime/AgentModels.swift @@ -28,6 +28,13 @@ public struct AgentRuntimeError: Error, LocalizedError, Equatable, Sendable { AgentRuntimeError(code: "unauthorized", message: message) } + public static func memoryNotConfigured() -> AgentRuntimeError { + AgentRuntimeError( + code: "memory_not_configured", + message: "This runtime was created without a memory store configuration." + ) + } + public static func invalidMessageContent() -> AgentRuntimeError { AgentRuntimeError( code: "invalid_message_content", diff --git a/Sources/CodexKit/Runtime/AgentRuntime.swift b/Sources/CodexKit/Runtime/AgentRuntime.swift index dab3e5d..abd5a13 100644 --- a/Sources/CodexKit/Runtime/AgentRuntime.swift +++ b/Sources/CodexKit/Runtime/AgentRuntime.swift @@ -431,6 +431,7 @@ public actor AgentRuntime { await self.consumeTurnStream( turnStream, for: threadID, + userMessage: userMessage, session: turnSession, resolvedTurnSkills: resolvedTurnSkills, continuation: continuation @@ -504,6 +505,7 @@ public actor AgentRuntime { private func consumeTurnStream( _ turnStream: any AgentTurnStreaming, for threadID: String, + userMessage: AgentMessage, session: ChatGPTSession, resolvedTurnSkills: ResolvedTurnSkills, continuation: AsyncThrowingStream.Continuation @@ -513,6 +515,7 @@ public actor AgentRuntime { } else { nil } + var assistantMessages: [AgentMessage] = [] do { for try await backendEvent in turnStream.events { @@ -531,6 +534,9 @@ public actor AgentRuntime { case let .assistantMessageCompleted(message): try await appendMessage(message) + if message.role == .assistant { + assistantMessages.append(message) + } continuation.yield(.messageCommitted(message)) case let .toolCallRequested(invocation): @@ -568,6 +574,11 @@ public actor AgentRuntime { } try await setThreadStatus(.idle, for: threadID) + await automaticallyCaptureMemoriesIfConfigured( + for: threadID, + userMessage: userMessage, + assistantMessages: assistantMessages + ) continuation.yield(.threadStatusChanged(threadID: threadID, status: .idle)) continuation.yield(.turnCompleted(summary)) } @@ -587,6 +598,71 @@ public actor AgentRuntime { } } + private func automaticallyCaptureMemoriesIfConfigured( + for threadID: String, + userMessage: AgentMessage, + assistantMessages: [AgentMessage] + ) async { + guard let memoryConfiguration, + let policy = memoryConfiguration.automaticCapturePolicy + else { + return + } + + guard let thread = thread(for: threadID) else { + return + } + + if policy.requiresThreadMemoryContext, thread.memoryContext == nil { + return + } + + let source: MemoryCaptureSource + let sourceDescription: String + switch policy.source { + case .lastTurn: + let turnMessages = [userMessage] + assistantMessages.filter { $0.threadID == threadID } + guard turnMessages.contains(where: { $0.role == .assistant }) else { + return + } + source = .messages(turnMessages) + sourceDescription = "last_turn" + + case let .threadHistory(maxMessages): + source = .threadHistory(maxMessages: maxMessages) + sourceDescription = "thread_history_\(max(1, maxMessages))" + } + + if let observer = memoryConfiguration.observer { + await observer.handle( + event: .captureStarted( + threadID: threadID, + sourceDescription: sourceDescription + ) + ) + } + + do { + let result = try await captureMemories( + from: source, + for: threadID, + options: policy.options + ) + if let observer = memoryConfiguration.observer { + await observer.handle(event: .captureSucceeded(threadID: threadID, result: result)) + } + } catch { + if let observer = memoryConfiguration.observer { + await observer.handle( + event: .captureFailed( + threadID: threadID, + message: error.localizedDescription + ) + ) + } + } + } + private func resolveToolInvocation( _ invocation: ToolInvocation, session: ChatGPTSession, @@ -695,6 +771,122 @@ public actor AgentRuntime { return thread.memoryContext } + public func memoryWriter( + defaults: MemoryWriterDefaults = .init() + ) throws -> MemoryWriter { + guard let memoryConfiguration else { + throw AgentRuntimeError.memoryNotConfigured() + } + + return MemoryWriter( + store: memoryConfiguration.store, + defaults: defaults + ) + } + + public func memoryWriter( + for threadID: String, + defaults: MemoryWriterDefaults = .init() + ) throws -> MemoryWriter { + guard let thread = thread(for: threadID) else { + throw AgentRuntimeError.threadNotFound(threadID) + } + + let inheritedDefaults: MemoryWriterDefaults + if let memoryContext = thread.memoryContext { + inheritedDefaults = MemoryWriterDefaults( + namespace: memoryContext.namespace, + scope: memoryContext.scopes.count == 1 ? memoryContext.scopes[0] : nil, + kind: memoryContext.kinds.count == 1 ? memoryContext.kinds[0] : nil, + tags: memoryContext.tags, + relatedIDs: memoryContext.relatedIDs + ) + } else { + inheritedDefaults = .init() + } + + return try memoryWriter( + defaults: defaults.fillingMissingValues(from: inheritedDefaults) + ) + } + + public func captureMemories( + from source: MemoryCaptureSource = .threadHistory(), + for threadID: String, + options: MemoryCaptureOptions = .init(), + decoder: JSONDecoder = JSONDecoder() + ) async throws -> MemoryCaptureResult { + guard let thread = thread(for: threadID) else { + throw AgentRuntimeError.threadNotFound(threadID) + } + + let sourceText = formattedMemoryCaptureSource( + source, + threadID: threadID + ) + guard !sourceText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + return MemoryCaptureResult( + sourceText: sourceText, + drafts: [], + records: [] + ) + } + + let writer = try memoryWriter( + for: threadID, + defaults: options.defaults + ) + let request = UserMessageRequest( + text: MemoryExtractionDraftResponse.prompt( + sourceText: sourceText, + maxMemories: max(1, options.maxMemories) + ) + ) + let session = try await sessionManager.requireSession() + let turnStart = try await beginTurnWithUnauthorizedRecovery( + thread: thread, + history: [], + message: request, + instructions: options.instructions ?? MemoryExtractionDraftResponse.instructions, + responseFormat: MemoryExtractionDraftResponse.responseFormat( + maxMemories: max(1, options.maxMemories) + ), + tools: [], + session: session + ) + let assistantMessage = try await collectFinalAssistantMessage( + from: turnStart.turnStream + ) + let payload = Data(assistantMessage.text.trimmingCharacters(in: .whitespacesAndNewlines).utf8) + + let extraction: MemoryExtractionDraftResponse + do { + extraction = try decoder.decode(MemoryExtractionDraftResponse.self, from: payload) + } catch { + throw AgentRuntimeError.structuredOutputDecodingFailed( + typeName: "MemoryExtractionDraftResponse", + underlyingMessage: error.localizedDescription + ) + } + + let drafts = extraction.memories.map(\.memoryDraft) + var records: [MemoryRecord] = [] + records.reserveCapacity(drafts.count) + for draft in drafts { + if draft.dedupeKey != nil { + records.append(try await writer.upsert(draft)) + } else { + records.append(try await writer.put(draft)) + } + } + + return MemoryCaptureResult( + sourceText: sourceText, + drafts: drafts, + records: records + ) + } + public func setSkillIDs( _ skillIDs: [String], for threadID: String @@ -798,6 +990,71 @@ public actor AgentRuntime { return latestAssistantMessage } + private func collectFinalAssistantMessage( + from turnStream: any AgentTurnStreaming + ) async throws -> AgentMessage { + var latestAssistantMessage: AgentMessage? + + for try await event in turnStream.events { + switch event { + case let .assistantMessageCompleted(message): + if message.role == .assistant { + latestAssistantMessage = message + } + + case let .toolCallRequested(invocation): + try await turnStream.submitToolResult( + .failure( + invocation: invocation, + message: "Automatic memory capture does not allow tool calls." + ), + for: invocation.id + ) + + default: + break + } + } + + guard let latestAssistantMessage else { + throw AgentRuntimeError.assistantResponseMissing() + } + + return latestAssistantMessage + } + + private func formattedMemoryCaptureSource( + _ source: MemoryCaptureSource, + threadID: String + ) -> String { + switch source { + case let .threadHistory(maxMessages): + let history = Array((state.messagesByThread[threadID] ?? []).suffix(max(1, maxMessages))) + return formattedMemoryTranscript(from: history) + + case let .messages(messages): + return formattedMemoryTranscript(from: messages) + + case let .text(text): + return text.trimmingCharacters(in: .whitespacesAndNewlines) + } + } + + private func formattedMemoryTranscript(from messages: [AgentMessage]) -> String { + messages + .map { message in + let role = message.role.rawValue.capitalized + let text = message.displayText.trimmingCharacters(in: .whitespacesAndNewlines) + + if text.isEmpty, !message.images.isEmpty { + return "\(role): [\(message.images.count) image attachment(s)]" + } + + return "\(role): \(text)" + } + .joined(separator: "\n") + } + private func persistState() async throws { try await stateStore.saveState(state) } diff --git a/Tests/CodexKitTests/AgentRuntimeTests.swift b/Tests/CodexKitTests/AgentRuntimeTests.swift index e6f570b..060e7fc 100644 --- a/Tests/CodexKitTests/AgentRuntimeTests.swift +++ b/Tests/CodexKitTests/AgentRuntimeTests.swift @@ -724,19 +724,19 @@ final class AgentRuntimeTests: XCTestCase { ) let store = InMemoryMemoryStore(initialRecords: [ MemoryRecord( - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", - summary: "Eleanor remembers being overruled on the trade bill.", - evidence: ["She warned the player twice before being ignored."], + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + summary: "Health Coach should use direct accountability when the user is behind on steps.", + evidence: ["The user responds better to blunt coaching than soft encouragement."], importance: 0.9, - tags: ["trade"] + tags: ["steps"] ), MemoryRecord( - namespace: "oval-office", - scope: "actor:sophia_ramirez", - kind: "grievance", - summary: "Sophia is focused on education messaging.", + namespace: "demo-assistant", + scope: "feature:travel-planner", + kind: "preference", + summary: "Travel Planner should keep itineraries compact and transit-aware.", importance: 0.8 ), ]) @@ -758,28 +758,27 @@ final class AgentRuntimeTests: XCTestCase { let thread = try await runtime.createThread( title: "Memory", memoryContext: AgentMemoryContext( - namespace: "oval-office", - scopes: ["actor:eleanor_price"] + namespace: "demo-assistant", + scopes: ["feature:health-coach"] ) ) let preview = try await runtime.memoryQueryPreview( for: thread.id, - request: UserMessageRequest(text: "What does Eleanor still remember about trade?") + request: UserMessageRequest(text: "How should the health coach respond when the user is behind on steps?") ) - XCTAssertEqual(preview?.matches.map(\.record.scope.rawValue), ["actor:eleanor_price"]) + XCTAssertEqual(preview?.matches.map(\.record.scope.rawValue), ["feature:health-coach"]) - let stream = try await runtime.sendMessage( - UserMessageRequest(text: "What does Eleanor still remember about trade?"), + _ = try await runtime.sendMessage( + UserMessageRequest(text: "How should the health coach respond when the user is behind on steps?"), in: thread.id ) - for try await _ in stream {} let instructions = await backend.receivedInstructions() let resolved = try XCTUnwrap(instructions.last) XCTAssertTrue(resolved.contains("Relevant Memory:")) - XCTAssertTrue(resolved.contains("Eleanor remembers being overruled on the trade bill.")) - XCTAssertFalse(resolved.contains("Sophia is focused on education messaging.")) + XCTAssertTrue(resolved.contains("Health Coach should use direct accountability when the user is behind on steps.")) + XCTAssertFalse(resolved.contains("Travel Planner should keep itineraries compact and transit-aware.")) } func testRuntimeMemorySelectionCanReplaceOrDisableThreadDefaults() async throws { @@ -788,16 +787,16 @@ final class AgentRuntimeTests: XCTestCase { ) let store = InMemoryMemoryStore(initialRecords: [ MemoryRecord( - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", - summary: "Eleanor grievance." + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + summary: "Health Coach preference." ), MemoryRecord( - namespace: "oval-office", - scope: "actor:sophia_ramirez", - kind: "grievance", - summary: "Sophia grievance." + namespace: "demo-assistant", + scope: "feature:travel-planner", + kind: "preference", + summary: "Travel Planner preference." ), ]) let runtime = try AgentRuntime(configuration: .init( @@ -818,39 +817,174 @@ final class AgentRuntimeTests: XCTestCase { let thread = try await runtime.createThread( title: "Scoped Memory", memoryContext: AgentMemoryContext( - namespace: "oval-office", - scopes: ["actor:eleanor_price"] + namespace: "demo-assistant", + scopes: ["feature:health-coach"] ) ) - let replaceStream = try await runtime.sendMessage( + _ = try await runtime.sendMessage( UserMessageRequest( - text: "Use Sophia memory instead.", + text: "Use travel planner memory instead.", memorySelection: MemorySelection( mode: .replace, - scopes: ["actor:sophia_ramirez"] + scopes: ["feature:travel-planner"] ) ), in: thread.id ) - for try await _ in replaceStream {} - let disableStream = try await runtime.sendMessage( + _ = try await runtime.sendMessage( UserMessageRequest( text: "Now disable memory.", memorySelection: MemorySelection(mode: .disable) ), in: thread.id ) - for try await _ in disableStream {} let instructions = await backend.receivedInstructions() XCTAssertEqual(instructions.count, 2) - XCTAssertTrue(instructions[0].contains("Sophia grievance.")) - XCTAssertFalse(instructions[0].contains("Eleanor grievance.")) + XCTAssertTrue(instructions[0].contains("Travel Planner preference.")) + XCTAssertFalse(instructions[0].contains("Health Coach preference.")) XCTAssertFalse(instructions[1].contains("Relevant Memory:")) } + func testRuntimeCanAutomaticallyCaptureMemoriesFromTranscript() async throws { + let backend = InMemoryAgentBackend( + structuredResponseText: """ + {"memories":[{"summary":"Health Coach should use direct accountability when step pace is low.","scope":"feature:health-coach","kind":"preference","evidence":["The user asked for blunt reminders when behind on steps."],"importance":0.92,"tags":["steps","tone"],"relatedIDs":["goal-10000"],"dedupeKey":"health-coach-direct-accountability"},{"summary":"Travel Planner should keep itineraries compact and transit-aware.","scope":"feature:travel-planner","kind":"preference","evidence":["The user dislikes sprawling travel plans."],"importance":0.81,"tags":["travel"],"relatedIDs":["travel-style-compact"],"dedupeKey":"travel-planner-compact-itinerary"}]} + """ + ) + let store = InMemoryMemoryStore() + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: backend, + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init(store: store) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Auto Memory", + memoryContext: AgentMemoryContext( + namespace: "demo-assistant", + scopes: ["feature:health-coach", "feature:travel-planner"] + ) + ) + + let capture = try await runtime.captureMemories( + from: .text(""" + User: Be direct when I am behind on steps. + User: Keep travel itineraries compact and transit-aware. + """), + for: thread.id, + options: .init( + defaults: .init(namespace: "demo-assistant"), + maxMemories: 3 + ) + ) + + XCTAssertEqual(capture.records.count, 2) + XCTAssertEqual(capture.records.map(\.scope.rawValue).sorted(), ["feature:health-coach", "feature:travel-planner"]) + + let stored = try await store.query( + MemoryQuery( + namespace: "demo-assistant", + scopes: ["feature:health-coach", "feature:travel-planner"], + text: "direct steps transit itinerary", + limit: 10, + maxCharacters: 1000 + ) + ) + XCTAssertEqual(stored.matches.count, 2) + + let formats = await backend.receivedResponseFormats() + XCTAssertEqual(formats.last??.name, "memory_extraction_batch") + } + + func testRuntimeCanAutomaticallyCaptureMemoryAfterSuccessfulTurn() async throws { + let backend = InMemoryAgentBackend( + structuredResponseText: """ + {"memories":[{"summary":"Health Coach should use direct accountability when the user falls behind on steps.","scope":"feature:health-coach","kind":"preference","evidence":["The user said blunt reminders work better than soft encouragement."],"importance":0.94,"tags":["steps","tone"],"relatedIDs":["goal-10000"],"dedupeKey":"health-coach-auto-capture"}]} + """ + ) + let store = InMemoryMemoryStore() + let observer = RecordingMemoryObserver() + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: backend, + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init( + store: store, + observer: observer, + automaticCapturePolicy: .init( + source: .lastTurn, + options: .init( + defaults: .init(namespace: "demo-assistant"), + maxMemories: 2 + ) + ) + ) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Auto Policy", + memoryContext: AgentMemoryContext( + namespace: "demo-assistant", + scopes: ["feature:health-coach"] + ) + ) + + _ = try await runtime.sendMessage( + UserMessageRequest(text: "If I am behind on steps, be direct and blunt with me."), + in: thread.id + ) + + let stored = try await store.query( + MemoryQuery( + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + text: "direct blunt steps", + limit: 10, + maxCharacters: 1000 + ) + ) + XCTAssertEqual(stored.matches.count, 1) + XCTAssertEqual(stored.matches[0].record.dedupeKey, "health-coach-auto-capture") + + let formats = await backend.receivedResponseFormats() + XCTAssertEqual(formats.count, 2) + XCTAssertNil(formats.first!) + XCTAssertEqual(formats.last??.name, "memory_extraction_batch") + + let events = await observer.events() + XCTAssertEqual(events.count, 2) + guard case let .captureStarted(captureThreadID, sourceDescription) = events[0] else { + return XCTFail("Expected captureStarted event.") + } + XCTAssertEqual(captureThreadID, thread.id) + XCTAssertEqual(sourceDescription, "last_turn") + guard case let .captureSucceeded(succeededThreadID, result) = events[1] else { + return XCTFail("Expected captureSucceeded event.") + } + XCTAssertEqual(succeededThreadID, thread.id) + XCTAssertEqual(result.records.count, 1) + } + func testRuntimeGracefullyDegradesWhenMemoryStoreFails() async throws { let backend = InMemoryAgentBackend( baseInstructions: "Base host instructions." @@ -873,16 +1007,15 @@ final class AgentRuntimeTests: XCTestCase { let thread = try await runtime.createThread( title: "Graceful", memoryContext: AgentMemoryContext( - namespace: "oval-office", - scopes: ["actor:eleanor_price"] + namespace: "demo-assistant", + scopes: ["feature:health-coach"] ) ) - let stream = try await runtime.sendMessage( + _ = try await runtime.sendMessage( UserMessageRequest(text: "This should still work."), in: thread.id ) - for try await _ in stream {} let instructions = await backend.receivedInstructions() let resolved = try XCTUnwrap(instructions.last) @@ -895,9 +1028,9 @@ final class AgentRuntimeTests: XCTestCase { ) let store = InMemoryMemoryStore(initialRecords: [ MemoryRecord( - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", summary: "Observed memory." ), ]) @@ -923,29 +1056,89 @@ final class AgentRuntimeTests: XCTestCase { let thread = try await runtime.createThread( title: "Observed", memoryContext: AgentMemoryContext( - namespace: "oval-office", - scopes: ["actor:eleanor_price"] + namespace: "demo-assistant", + scopes: ["feature:health-coach"] ) ) - let stream = try await runtime.sendMessage( + _ = try await runtime.sendMessage( UserMessageRequest(text: "Use memory."), in: thread.id ) - for try await _ in stream {} let events = await observer.events() XCTAssertEqual(events.count, 2) guard case let .queryStarted(startedQuery) = events[0] else { return XCTFail("Expected queryStarted event.") } - XCTAssertEqual(startedQuery.namespace, "oval-office") + XCTAssertEqual(startedQuery.namespace, "demo-assistant") guard case let .querySucceeded(_, result) = events[1] else { return XCTFail("Expected querySucceeded event.") } XCTAssertEqual(result.matches.count, 1) } + func testRuntimeProvidesThreadAwareMemoryWriterDefaults() async throws { + let store = InMemoryMemoryStore() + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: InMemoryAgentBackend(), + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore(), + memory: .init(store: store) + )) + + _ = try await runtime.restore() + _ = try await runtime.signIn() + + let thread = try await runtime.createThread( + title: "Thread Memory Writer", + memoryContext: AgentMemoryContext( + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + kinds: ["preference"], + tags: ["steps"], + relatedIDs: ["goal-10000"] + ) + ) + + let writer = try await runtime.memoryWriter(for: thread.id) + let record = try await writer.put( + MemoryDraft( + summary: "The demo user responds better to direct step reminders." + ) + ) + + XCTAssertEqual(record.namespace, "demo-assistant") + XCTAssertEqual(record.scope, "feature:health-coach") + XCTAssertEqual(record.kind, "preference") + XCTAssertEqual(record.tags, ["steps"]) + XCTAssertEqual(record.relatedIDs, ["goal-10000"]) + } + + func testRuntimeMemoryWriterThrowsWhenMemoryIsNotConfigured() async throws { + let runtime = try AgentRuntime(configuration: .init( + authProvider: DemoChatGPTAuthProvider(), + secureStore: KeychainSessionSecureStore( + service: "CodexKitTests.ChatGPTSession", + account: UUID().uuidString + ), + backend: InMemoryAgentBackend(), + approvalPresenter: AutoApprovalPresenter(), + stateStore: InMemoryRuntimeStateStore() + )) + + _ = try await runtime.restore() + + await XCTAssertThrowsErrorAsync(try await runtime.memoryWriter()) { error in + XCTAssertEqual(error as? AgentRuntimeError, .memoryNotConfigured()) + } + } + func testResolvedInstructionsPreviewThrowsForMissingThread() async throws { let runtime = try AgentRuntime(configuration: .init( authProvider: DemoChatGPTAuthProvider(), @@ -1176,8 +1369,8 @@ final class AgentRuntimeTests: XCTestCase { account: UUID().uuidString ) let memoryContext = AgentMemoryContext( - namespace: "oval-office", - scopes: ["actor:eleanor_price", "world:public"], + namespace: "demo-assistant", + scopes: ["feature:health-coach", "feature:travel-planner"], readBudget: .init(maxItems: 4, maxCharacters: 800) ) diff --git a/Tests/CodexKitTests/MemoryStoreTests.swift b/Tests/CodexKitTests/MemoryStoreTests.swift index b66655a..26157ef 100644 --- a/Tests/CodexKitTests/MemoryStoreTests.swift +++ b/Tests/CodexKitTests/MemoryStoreTests.swift @@ -10,15 +10,15 @@ final class MemoryStoreTests: XCTestCase { let store = try SQLiteMemoryStore(url: url) let record = MemoryRecord( - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", - summary: "Eleanor remembers being overruled on the trade bill.", - evidence: ["The player dismissed her warning on day 12."], + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + summary: "Health Coach should use direct accountability when the user is behind on steps.", + evidence: ["The user ignores soft reminders late in the day."], importance: 0.9, - tags: ["trade", "advisors"], - relatedIDs: ["bill-12"], - dedupeKey: "eleanor-trade-day-12" + tags: ["steps", "tone"], + relatedIDs: ["goal-10000"], + dedupeKey: "health-coach-direct-accountability" ) try await store.put(record) @@ -26,9 +26,9 @@ final class MemoryStoreTests: XCTestCase { let reloaded = try SQLiteMemoryStore(url: url) let result = try await reloaded.query( MemoryQuery( - namespace: "oval-office", - scopes: ["actor:eleanor_price"], - text: "overruled trade warning", + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + text: "direct steps reminder", limit: 5, maxCharacters: 600 ) @@ -37,7 +37,7 @@ final class MemoryStoreTests: XCTestCase { XCTAssertEqual(result.matches.map(\.record.id), [record.id]) XCTAssertGreaterThan(result.matches[0].explanation.textScore, 0) - let diagnostics = try await reloaded.diagnostics(namespace: "oval-office") + let diagnostics = try await reloaded.diagnostics(namespace: "demo-assistant") XCTAssertEqual(diagnostics.implementation, "sqlite") XCTAssertEqual(diagnostics.schemaVersion, 1) } @@ -74,8 +74,8 @@ final class MemoryStoreTests: XCTestCase { let store = try SQLiteMemoryStore(url: url) let existing = MemoryRecord( id: "memory-1", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:health-coach", kind: "fact", summary: "Existing memory." ) @@ -85,15 +85,15 @@ final class MemoryStoreTests: XCTestCase { try await store.putMany([ MemoryRecord( id: "memory-2", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:health-coach", kind: "fact", summary: "Should roll back." ), MemoryRecord( id: "memory-1", - namespace: "oval-office", - scope: "actor:sophia_ramirez", + namespace: "demo-assistant", + scope: "feature:travel-planner", kind: "fact", summary: "Duplicate id." ), @@ -107,7 +107,7 @@ final class MemoryStoreTests: XCTestCase { let result = try await store.query( MemoryQuery( - namespace: "oval-office", + namespace: "demo-assistant", scopes: [], limit: 10, maxCharacters: 1000 @@ -122,29 +122,29 @@ final class MemoryStoreTests: XCTestCase { try await store.upsert( MemoryRecord( id: "memory-1", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:travel-planner", kind: "quote_record", - summary: "Initial quote." + summary: "Initial itinerary note." ), - dedupeKey: "press-quote-17" + dedupeKey: "travel-plan-17" ) try await store.upsert( MemoryRecord( id: "memory-2", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:travel-planner", kind: "quote_record", - summary: "Updated quote after retry." + summary: "Updated itinerary note after retry." ), - dedupeKey: "press-quote-17" + dedupeKey: "travel-plan-17" ) let result = try await store.query( MemoryQuery( - namespace: "oval-office", - scopes: ["actor:eleanor_price"], + namespace: "demo-assistant", + scopes: ["feature:travel-planner"], limit: 10, maxCharacters: 1000 ) @@ -152,7 +152,75 @@ final class MemoryStoreTests: XCTestCase { XCTAssertEqual(result.matches.count, 1) XCTAssertEqual(result.matches[0].record.id, "memory-2") - XCTAssertEqual(result.matches[0].record.dedupeKey, "press-quote-17") + XCTAssertEqual(result.matches[0].record.dedupeKey, "travel-plan-17") + } + + func testMemoryWriterAppliesDefaultsAndUpsertsDraft() async throws { + let store = InMemoryMemoryStore() + let writer = MemoryWriter( + store: store, + defaults: MemoryWriterDefaults( + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + importance: 0.6, + tags: ["demo", "health"], + relatedIDs: ["goal-10000"] + ) + ) + + let record = try await writer.upsert( + MemoryDraft( + summary: "Use direct accountability when the user is behind on steps.", + evidence: ["The user follows through more often with blunt reminders."], + importance: 0.95, + tags: ["tone"], + dedupeKey: "health-coach-direct-tone" + ) + ) + + XCTAssertEqual(record.namespace, "demo-assistant") + XCTAssertEqual(record.scope, "feature:health-coach") + XCTAssertEqual(record.kind, "preference") + XCTAssertEqual(record.tags, ["demo", "health", "tone"]) + XCTAssertEqual(record.relatedIDs, ["goal-10000"]) + XCTAssertEqual(record.dedupeKey, "health-coach-direct-tone") + + let result = try await store.query( + MemoryQuery( + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + text: "direct steps accountability", + limit: 10, + maxCharacters: 1000 + ) + ) + XCTAssertEqual(result.matches.map(\.record.id), [record.id]) + } + + func testMemoryWriterThrowsWhenRequiredDefaultsAreMissing() async throws { + let writer = MemoryWriter(store: InMemoryMemoryStore()) + + XCTAssertThrowsError( + try writer.resolve( + MemoryDraft(summary: "Missing namespace and scope.") + ) + ) { error in + XCTAssertEqual(error as? MemoryAuthoringError, .missingNamespace) + } + + let namespaceOnlyWriter = MemoryWriter( + store: InMemoryMemoryStore(), + defaults: MemoryWriterDefaults(namespace: "demo-assistant") + ) + + XCTAssertThrowsError( + try namespaceOnlyWriter.resolve( + MemoryDraft(summary: "Still missing scope.") + ) + ) { error in + XCTAssertEqual(error as? MemoryAuthoringError, .missingScope) + } } func testStoreInspectionAPIsReturnRecordsAndDiagnostics() async throws { @@ -160,15 +228,15 @@ final class MemoryStoreTests: XCTestCase { try await store.putMany([ MemoryRecord( id: "active-memory", - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", summary: "Active memory." ), MemoryRecord( id: "archived-memory", - namespace: "oval-office", - scope: "world:press", + namespace: "demo-assistant", + scope: "feature:travel-planner", kind: "summary", summary: "Archived memory.", status: .archived @@ -177,24 +245,24 @@ final class MemoryStoreTests: XCTestCase { let fetched = try await store.record( id: "active-memory", - namespace: "oval-office" + namespace: "demo-assistant" ) - XCTAssertEqual(fetched?.kind, "grievance") + XCTAssertEqual(fetched?.kind, "preference") let listed = try await store.list( - namespace: "oval-office", + namespace: "demo-assistant", includeArchived: true, limit: 10 ) XCTAssertEqual(listed.map(\.id).sorted(), ["active-memory", "archived-memory"]) - let diagnostics = try await store.diagnostics(namespace: "oval-office") + let diagnostics = try await store.diagnostics(namespace: "demo-assistant") XCTAssertEqual(diagnostics.implementation, "in_memory") XCTAssertNil(diagnostics.schemaVersion) XCTAssertEqual(diagnostics.totalRecords, 2) XCTAssertEqual(diagnostics.activeRecords, 1) XCTAssertEqual(diagnostics.archivedRecords, 1) - XCTAssertEqual(diagnostics.countsByScope["actor:eleanor_price"], 1) + XCTAssertEqual(diagnostics.countsByScope["feature:health-coach"], 1) XCTAssertEqual(diagnostics.countsByKind["summary"], 1) } @@ -203,46 +271,46 @@ final class MemoryStoreTests: XCTestCase { try await store.putMany([ MemoryRecord( id: "match-1", - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", - summary: "Eleanor is still angry about the farm subsidy reversal.", - evidence: ["She warned the player twice before being ignored."], + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + summary: "Health Coach should push for a 15 minute walk when step pace is low.", + evidence: ["The user follows through when told to walk before dinner."], importance: 0.95, - tags: ["farm", "economy"], - relatedIDs: ["policy-farm"] + tags: ["steps", "walk"], + relatedIDs: ["goal-10000"] ), MemoryRecord( id: "match-2", - namespace: "oval-office", - scope: "actor:eleanor_price", - kind: "grievance", - summary: "A second long grievance that should be trimmed by the budget.", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "preference", + summary: "A second long coaching note that should be trimmed by the budget.", evidence: ["This extra line makes the rendered memory longer than the cap."], importance: 0.80, - tags: ["farm"], - relatedIDs: ["policy-farm"] + tags: ["steps"], + relatedIDs: ["goal-10000"] ), MemoryRecord( id: "other-scope", - namespace: "oval-office", - scope: "actor:sophia_ramirez", - kind: "grievance", - summary: "Sophia has a separate grievance.", + namespace: "demo-assistant", + scope: "feature:travel-planner", + kind: "preference", + summary: "Travel Planner prefers early museum starts.", importance: 0.99, - tags: ["farm"], - relatedIDs: ["policy-farm"] + tags: ["steps"], + relatedIDs: ["goal-10000"] ), ]) let result = try await store.query( MemoryQuery( - namespace: "oval-office", - scopes: ["actor:eleanor_price"], - text: "farm grievance warning", - kinds: ["grievance"], - tags: ["farm"], - relatedIDs: ["policy-farm"], + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + text: "steps walk dinner", + kinds: ["preference"], + tags: ["steps"], + relatedIDs: ["goal-10000"], minImportance: 0.5, limit: 10, maxCharacters: 220 @@ -257,6 +325,79 @@ final class MemoryStoreTests: XCTestCase { XCTAssertGreaterThan(result.matches[0].explanation.relatedIDBoost, 0) } + func testInMemoryQueryPrefersHigherTextOverlap() async throws { + let store = InMemoryMemoryStore() + try await store.putMany([ + MemoryRecord( + id: "strong-match", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "fact", + summary: "Health Coach should ask for a brisk evening walk when steps are low." + ), + MemoryRecord( + id: "weak-match", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "fact", + summary: "Health Coach checked in with the user today." + ), + ]) + + let result = try await store.query( + MemoryQuery( + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + text: "brisk walk steps", + limit: 10, + maxCharacters: 1000 + ) + ) + + XCTAssertEqual(result.matches.first?.record.id, "strong-match") + XCTAssertGreaterThan( + result.matches.first?.explanation.textScore ?? 0, + result.matches.dropFirst().first?.explanation.textScore ?? 0 + ) + } + + func testQuerySkipsOversizedTopCandidateAndKeepsSmallerMatch() async throws { + let store = InMemoryMemoryStore() + let oversizedSummary = String(repeating: "trade ", count: 80) + + try await store.putMany([ + MemoryRecord( + id: "oversized", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "fact", + summary: oversizedSummary, + importance: 1.0 + ), + MemoryRecord( + id: "fits-budget", + namespace: "demo-assistant", + scope: "feature:health-coach", + kind: "fact", + summary: "Short step reminder for Health Coach.", + importance: 0.2 + ), + ]) + + let result = try await store.query( + MemoryQuery( + namespace: "demo-assistant", + scopes: ["feature:health-coach"], + text: "steps", + limit: 10, + maxCharacters: 120 + ) + ) + + XCTAssertEqual(result.matches.map(\.record.id), ["fits-budget"]) + XCTAssertTrue(result.truncated) + } + func testCompactArchivesSourcesAndPruneExpiredSkipsPinned() async throws { let store = InMemoryMemoryStore() let now = Date() @@ -264,22 +405,22 @@ final class MemoryStoreTests: XCTestCase { try await store.putMany([ MemoryRecord( id: "source-1", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:health-coach", kind: "fact", summary: "Source memory one." ), MemoryRecord( id: "source-2", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:health-coach", kind: "fact", summary: "Source memory two." ), MemoryRecord( id: "expired-pinned", - namespace: "oval-office", - scope: "world:press", + namespace: "demo-assistant", + scope: "feature:travel-planner", kind: "summary", summary: "Pinned memory should survive pruning.", expiresAt: now.addingTimeInterval(-60), @@ -287,8 +428,8 @@ final class MemoryStoreTests: XCTestCase { ), MemoryRecord( id: "expired-unpinned", - namespace: "oval-office", - scope: "world:press", + namespace: "demo-assistant", + scope: "feature:travel-planner", kind: "summary", summary: "Unpinned memory should be removed.", expiresAt: now.addingTimeInterval(-60) @@ -299,10 +440,10 @@ final class MemoryStoreTests: XCTestCase { MemoryCompactionRequest( replacement: MemoryRecord( id: "replacement", - namespace: "oval-office", - scope: "actor:eleanor_price", + namespace: "demo-assistant", + scope: "feature:health-coach", kind: "summary", - summary: "Compacted grievance summary." + summary: "Compacted health coach summary." ), sourceIDs: ["source-1", "source-2"] ) @@ -310,8 +451,8 @@ final class MemoryStoreTests: XCTestCase { let active = try await store.query( MemoryQuery( - namespace: "oval-office", - scopes: ["actor:eleanor_price"], + namespace: "demo-assistant", + scopes: ["feature:health-coach"], limit: 10, maxCharacters: 1000 ) @@ -320,14 +461,14 @@ final class MemoryStoreTests: XCTestCase { let prunedCount = try await store.pruneExpired( now: now, - namespace: "oval-office" + namespace: "demo-assistant" ) XCTAssertEqual(prunedCount, 1) let remaining = try await store.query( MemoryQuery( - namespace: "oval-office", - scopes: ["world:press"], + namespace: "demo-assistant", + scopes: ["feature:travel-planner"], limit: 10, maxCharacters: 1000 ) From e81a42a64fbcc2400566054ef1a658450a477026 Mon Sep 17 00:00:00 2001 From: Timothy Zelinsky Date: Sun, 22 Mar 2026 18:11:19 +1100 Subject: [PATCH 4/4] Fix automatic memory capture CI test --- Tests/CodexKitTests/AgentRuntimeTests.swift | 25 ++++++++++++--------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/Tests/CodexKitTests/AgentRuntimeTests.swift b/Tests/CodexKitTests/AgentRuntimeTests.swift index 060e7fc..2e81706 100644 --- a/Tests/CodexKitTests/AgentRuntimeTests.swift +++ b/Tests/CodexKitTests/AgentRuntimeTests.swift @@ -972,17 +972,22 @@ final class AgentRuntimeTests: XCTestCase { XCTAssertEqual(formats.last??.name, "memory_extraction_batch") let events = await observer.events() - XCTAssertEqual(events.count, 2) - guard case let .captureStarted(captureThreadID, sourceDescription) = events[0] else { - return XCTFail("Expected captureStarted event.") - } - XCTAssertEqual(captureThreadID, thread.id) - XCTAssertEqual(sourceDescription, "last_turn") - guard case let .captureSucceeded(succeededThreadID, result) = events[1] else { - return XCTFail("Expected captureSucceeded event.") + let captureEvents = events.compactMap { event -> (String, String?, Int?)? in + switch event { + case let .captureStarted(threadID, sourceDescription): + return (threadID, sourceDescription, nil) + case let .captureSucceeded(threadID, result): + return (threadID, nil, result.records.count) + default: + return nil + } } - XCTAssertEqual(succeededThreadID, thread.id) - XCTAssertEqual(result.records.count, 1) + + XCTAssertEqual(captureEvents.count, 2) + XCTAssertEqual(captureEvents[0].0, thread.id) + XCTAssertEqual(captureEvents[0].1, "last_turn") + XCTAssertEqual(captureEvents[1].0, thread.id) + XCTAssertEqual(captureEvents[1].2, 1) } func testRuntimeGracefullyDegradesWhenMemoryStoreFails() async throws {