diff --git a/brain-bar/Sources/BrainBar/BrainDatabase.swift b/brain-bar/Sources/BrainBar/BrainDatabase.swift index e560fd98d..750f843d5 100644 --- a/brain-bar/Sources/BrainBar/BrainDatabase.swift +++ b/brain-bar/Sources/BrainBar/BrainDatabase.swift @@ -780,6 +780,9 @@ final class BrainDatabase: @unchecked Sendable { // already scrubbed, so replay would otherwise find nothing to record. let secretScrubRedactions: [String]? let secretScrubQuarantineCount: Int? + // #655 Part B: the caller's self-reported model. Absent on records queued + // before this field existed, which replay exactly as they did before. + let callerModel: String? enum CodingKeys: String, CodingKey { case content @@ -795,6 +798,7 @@ final class BrainDatabase: @unchecked Sendable { case conversationID = "conversation_id" case secretScrubRedactions = "secret_scrub_redactions" case secretScrubQuarantineCount = "secret_scrub_quarantine_count" + case callerModel = "caller_model" } init( @@ -810,7 +814,8 @@ final class BrainDatabase: @unchecked Sendable { chunkID: String? = nil, conversationID: String? = nil, secretScrubRedactions: [String]? = nil, - secretScrubQuarantineCount: Int? = nil + secretScrubQuarantineCount: Int? = nil, + callerModel: String? = nil ) { self.content = content self.tags = tags @@ -825,6 +830,7 @@ final class BrainDatabase: @unchecked Sendable { self.conversationID = conversationID self.secretScrubRedactions = secretScrubRedactions self.secretScrubQuarantineCount = secretScrubQuarantineCount + self.callerModel = BrainDatabase.normalizedCallerModel(callerModel) } init(from decoder: Decoder) throws { @@ -842,6 +848,11 @@ final class BrainDatabase: @unchecked Sendable { conversationID = try container.decodeIfPresent(String.self, forKey: .conversationID) secretScrubRedactions = try container.decodeIfPresent([String].self, forKey: .secretScrubRedactions) secretScrubQuarantineCount = try container.decodeIfPresent(Int.self, forKey: .secretScrubQuarantineCount) + // Lenient on purpose: a malformed model on a queued line drops the model, + // never the memory. A strict decode would strand the whole line. + callerModel = BrainDatabase.normalizedCallerModel( + try? container.decodeIfPresent(String.self, forKey: .callerModel) + ) } /// This item with content and tags scrubbed and the findings folded in. @@ -864,7 +875,8 @@ final class BrainDatabase: @unchecked Sendable { chunkID: chunkID, conversationID: conversationID, secretScrubRedactions: result.metadata["secret_scrub_redactions"] as? [String], - secretScrubQuarantineCount: result.metadata["secret_scrub_quarantine_count"] as? Int + secretScrubQuarantineCount: result.metadata["secret_scrub_quarantine_count"] as? Int, + callerModel: callerModel ) } @@ -875,6 +887,48 @@ final class BrainDatabase: @unchecked Sendable { if let secretScrubQuarantineCount { metadata["secret_scrub_quarantine_count"] = secretScrubQuarantineCount } return metadata } + + /// Everything a replayed row records: scrub findings plus caller model provenance, + /// so a replayed row carries the same metadata a direct store would have. + var replayMetadata: [String: Any] { + scrubMetadata.merging(BrainDatabase.callerModelMetadata(callerModel)) { current, _ in current } + } + } + + /// Longest caller-reported model accepted, in Unicode scalars. Mirrors the + /// `caller_model` maxLength advertised on the brain_store inputSchema. + static let callerModelMaxLength = 128 + + /// The caller's self-reported model, trimmed, or nil when it cannot be trusted as one. + /// Nothing is inferred: absent, blank, over-long, control/format characters, the + /// `` placeholder, or anything the secret scrubber would touch → unknown. + static func normalizedCallerModel(_ raw: String?) -> String? { + guard let model = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !model.isEmpty else { + return nil + } + guard model.unicodeScalars.count <= callerModelMaxLength, model != "" else { + return nil + } + for scalar in model.unicodeScalars { + switch scalar.properties.generalCategory { + case .control, .format, .lineSeparator, .paragraphSeparator, .surrogate, .privateUse, .unassigned: + return nil + default: + continue + } + } + let scrubbed = SecretScrubber.scrub(model) + guard scrubbed.text == model, scrubbed.quarantineCount == 0 else { + return nil + } + return model + } + + /// Authoring-model provenance for a brain_store row. Written only when the caller + /// supplied a valid model; never the enrichment_model / enrichment_backend columns. + static func callerModelMetadata(_ raw: String?) -> [String: Any] { + guard let model = normalizedCallerModel(raw) else { return [:] } + return ["model": model, "model_source": "caller_self_reported"] } private var db: OpaquePointer? @@ -1732,6 +1786,7 @@ final class BrainDatabase: @unchecked Sendable { source: String, project: String? = nil, conversationID: String? = nil, + callerModel: String? = nil, busyTimeoutMillis: Int32 = 30_000, retries: Int = 3 ) throws -> StoreWriteOutcome { @@ -1756,7 +1811,8 @@ final class BrainDatabase: @unchecked Sendable { conversationID: conversationID, refreshStatistics: true, retries: retries, - busyTimeoutMillis: busyTimeoutMillis + busyTimeoutMillis: busyTimeoutMillis, + metadata: Self.callerModelMetadata(callerModel) ) return .stored(stored) } catch { @@ -1771,7 +1827,8 @@ final class BrainDatabase: @unchecked Sendable { project: project, createdAt: createdAt, chunkID: chunkID, - conversationID: conversationID + conversationID: conversationID, + callerModel: callerModel ) return .queued(queueID: queued.queueID, queuedAt: queued.queuedAt, chunkID: queued.chunkID) } @@ -1828,6 +1885,7 @@ final class BrainDatabase: @unchecked Sendable { createdAt: String? = nil, chunkID: String? = nil, conversationID: String? = nil, + callerModel: String? = nil, reason: String = "DB_BUSY" ) throws -> (queueID: String, queuedAt: String, chunkID: String) { try Self.queuePendingStore( @@ -1840,6 +1898,7 @@ final class BrainDatabase: @unchecked Sendable { createdAt: createdAt, chunkID: chunkID, conversationID: conversationID, + callerModel: callerModel, reason: reason, sync: pendingStoreSyncForTesting, directorySync: pendingStoreDirectorySyncForTesting @@ -1865,6 +1924,7 @@ final class BrainDatabase: @unchecked Sendable { createdAt: String? = nil, chunkID: String? = nil, conversationID: String? = nil, + callerModel: String? = nil, reason: String = "DB_BUSY", sync: ((Int32) throws -> Void)? = nil, directorySync: ((URL) throws -> Void)? = nil @@ -1890,7 +1950,8 @@ final class BrainDatabase: @unchecked Sendable { telemetry: DeferredStoreTelemetry.Origin(epoch: Date().timeIntervalSince1970, reason: reason), createdAt: createdAt, chunkID: chunkID, - conversationID: conversationID + conversationID: conversationID, + callerModel: callerModel ).scrubbed() var line = try JSONEncoder().encode(item) line.append(0x0A) @@ -1976,7 +2037,7 @@ final class BrainDatabase: @unchecked Sendable { refreshStatistics: false, retries: retries, busyTimeoutMillis: busyTimeoutMillis, - metadata: item.scrubMetadata + metadata: item.replayMetadata ) if stored.outcome.storedNew { DeferredStoreTelemetry.record("committed", dbPath: self.path, queueID: queueID, @@ -4968,7 +5029,8 @@ final class BrainDatabase: @unchecked Sendable { chunkID: chunkID, conversationID: item.conversationID, secretScrubRedactions: item.secretScrubRedactions, - secretScrubQuarantineCount: item.secretScrubQuarantineCount + secretScrubQuarantineCount: item.secretScrubQuarantineCount, + callerModel: item.callerModel ) return try? JSONEncoder().encode(replayItem) } diff --git a/brain-bar/Sources/BrainBar/MCPRouter.swift b/brain-bar/Sources/BrainBar/MCPRouter.swift index 49623d3f2..092b8d1a1 100644 --- a/brain-bar/Sources/BrainBar/MCPRouter.swift +++ b/brain-bar/Sources/BrainBar/MCPRouter.swift @@ -239,6 +239,7 @@ final class MCPRouter: @unchecked Sendable { private static let stringMaxLengths: [String: Int] = [ "action": 64, "agent_id": 128, + "caller_model": BrainDatabase.callerModelMaxLength, "chunk_id": chunkIdentifierMaxLength, "content": 200_000, "detail": 16, @@ -355,12 +356,25 @@ final class MCPRouter: @unchecked Sendable { "name": name, "description": description, ] - if let inputSchema = definition["inputSchema"] as? [String: Any] { + if var inputSchema = definition["inputSchema"] as? [String: Any] { + if let hidden = fullPaletteOnlyProperties[name], + var properties = inputSchema["properties"] as? [String: Any] { + for property in hidden { properties.removeValue(forKey: property) } + inputSchema["properties"] = properties + } compact["inputSchema"] = removingDescriptions(from: inputSchema) } return compact } + /// Optional properties advertised only once the palette is expanded. They are still + /// accepted and validated against the full schema from a core session. The core boot + /// payload has a ratified 1600-byte budget that main already fills (#655 Part B: + /// advertising caller_model in core costs 49 bytes and is an owner decision). + private static let fullPaletteOnlyProperties: [String: Set] = [ + "brain_store": ["caller_model"], + ] + private static func removingDescriptions(from value: Any) -> Any { if let dictionary = value as? [String: Any] { return dictionary.reduce(into: [String: Any]()) { result, entry in @@ -1047,6 +1061,9 @@ final class MCPRouter: @unchecked Sendable { let tags = args["tags"] as? [String] ?? [] let importance = args["importance"] as? Int ?? 5 let project = args["project"] as? String + // Self-reported by the caller (MCP cannot see who is calling). Never inferred + // from agent_id or the session; an invalid value is dropped, not stored. + let callerModel = args["caller_model"] as? String if deferForBackup { return try queueBrainStore( content: content, @@ -1055,6 +1072,7 @@ final class MCPRouter: @unchecked Sendable { source: "mcp", project: project, conversationID: session.conversationID, + callerModel: callerModel, reason: "BACKUP_SNAPSHOT" ) } @@ -1066,6 +1084,7 @@ final class MCPRouter: @unchecked Sendable { source: "mcp", project: project, conversationID: session.conversationID, + callerModel: callerModel, reason: "DB_NOT_OPEN" ) } @@ -1078,6 +1097,7 @@ final class MCPRouter: @unchecked Sendable { source: "mcp", project: project, conversationID: session.conversationID, + callerModel: callerModel, busyTimeoutMillis: Self.mcpStoreBusyTimeoutMillis, retries: Self.mcpStoreRetries ) { @@ -1137,6 +1157,7 @@ final class MCPRouter: @unchecked Sendable { source: "mcp", project: project, conversationID: session.conversationID, + callerModel: callerModel, reason: "DB_NOT_OPEN" ) } @@ -1284,6 +1305,7 @@ final class MCPRouter: @unchecked Sendable { source: String, project: String?, conversationID: String, + callerModel: String? = nil, reason: String ) throws -> ToolOutput { guard let dbPath else { @@ -1297,6 +1319,7 @@ final class MCPRouter: @unchecked Sendable { source: source, project: project, conversationID: conversationID, + callerModel: callerModel, reason: reason ) return queuedBrainStoreOutput( @@ -1929,6 +1952,7 @@ final class MCPRouter: @unchecked Sendable { "tags": ["type": "array", "items": ["type": "string"], "description": "Tags for categorization"], "importance": ["type": "integer", "description": "Importance score (1-10)"], "project": ["type": "string", "description": "Project context for the stored memory"], + "caller_model": ["type": "string", "description": "Optional: your own model ID, self-reported. Stored as metadata.model; never inferred."], ] as [String: Any], "required": ["content"] ] as [String: Any]) diff --git a/brain-bar/Tests/BrainBarTests/StoreCallerModelTests.swift b/brain-bar/Tests/BrainBarTests/StoreCallerModelTests.swift new file mode 100644 index 000000000..e6f008bf3 --- /dev/null +++ b/brain-bar/Tests/BrainBarTests/StoreCallerModelTests.swift @@ -0,0 +1,397 @@ +// StoreCallerModelTests.swift — #655 Part B: brain_store records the caller's +// self-reported model, on every write path, and never invents one. +// +// Everything here drives the served surface: the brain_store wire arguments, the +// durable pending-stores.jsonl file and the chunks row. Model strings are +// synthetic fixtures, not real model IDs. + +import Foundation +import SQLite3 +import XCTest +@testable import BrainBar + +final class StoreCallerModelTests: XCTestCase { + private let modelA = "fixture-model-a" + private let modelB = "fixture-model-b" + + private var tempDir: URL! + private var dbPath: String! + + override func setUpWithError() throws { + tempDir = brainBarTestTemporaryDirectory() + .appendingPathComponent("brainbar-caller-model-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) + dbPath = tempDir.appendingPathComponent("brainbar.db").path + } + + override func tearDownWithError() throws { + try? FileManager.default.removeItem(at: tempDir) + } + + private var queuePath: URL { tempDir.appendingPathComponent("pending-stores.jsonl") } + + // MARK: helpers + + /// A router whose deferred drains never run on their own, so each test decides + /// when replay happens. + private func makeRouter(_ db: BrainDatabase, drainQueue: DispatchQueue) -> MCPRouter { + let router = MCPRouter(profile: "full", dbPath: dbPath, pendingStoreDrainQueue: drainQueue) + router.setDatabase(db) + return router + } + + private func heldDrainQueue() -> DispatchQueue { + let queue = DispatchQueue(label: "caller-model-held-drain-\(UUID().uuidString)") + queue.suspend() + return queue + } + + private func store( + _ router: MCPRouter, + _ arguments: [String: Any], + session: MCPRouter.PaletteSession? = nil, + deferForBackup: Bool = false, + id: Int = 1 + ) throws -> [String: Any] { + let response = router.handle( + [ + "jsonrpc": "2.0", + "id": id, + "method": "tools/call", + "params": ["name": "brain_store", "arguments": arguments] as [String: Any], + ], + session: session, + deferStoreForBackup: deferForBackup + ) + return try XCTUnwrap(response["result"] as? [String: Any]) + } + + private func metadata(ofChunk chunkID: String) throws -> [String: Any] { + var handle: OpaquePointer? + guard sqlite3_open_v2(dbPath, &handle, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + throw NSError(domain: "StoreCallerModelTests", code: 1) + } + defer { sqlite3_close(handle) } + var statement: OpaquePointer? + // Chunk IDs here are server-made or fixture literals; no quoting is needed. + let sql = "SELECT metadata FROM chunks WHERE id = '\(chunkID)'" + guard sqlite3_prepare_v2(handle, sql, -1, &statement, nil) == SQLITE_OK else { + throw NSError(domain: "StoreCallerModelTests", code: 2) + } + defer { sqlite3_finalize(statement) } + guard sqlite3_step(statement) == SQLITE_ROW, let text = sqlite3_column_text(statement, 0) else { + throw NSError(domain: "StoreCallerModelTests", code: 3, userInfo: [NSLocalizedDescriptionKey: "no row \(chunkID)"]) + } + let json = String(cString: text) + return try XCTUnwrap(JSONSerialization.jsonObject(with: Data(json.utf8)) as? [String: Any], json) + } + + private func chunkCount() throws -> Int { + var handle: OpaquePointer? + guard sqlite3_open_v2(dbPath, &handle, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + throw NSError(domain: "StoreCallerModelTests", code: 4) + } + defer { sqlite3_close(handle) } + var statement: OpaquePointer? + guard sqlite3_prepare_v2(handle, "SELECT COUNT(*) FROM chunks", -1, &statement, nil) == SQLITE_OK, + sqlite3_step(statement) == SQLITE_ROW else { + sqlite3_finalize(statement) + throw NSError(domain: "StoreCallerModelTests", code: 5) + } + defer { sqlite3_finalize(statement) } + return Int(sqlite3_column_int64(statement, 0)) + } + + private func assertProvenance( + _ metadata: [String: Any], + model: String, + file: StaticString = #filePath, + line: UInt = #line + ) { + XCTAssertEqual(metadata["model"] as? String, model, "metadata.model", file: file, line: line) + XCTAssertEqual( + metadata["model_source"] as? String, "caller_self_reported", "metadata.model_source", file: file, line: line + ) + } + + private func assertNoProvenance(_ metadata: [String: Any], file: StaticString = #filePath, line: UInt = #line) { + XCTAssertNil(metadata["model"], "no model may be recorded: \(metadata)", file: file, line: line) + XCTAssertNil(metadata["model_source"], "no model_source may be recorded: \(metadata)", file: file, line: line) + } + + // MARK: schema + + func testBrainStoreSchemaAdvertisesOptionalBoundedCallerModel() throws { + let router = MCPRouter(profile: "full") + let response = router.handle(["jsonrpc": "2.0", "id": 1, "method": "tools/list"]) + let result = try XCTUnwrap(response["result"] as? [String: Any]) + let tools = try XCTUnwrap(result["tools"] as? [[String: Any]]) + let store = try XCTUnwrap(tools.first { ($0["name"] as? String) == "brain_store" }) + let schema = try XCTUnwrap(store["inputSchema"] as? [String: Any]) + let properties = try XCTUnwrap(schema["properties"] as? [String: Any]) + let callerModel = try XCTUnwrap(properties["caller_model"] as? [String: Any], "caller_model must be on the schema") + XCTAssertEqual(callerModel["type"] as? String, "string") + XCTAssertEqual(callerModel["maxLength"] as? Int, 128) + XCTAssertEqual(schema["required"] as? [String], ["content"], "caller_model must stay optional") + } + + /// The core boot payload has no room for another property (ratified 1600-byte budget), + /// so core does not advertise caller_model, but a core session can still use it. + func testCorePaletteHidesCallerModelButStillAcceptsIt() throws { + let core = MCPRouter(profile: "core", dbPath: dbPath) + let listed = core.handle(["jsonrpc": "2.0", "id": 1, "method": "tools/list"]) + let tools = try XCTUnwrap((listed["result"] as? [String: Any])?["tools"] as? [[String: Any]]) + let storeTool = try XCTUnwrap(tools.first { ($0["name"] as? String) == "brain_store" }) + let properties = try XCTUnwrap((storeTool["inputSchema"] as? [String: Any])?["properties"] as? [String: Any]) + XCTAssertNil(properties["caller_model"], "core must not spend boot bytes on caller_model") + XCTAssertNotNil(properties["content"]) + + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = MCPRouter(profile: "core", dbPath: dbPath, pendingStoreDrainQueue: drain) + router.setDatabase(db) + let result = try store(router, ["content": "core session with a model", "caller_model": modelA]) + XCTAssertEqual(result["status"] as? String, "STORED") + assertProvenance(try metadata(ofChunk: try XCTUnwrap(result["chunk_id"] as? String)), model: modelA) + + let oversized = try store(router, ["content": "core oversized", "caller_model": String(repeating: "m", count: 129)]) + XCTAssertEqual(oversized["isError"] as? Bool, true, "core still validates against the full schema") + } + + // MARK: direct path + + func testDirectStoreRecordsSelfReportedModel() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + let result = try store(router, ["content": "direct with a model", "caller_model": " \(modelA) "]) + XCTAssertEqual(result["status"] as? String, "STORED") + let chunkID = try XCTUnwrap(result["chunk_id"] as? String) + assertProvenance(try metadata(ofChunk: chunkID), model: modelA) + } + + func testAbsentCallerModelRecordsNothing() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + let result = try store(router, ["content": "direct without a model", "tags": ["agent:someone"]]) + XCTAssertEqual(result["status"] as? String, "STORED") + assertNoProvenance(try metadata(ofChunk: try XCTUnwrap(result["chunk_id"] as? String))) + } + + /// Untrustworthy strings are dropped, and the memory is still stored. + func testUntrustworthyCallerModelIsDroppedButMemoryIsStored() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + let untrusted: [(String, String)] = [ + ("empty", ""), + ("blank", " \t "), + ("synthetic placeholder", ""), + ("control character", "fixture\u{0007}model"), + ("embedded newline", "fixture\nmodel"), + ("zero-width format character", "fixture\u{200B}model"), + ("bidi override", "fixture\u{202E}model"), + ("secret-shaped", "sk-" + String(repeating: "0", count: 40)), + ("over the scalar cap via combining marks", "m" + String(repeating: "\u{0301}", count: 128)), + ] + for (index, (label, value)) in untrusted.enumerated() { + let result = try store(router, ["content": "memory \(index) with \(label) model", "caller_model": value], id: index + 1) + XCTAssertEqual(result["status"] as? String, "STORED", "\(label): the memory must still be stored") + let chunkID = try XCTUnwrap(result["chunk_id"] as? String, label) + assertNoProvenance(try metadata(ofChunk: chunkID)) + } + } + + /// Wrong type or over the advertised maxLength is a bad request: REJECTED, nothing stored. + func testMalformedOrOversizedCallerModelIsRejectedAndStoresNothing() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + let malformed: [(String, Any)] = [ + ("number", 42), + ("array", [modelA]), + ("object", ["name": modelA]), + ("129 characters", String(repeating: "m", count: 129)), + ] + for (label, value) in malformed { + let result = try store(router, ["content": "rejected \(label)", "caller_model": value]) + XCTAssertEqual(result["isError"] as? Bool, true, "\(label) must be rejected") + XCTAssertNil(result["status"], "\(label): REJECTED carries no status") + XCTAssertNil(result["chunk_id"], "\(label): REJECTED carries no chunk_id") + let text = ((result["content"] as? [[String: Any]]) ?? []).compactMap { $0["text"] as? String }.joined() + XCTAssertTrue(text.contains("REJECTED"), "\(label): \(text)") + XCTAssertTrue(text.contains("caller_model"), "\(label): the reason must name the field: \(text)") + } + XCTAssertEqual(try chunkCount(), 0) + XCTAssertFalse(FileManager.default.fileExists(atPath: queuePath.path), "a rejected store must not queue") + } + + func testRetryInTheSameSessionIsIdempotentAndKeepsTheFirstModel() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + let session = router.makePaletteSession() + + let first = try store(router, ["content": "same memory", "caller_model": modelA], session: session, id: 1) + let second = try store(router, ["content": "same memory", "caller_model": modelB], session: session, id: 2) + XCTAssertEqual(first["status"] as? String, "STORED") + XCTAssertEqual(second["status"] as? String, "DUPLICATE") + let chunkID = try XCTUnwrap(first["chunk_id"] as? String) + XCTAssertEqual(second["chunk_id"] as? String, chunkID) + XCTAssertEqual(try chunkCount(), 1) + assertProvenance(try metadata(ofChunk: chunkID), model: modelA) + } + + // MARK: queue + replay + + func testBusyQueuedStoreCarriesModelThroughDurableQueueAndReplay() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + db.failNextStoreWithBusyForTesting = true + let result = try store(router, ["content": "queued under a busy db", "caller_model": modelB]) + XCTAssertEqual(result["status"] as? String, "DEFERRED") + let chunkID = try XCTUnwrap(result["chunk_id"] as? String) + + let line = try XCTUnwrap( + String(contentsOf: queuePath, encoding: .utf8).split(separator: "\n").first.map(String.init) + ) + let queued = try XCTUnwrap(JSONSerialization.jsonObject(with: Data(line.utf8)) as? [String: Any]) + XCTAssertEqual(queued["caller_model"] as? String, modelB, "the durable queue record must keep the model") + + XCTAssertEqual(db.flushPendingStores().map(\.storedChunk.chunkID), [chunkID]) + assertProvenance(try metadata(ofChunk: chunkID), model: modelB) + XCTAssertNotNil(try metadata(ofChunk: chunkID)["brainbar_queue_id"], "queue provenance still recorded") + } + + func testBackupDeferredStoreCarriesModelThroughReplay() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + let result = try store(router, ["content": "deferred for a backup", "caller_model": modelA], deferForBackup: true) + XCTAssertEqual(result["status"] as? String, "DEFERRED") + let chunkID = try XCTUnwrap(result["chunk_id"] as? String) + XCTAssertEqual(db.flushPendingStores().map(\.storedChunk.chunkID), [chunkID]) + assertProvenance(try metadata(ofChunk: chunkID), model: modelA) + } + + func testDirectAndReplayedRowsRecordIdenticalProvenance() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + let direct = try store(router, ["content": "parity direct", "caller_model": modelA], id: 1) + db.failNextStoreWithBusyForTesting = true + let deferred = try store(router, ["content": "parity queued", "caller_model": modelA], id: 2) + XCTAssertEqual(deferred["status"] as? String, "DEFERRED") + let deferredID = try XCTUnwrap(deferred["chunk_id"] as? String) + XCTAssertEqual(db.flushPendingStores().map(\.storedChunk.chunkID), [deferredID]) + + let directMeta = try metadata(ofChunk: try XCTUnwrap(direct["chunk_id"] as? String)) + let replayMeta = try metadata(ofChunk: deferredID) + XCTAssertEqual(directMeta["model"] as? String, replayMeta["model"] as? String) + XCTAssertEqual(directMeta["model_source"] as? String, replayMeta["model_source"] as? String) + assertProvenance(replayMeta, model: modelA) + } + + func testLineSkippedByAFlushKeepsItsModelWhenRewritten() throws { + // Created first so it is resumed after the database closes, as in the reliability tests. + let drain = heldDrainQueue() + defer { drain.resume() } + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let router = makeRouter(db, drainQueue: drain) + + db.failNextStoreWithBusyForTesting = true + let result = try store(router, ["content": "kept queued once", "caller_model": modelB]) + let chunkID = try XCTUnwrap(result["chunk_id"] as? String) + + XCTAssertTrue(db.flushPendingStores(excludingChunkIDs: [chunkID]).isEmpty) + let rewritten = try String(contentsOf: queuePath, encoding: .utf8) + XCTAssertTrue(rewritten.contains("\"caller_model\":\"\(modelB)\""), rewritten) + + XCTAssertEqual(db.flushPendingStores().map(\.storedChunk.chunkID), [chunkID]) + assertProvenance(try metadata(ofChunk: chunkID), model: modelB) + } + + // MARK: hand-written queue lines (old records, malformed records) + + private func writeQueueLines(_ records: [[String: Any]]) throws { + var data = Data() + for record in records { + data.append(try JSONSerialization.data(withJSONObject: record)) + data.append(0x0A) + } + try data.write(to: queuePath) + } + + func testQueuedRecordsReplayWithValidatedModelOnly() throws { + try writeQueueLines([ + // Written before caller_model existed: replays exactly as before. + ["content": "legacy line", "tags": [], "importance": 5, "source": "mcp", "chunk_id": "cm-legacy"], + ["content": "valid model line", "tags": [], "importance": 5, "source": "mcp", + "chunk_id": "cm-valid", "caller_model": modelA], + // Malformed model values: the model is dropped, the memory is not stranded. + ["content": "numeric model line", "tags": [], "importance": 5, "source": "mcp", + "chunk_id": "cm-number", "caller_model": 42], + ["content": "null model line", "tags": [], "importance": 5, "source": "mcp", + "chunk_id": "cm-null", "caller_model": NSNull()], + ["content": "control model line", "tags": [], "importance": 5, "source": "mcp", + "chunk_id": "cm-control", "caller_model": "fixture\u{0000}model"], + ["content": "oversized model line", "tags": [], "importance": 5, "source": "mcp", + "chunk_id": "cm-oversized", "caller_model": String(repeating: "m", count: 4_096)], + ]) + + let db = BrainDatabase(path: dbPath) + defer { db.close() } + let flushed = db.flushPendingStores().map(\.storedChunk.chunkID) + XCTAssertEqual( + Set(flushed), ["cm-legacy", "cm-valid", "cm-number", "cm-null", "cm-control", "cm-oversized"], + "every line must replay; a bad model never strands a memory" + ) + XCTAssertFalse(FileManager.default.fileExists(atPath: queuePath.path), "nothing may remain queued") + + assertProvenance(try metadata(ofChunk: "cm-valid"), model: modelA) + for chunkID in ["cm-legacy", "cm-number", "cm-null", "cm-control", "cm-oversized"] { + assertNoProvenance(try metadata(ofChunk: chunkID)) + } + + // Replay is idempotent: a second flush finds nothing and changes nothing. + XCTAssertTrue(db.flushPendingStores().isEmpty) + XCTAssertEqual(try chunkCount(), 6) + assertProvenance(try metadata(ofChunk: "cm-valid"), model: modelA) + } +} diff --git a/docs/mcp-tools.md b/docs/mcp-tools.md index 8632c8927..84ce8a4d6 100644 --- a/docs/mcp-tools.md +++ b/docs/mcp-tools.md @@ -79,6 +79,7 @@ Persist a memory (idea, decision, learning, mistake, etc.) for future retrieval. | `tags` | array[string] | No | Tags for categorization | | `importance` | integer | No | Importance score 1-10 (auto-detected if omitted) | | `entity_id` | string | No | Link memory to a KG entity | +| `caller_model` | string | No | Your own model ID, self-reported (max 128). Stored as `metadata.model` with `metadata.model_source = "caller_self_reported"` on direct, queued and replayed writes. Blank, control/format characters, `` or secret-shaped values are dropped and the memory is still stored. A non-string or over-long value is `REJECTED`. Never inferred, and unrelated to the enrichment columns. Advertised once the palette is expanded (or under the full profile); the core boot palette accepts it but does not list it, to stay inside the 1600-byte boot budget. | Issue type supports lifecycle tracking: status (open→in_progress→done→archived), severity (critical/high/medium/low), and code references (file_path, function_name, line_number).