diff --git a/brain-bar/Sources/BrainBar/BrainBarServer.swift b/brain-bar/Sources/BrainBar/BrainBarServer.swift index d1ccdfe5..b6611598 100644 --- a/brain-bar/Sources/BrainBar/BrainBarServer.swift +++ b/brain-bar/Sources/BrainBar/BrainBarServer.swift @@ -150,6 +150,7 @@ final class BrainBarServer: @unchecked Sendable { struct ClientState { var source: DispatchSourceRead var framing: MCPFraming + let paletteSession: MCPRouter.PaletteSession /// Whether this client uses Content-Length framing (LSP-style). /// false = newline-delimited JSON-RPC (Claude Code v2.1+). var usesContentLengthFraming: Bool = true @@ -441,7 +442,11 @@ final class BrainBarServer: @unchecked Sendable { } readSource.resume() - clients[clientFD] = ClientState(source: readSource, framing: MCPFraming()) + clients[clientFD] = ClientState( + source: readSource, + framing: MCPFraming(), + paletteSession: router.makePaletteSession() + ) NSLog("[BrainBar] Client connected (fd: %d)", clientFD) debugLog("CLIENT CONNECTED fd=\(clientFD) (total clients: \(clients.count))") } @@ -493,16 +498,29 @@ final class BrainBarServer: @unchecked Sendable { } private func handleMessage(fd: Int32, request: [String: Any]) -> [String: Any] { + let paletteSession = clients[fd]?.paletteSession if let toolCall = parseToolCall(request) { switch toolCall.name { case "brain_subscribe": + guard let paletteSession, + router.isToolExposed(toolCall.name, session: paletteSession) else { + return router.handle(request, session: paletteSession) + } return handleSubscribeTool(fd: fd, id: request["id"], arguments: toolCall.arguments) case "brain_unsubscribe": + guard let paletteSession, + router.isToolExposed(toolCall.name, session: paletteSession) else { + return router.handle(request, session: paletteSession) + } return handleUnsubscribeTool(fd: fd, id: request["id"], arguments: toolCall.arguments) case "brain_ack": + guard let paletteSession, + router.isToolExposed(toolCall.name, session: paletteSession) else { + return router.handle(request, session: paletteSession) + } return handleAckTool(id: request["id"], arguments: toolCall.arguments) default: - let response = router.handle(request) + let response = router.handle(request, session: paletteSession) if toolCall.name == "brain_store", !isToolError(response) { publishStoredChunks(response: response, arguments: toolCall.arguments) } @@ -520,7 +538,7 @@ final class BrainBarServer: @unchecked Sendable { break } } - return router.handle(request) + return router.handle(request, session: paletteSession) } @discardableResult diff --git a/brain-bar/Sources/BrainBar/MCPRouter.swift b/brain-bar/Sources/BrainBar/MCPRouter.swift index c1f7d9ad..c7f57078 100644 --- a/brain-bar/Sources/BrainBar/MCPRouter.swift +++ b/brain-bar/Sources/BrainBar/MCPRouter.swift @@ -10,6 +10,32 @@ import Foundation final class MCPRouter: @unchecked Sendable { + private enum ToolProfile { + case core + case full + } + + private static let profileEnvironmentKey = "BRAINLAYER_MCP_PROFILE" + private static let coreToolNames: Set = [ + "brain_search", + "brain_store", + "brain_recall", + "brain_expand", + ] + private static let coreToolDescriptions: [String: String] = [ + "brain_search": "Search memory.", + "brain_store": "Store memory.", + "brain_recall": "Recall context.", + "brain_expand": "Expand chunk.", + ] + nonisolated(unsafe) private static let expandPaletteToolDefinition: [String: Any] = [ + "name": "expand_palette", + "description": "Expose all tools.", + "inputSchema": [ + "type": "object", + ] as [String: Any], + ] + // Keep contested MCP writes interactive: make one short attempt, then queue // for replay so brain_store returns well under the 1s prompt-queue budget. // Longer contention handling belongs in the deferred single-writer/backpressure @@ -29,6 +55,25 @@ final class MCPRouter: @unchecked Sendable { private static let pendingStoreDrainInitialDelay: TimeInterval = 0.25 private static let pendingStoreDrainMaxDelay: TimeInterval = 30.0 + final class PaletteSession: @unchecked Sendable { + private let lock = NSLock() + private var expanded = false + + func isExpanded() -> Bool { + lock.lock() + defer { lock.unlock() } + return expanded + } + + func expand() -> Bool { + lock.lock() + defer { lock.unlock() } + guard !expanded else { return false } + expanded = true + return true + } + } + private struct ToolOutput { let text: String let metadata: [String: Any] @@ -94,6 +139,8 @@ final class MCPRouter: @unchecked Sendable { private let hybridSearchClient: HybridSearchClientProtocol? private let dbPath: String? private let hybridSearchBudget: TimeInterval + private let toolProfile: ToolProfile + private let defaultPaletteSession = PaletteSession() let entityCache = EntityCache() private static let defaultStringMaxLength = 256 private static let defaultStringArrayMaxItems = 100 @@ -121,15 +168,108 @@ final class MCPRouter: @unchecked Sendable { ] init( + profile: String? = nil, hybridSearchClient: HybridSearchClientProtocol? = nil, hybridSearchBudget: TimeInterval = 0.8, dbPath: String? = nil ) { + self.toolProfile = Self.resolveToolProfile(profile) self.hybridSearchClient = hybridSearchClient self.hybridSearchBudget = max(0.001, hybridSearchBudget) self.dbPath = dbPath } + private static func resolveToolProfile(_ explicitProfile: String?) -> ToolProfile { + let environmentProfile = ProcessInfo.processInfo.environment[profileEnvironmentKey] + let rawProfile = explicitProfile ?? environmentProfile + let normalized = rawProfile? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() ?? "" + + switch normalized { + case "full", "operator": + return .full + case "", "core": + return .core + default: + if explicitProfile == nil, let rawProfile = environmentProfile { + let message = "BrainBar: unknown \(profileEnvironmentKey)=\(rawProfile); using core profile\n" + FileHandle.standardError.write(Data(message.utf8)) + } + return .core + } + } + + func makePaletteSession() -> PaletteSession { + PaletteSession() + } + + private func exposedToolDefinitions(for session: PaletteSession) -> [[String: Any]] { + if toolProfile == .full || session.isExpanded() { + return Self.toolDefinitions + } + + let coreDefinitions = Self.toolDefinitions.filter { definition in + guard let name = definition["name"] as? String else { return false } + return Self.coreToolNames.contains(name) + }.map(Self.compactCoreToolDefinition) + return coreDefinitions + [Self.expandPaletteToolDefinition] + } + + func isToolExposed(_ name: String, session: PaletteSession) -> Bool { + exposedToolDefinitions(for: session).contains { ($0["name"] as? String) == name } + } + + private static func compactCoreToolDefinition(_ definition: [String: Any]) -> [String: Any] { + guard let name = definition["name"] as? String else { return definition } + + var compact: [String: Any] = ["name": name] + compact["description"] = coreToolDescriptions[name] + if let inputSchema = definition["inputSchema"] as? [String: Any] { + compact["inputSchema"] = removingDescriptions(from: inputSchema) + } + return compact + } + + private static func removingDescriptions(from value: Any) -> Any { + if let dictionary = value as? [String: Any] { + return dictionary.reduce(into: [String: Any]()) { result, entry in + guard entry.key != "description" else { return } + result[entry.key] = removingDescriptions(from: entry.value) + } + } + if let array = value as? [Any] { + return array.map(removingDescriptions(from:)) + } + return value + } + + private func expandPalette(session: PaletteSession) -> ToolOutput { + guard session.expand() else { + return ToolOutput( + text: "BrainLayer tool palette is already expanded.", + metadata: [ + "expanded": false, + "already_expanded": true, + "registered_tools": [String](), + ] + ) + } + + let deferredNames = Self.toolDefinitions.compactMap { definition -> String? in + guard let name = definition["name"] as? String else { return nil } + return Self.coreToolNames.contains(name) ? nil : name + } + return ToolOutput( + text: "Expanded BrainLayer tool palette.", + metadata: [ + "expanded": true, + "already_expanded": false, + "registered_tools": deferredNames, + ] + ) + } + /// Inject database for tool handlers + load entity cache. func setDatabase(_ db: BrainDatabase) { setDatabases(write: db, read: db) @@ -160,7 +300,8 @@ final class MCPRouter: @unchecked Sendable { /// Handle a parsed JSON-RPC request and return a response. /// Returns empty dict for notifications (no id). - func handle(_ request: [String: Any]) -> [String: Any] { + func handle(_ request: [String: Any], session: PaletteSession? = nil) -> [String: Any] { + let paletteSession = session ?? defaultPaletteSession guard let method = request["method"] as? String else { return jsonRPCError(id: request["id"], code: -32600, message: "Invalid request: missing method") } @@ -181,9 +322,13 @@ final class MCPRouter: @unchecked Sendable { // If a client sends this with an id, ack it so it doesn't hang. return jsonRPCResult(id: id, result: [:] as [String: Any]) case "tools/list": - return handleToolsList(id: id) + return handleToolsList(id: id, session: paletteSession) case "tools/call": - return handleToolsCall(id: id, params: request["params"] as? [String: Any] ?? [:]) + return handleToolsCall( + id: id, + params: request["params"] as? [String: Any] ?? [:], + session: paletteSession + ) case "resources/list": return handleResourcesList(id: id) case "prompts/list": @@ -204,7 +349,7 @@ final class MCPRouter: @unchecked Sendable { "result": [ "protocolVersion": "2024-11-05", "capabilities": [ - "tools": ["listChanged": false], + "tools": ["listChanged": true], "experimental": [ "claude/channel": [:] as [String: Any] ] @@ -219,12 +364,12 @@ final class MCPRouter: @unchecked Sendable { // MARK: - tools/list - private func handleToolsList(id: Any) -> [String: Any] { + private func handleToolsList(id: Any, session: PaletteSession) -> [String: Any] { return [ "jsonrpc": "2.0", "id": id, "result": [ - "tools": Self.toolDefinitions + "tools": exposedToolDefinitions(for: session) ] ] } @@ -236,15 +381,22 @@ final class MCPRouter: @unchecked Sendable { // MARK: - tools/call - private func handleToolsCall(id: Any, params: [String: Any]) -> [String: Any] { + private func handleToolsCall(id: Any, params: [String: Any], session: PaletteSession) -> [String: Any] { guard let toolName = params["name"] as? String else { return jsonRPCError(id: id, code: -32602, message: "Missing tool name") } let arguments = params["arguments"] as? [String: Any] ?? [:] + if toolName == "expand_palette" { + guard toolProfile == .core else { + return jsonRPCError(id: id, code: -32601, message: "Unknown tool: \(toolName)") + } + return toolCallResult(id: id, output: expandPalette(session: session)) + } + // Check tool exists - guard Self.toolDefinitions.contains(where: { ($0["name"] as? String) == toolName }) else { + guard isToolExposed(toolName, session: session) else { return jsonRPCError(id: id, code: -32601, message: "Unknown tool: \(toolName)") } @@ -252,19 +404,7 @@ final class MCPRouter: @unchecked Sendable { do { try Self.validate(arguments: arguments, for: toolName) let output = try dispatchTool(name: toolName, arguments: arguments) - var result: [String: Any] = [ - "content": [ - ["type": "text", "text": output.text] - ] - ] - for (key, value) in output.metadata { - result[key] = value - } - return [ - "jsonrpc": "2.0", - "id": id, - "result": result - ] + return toolCallResult(id: id, output: output) } catch { return [ "jsonrpc": "2.0", @@ -279,6 +419,22 @@ final class MCPRouter: @unchecked Sendable { } } + private func toolCallResult(id: Any, output: ToolOutput) -> [String: Any] { + var result: [String: Any] = [ + "content": [ + ["type": "text", "text": output.text] + ] + ] + for (key, value) in output.metadata { + result[key] = value + } + return [ + "jsonrpc": "2.0", + "id": id, + "result": result, + ] + } + private func dispatchTool(name: String, arguments: [String: Any]) throws -> ToolOutput { switch name { case "brain_search": diff --git a/brain-bar/Tests/BrainBarTests/MCPRouterTests.swift b/brain-bar/Tests/BrainBarTests/MCPRouterTests.swift index 27291e4d..60419d55 100644 --- a/brain-bar/Tests/BrainBarTests/MCPRouterTests.swift +++ b/brain-bar/Tests/BrainBarTests/MCPRouterTests.swift @@ -58,6 +58,27 @@ private func collectStringArrays(in schema: [String: Any], path: String = "") -> } final class MCPRouterTests: XCTestCase { + private static let coreToolNames = [ + "brain_search", + "brain_store", + "brain_recall", + "brain_expand", + "expand_palette", + ] + + private func listedTools(_ router: MCPRouter) -> [[String: Any]] { + let response = router.handle([ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + ]) + return (response["result"] as? [String: Any])?["tools"] as? [[String: Any]] ?? [] + } + + private func listedToolNames(_ router: MCPRouter) -> [String] { + listedTools(router).compactMap { $0["name"] as? String } + } + private func assertDeferredBrainStoreReceipt( _ result: [String: Any], chunkID: String, @@ -88,7 +109,7 @@ final class MCPRouterTests: XCTestCase { // MARK: - Initialize func testInitializeReturnsProtocolVersion() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 1, @@ -109,7 +130,8 @@ final class MCPRouterTests: XCTestCase { // Must declare tool capabilities let capabilities = result?["capabilities"] as? [String: Any] - XCTAssertNotNil(capabilities?["tools"]) + let tools = capabilities?["tools"] as? [String: Any] + XCTAssertEqual(tools?["listChanged"] as? Bool, true) XCTAssertNil(capabilities?["resources"], "Tag resources should not be advertised during initialize") let experimental = capabilities?["experimental"] as? [String: Any] XCTAssertEqual((experimental?["claude/channel"] as? [String: Any])?.isEmpty, true) @@ -123,7 +145,7 @@ final class MCPRouterTests: XCTestCase { try db.insertChunk(id: "tagged-1", content: "Tagged chunk", sessionId: "s1", project: "test", contentType: "assistant_text", importance: 5, tags: "[\"agent-message\"]") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -140,8 +162,75 @@ final class MCPRouterTests: XCTestCase { // MARK: - Tools list + func testCoreProfilesExposeOnlyCoreFourAndExpansionControl() throws { + XCTAssertEqual(listedToolNames(MCPRouter(profile: "core")), Self.coreToolNames) + for profile in ["", " ", "core", "bogus"] { + XCTAssertEqual(listedToolNames(MCPRouter(profile: profile)), Self.coreToolNames) + } + } + + func testEnvironmentProfileControlsDefaultRouter() throws { + let environmentKey = "BRAINLAYER_MCP_PROFILE" + let originalProfile = ProcessInfo.processInfo.environment[environmentKey] + defer { + if let originalProfile { + setenv(environmentKey, originalProfile, 1) + } else { + unsetenv(environmentKey) + } + } + + setenv(environmentKey, "operator", 1) + let canonicalNames = MCPRouter.toolDefinitions.compactMap { $0["name"] as? String } + XCTAssertEqual(listedToolNames(MCPRouter()), canonicalNames) + + setenv(environmentKey, "core", 1) + XCTAssertEqual(listedToolNames(MCPRouter()), Self.coreToolNames) + } + + func testFullAndOperatorProfilesPreserveCanonicalInventory() throws { + let canonicalNames = MCPRouter.toolDefinitions.compactMap { $0["name"] as? String } + + XCTAssertEqual(canonicalNames.count, 17) + XCTAssertEqual(listedToolNames(MCPRouter(profile: "full")), canonicalNames) + XCTAssertEqual(listedToolNames(MCPRouter(profile: "operator")), canonicalNames) + } + + func testCoreToolsListStaysWithinBootBudget() throws { + let tools = listedTools(MCPRouter(profile: "core")) + let data = try JSONSerialization.data(withJSONObject: tools, options: [.sortedKeys]) + + XCTAssertLessThanOrEqual(data.count, 1_500) + } + + func testCorePaletteExpandsOnceAndDispatchesDeferredTools() throws { + let router = MCPRouter(profile: "core") + + XCTAssertEqual(listedToolNames(router), Self.coreToolNames) + let deferredBeforeExpansion = router.handle(toolCall(id: 20, name: "brain_tags", arguments: [:])) + XCTAssertEqual((deferredBeforeExpansion["error"] as? [String: Any])?["code"] as? Int, -32601) + + let firstExpansion = router.handle(toolCall(id: 21, name: "expand_palette", arguments: [:])) + let firstResult = try XCTUnwrap(firstExpansion["result"] as? [String: Any]) + XCTAssertEqual(firstResult["expanded"] as? Bool, true) + XCTAssertEqual(firstResult["already_expanded"] as? Bool, false) + XCTAssertEqual((firstResult["registered_tools"] as? [String])?.count, 13) + XCTAssertEqual(listedToolNames(router), MCPRouter.toolDefinitions.compactMap { $0["name"] as? String }) + + let deferredAfterExpansion = router.handle(toolCall(id: 22, name: "brain_tags", arguments: [:])) + XCTAssertNil(deferredAfterExpansion["error"]) + XCTAssertNotNil(deferredAfterExpansion["result"]) + + let secondExpansion = router.handle(toolCall(id: 23, name: "expand_palette", arguments: [:])) + let secondResult = try XCTUnwrap(secondExpansion["result"] as? [String: Any]) + XCTAssertEqual(secondResult["expanded"] as? Bool, false) + XCTAssertEqual(secondResult["already_expanded"] as? Bool, true) + XCTAssertEqual(secondResult["registered_tools"] as? [String], []) + XCTAssertEqual(listedToolNames(router).count, 17) + } + func testToolsListReturnsAllTools() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 2, @@ -167,7 +256,7 @@ final class MCPRouterTests: XCTestCase { } func testEncodedToolsListEnvelopeStartsWithJSONRPC() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let response = router.handle([ "jsonrpc": "2.0", "id": 2, @@ -185,7 +274,7 @@ final class MCPRouterTests: XCTestCase { } func testToolsListPreservesCanonicalAnnotations() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let response = router.handle([ "jsonrpc": "2.0", "id": 2, @@ -204,7 +293,7 @@ final class MCPRouterTests: XCTestCase { } func testEachToolHasInputSchema() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 3, @@ -222,7 +311,7 @@ final class MCPRouterTests: XCTestCase { } func testEachToolHasExpectedAnnotations() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let response = router.handle([ "jsonrpc": "2.0", "id": 3, @@ -315,7 +404,7 @@ final class MCPRouterTests: XCTestCase { // MARK: - Tools call func testToolsCallDispatchesToHandler() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 4, @@ -363,7 +452,7 @@ final class MCPRouterTests: XCTestCase { ] ) ) - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle([ @@ -422,7 +511,7 @@ final class MCPRouterTests: XCTestCase { let helperText = "python helper canonical KG/search output" let helper = RecordingHybridSearchClient(response: HybridSearchResponse(text: helperText)) - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle([ @@ -460,7 +549,7 @@ final class MCPRouterTests: XCTestCase { ) let helper = RecordingHybridSearchClient() - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle([ @@ -506,7 +595,7 @@ No results found. """ ) ) - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle(toolCall(id: 179, name: "brain_search", arguments: [ @@ -527,7 +616,7 @@ No results found. let db = BrainDatabase(path: tempDB) defer { db.close() } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let fullContent = "Sensitive full content must not be echoed in the MCP store response" @@ -574,7 +663,7 @@ No results found. ) let helper = RecordingHybridSearchClient() - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle([ @@ -625,7 +714,7 @@ No results found. XCTAssertEqual(try sqlitePragma(handle: readDB.dbHandle, name: "busy_timeout"), "250") writeDB.close() - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabases(write: writeDB, read: readDB) let searchText = try toolText(router.handle(toolCall(id: 301, name: "brain_search", arguments: [ @@ -684,7 +773,7 @@ No results found. sql: "UPDATE chunks SET summary = 'Short generated summary' WHERE id = 'expand-full-target'" ) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let expandText = try toolText(router.handle(toolCall(id: 307, name: "brain_expand", arguments: [ "chunk_id": "expand-full-target" @@ -725,7 +814,7 @@ No results found. // here and failed with SQLITE_READONLY. XCTAssertEqual(try sqlitePragma(handle: readDB.dbHandle, name: "query_only"), "1") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabases(write: writeDB, read: readDB) let response = router.handle(toolCall(id: 401, name: "brain_search", arguments: [ @@ -768,7 +857,7 @@ No results found. response: HybridSearchResponse(text: "slow helper should not win"), delay: 0.50 ) - let router = MCPRouter(hybridSearchClient: helper, hybridSearchBudget: 0.05) + let router = MCPRouter(profile: "full", hybridSearchClient: helper, hybridSearchBudget: 0.05) router.setDatabases(write: writeDB, read: readDB) let started = Date() @@ -803,7 +892,7 @@ No results found. response: HybridSearchResponse(text: "helper should not be called while warming") ) helper.ready = false - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle(toolCall(id: 308, name: "brain_search", arguments: [ @@ -836,7 +925,7 @@ No results found. let helper = RecordingHybridSearchClient( response: HybridSearchResponse(text: "helper should not be called") ) - let router = MCPRouter(hybridSearchClient: helper) + let router = MCPRouter(profile: "full", hybridSearchClient: helper) router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -875,7 +964,7 @@ No results found. } try sqliteExec(path: tempDB, sql: "DELETE FROM chunks_fts_trigram") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -910,7 +999,7 @@ No results found. defer { db.close() } try sqliteExec(path: tempDB, sql: "DROP TABLE IF EXISTS chunks_fts_trigram") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -933,7 +1022,7 @@ No results found. } func testToolsCallUnknownToolReturnsError() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 5, @@ -952,7 +1041,7 @@ No results found. } func testBrainDigestRejectsOversizedContentAtSchemaLayer() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let response = router.handle([ "jsonrpc": "2.0", "id": 6, @@ -975,7 +1064,7 @@ No results found. } func testBrainDigestSchemaExposesProjectAndTitle() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let response = router.handle([ "jsonrpc": "2.0", "id": 7, @@ -999,7 +1088,7 @@ No results found. let db = BrainDatabase(path: tempDB) defer { db.close() } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1025,7 +1114,7 @@ No results found. } func testBrainSubscribeToolIsServerHandled() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 7, @@ -1045,7 +1134,7 @@ No results found. } func testBrainUnsubscribeToolIsServerHandled() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 8, @@ -1064,7 +1153,7 @@ No results found. } func testBrainAckToolIsServerHandled() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 9, @@ -1086,7 +1175,7 @@ No results found. // MARK: - Unknown method func testUnknownMethodReturnsError() throws { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "id": 6, @@ -1111,7 +1200,7 @@ No results found. try db.insertChunk(id: "f-1", content: "Socket handling code", sessionId: "s1", project: "brainbar", contentType: "assistant_text", importance: 5) try db.insertChunk(id: "f-2", content: "Socket connection code", sessionId: "s2", project: "other-proj", contentType: "assistant_text", importance: 5) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1142,7 +1231,7 @@ No results found. try db.insertChunk(id: "detail-1", content: "Socket handling code", sessionId: "s1", project: "brainbar", contentType: "assistant_text", importance: 5) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1172,7 +1261,7 @@ No results found. try db.insertChunk(id: "i-1", content: "Critical security finding", sessionId: "s1", project: "test", contentType: "assistant_text", importance: 9) try db.insertChunk(id: "i-2", content: "Security review notes", sessionId: "s2", project: "test", contentType: "assistant_text", importance: 3) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1204,7 +1293,7 @@ No results found. _ = try db.store(content: "Sagit meeting notes", tags: ["meeting"], importance: 6, source: "whatsapp") _ = try db.store(content: "Sagit meeting notes", tags: ["meeting"], importance: 6, source: "youtube") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1242,7 +1331,7 @@ No results found. importance: 8 ) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1279,7 +1368,7 @@ No results found. try db.insertEntity(id: "tool-1", type: "tool", name: "Claude Code") try db.insertRelation(sourceId: "proj-1", targetId: "tool-1", relationType: "used_by") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1310,7 +1399,7 @@ No results found. try db.insertChunk(id: "s-1", content: "Search result one", sessionId: "session-1", project: "brainlayer", contentType: "assistant_text", importance: 5) try db.insertChunk(id: "s-2", content: "Search result two", sessionId: "session-2", project: "orchestrator", contentType: "user_message", importance: 4) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1345,7 +1434,7 @@ No results found. timestamp: "2026-03-31T04:03:00.000Z" ) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1382,7 +1471,7 @@ No results found. } try db.acknowledge(agentID: "agent-1", seq: readSeq) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1416,7 +1505,7 @@ No results found. defer { db.close() } try db.insertChunk(id: "archive-target", content: "Archive this stale memory", sessionId: "s1", project: "test", contentType: "assistant_text", importance: 5) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let archiveResponse = router.handle([ @@ -1454,7 +1543,7 @@ No results found. try db.insertChunk(id: "old-chunk", content: "TechGym guidance old version", sessionId: "s1", project: "test", contentType: "assistant_text", importance: 5) try db.insertChunk(id: "new-chunk", content: "TechGym guidance new version", sessionId: "s2", project: "test", contentType: "assistant_text", importance: 9) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let supersedeResponse = router.handle([ @@ -1510,7 +1599,7 @@ No results found. try db.insertChunk(id: "mem-sagit-1", content: "Sagit Stern gave the TechGym lecture about search ranking.", sessionId: "s1", project: "test", contentType: "assistant_text", importance: 8) try db.linkEntityChunk(entityId: "person-sagit", chunkId: "mem-sagit-1", relevance: 0.9) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1544,7 +1633,7 @@ No results found. importance: 5 ) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let enrichResponse = router.handle([ "jsonrpc": "2.0", @@ -1597,7 +1686,7 @@ No results found. ) db.exec("UPDATE chunks SET enrich_status = 'duplicate' WHERE id = 'enrich-duplicate'") - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1630,7 +1719,7 @@ No results found. importance: 5 ) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1662,7 +1751,7 @@ No results found. importance: 5 ) - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ "jsonrpc": "2.0", @@ -1688,7 +1777,7 @@ No results found. let dbPath = tempDir.appendingPathComponent("brainbar.db").path let queuePath = tempDir.appendingPathComponent("pending-stores.jsonl") - let router = MCPRouter(dbPath: dbPath) + let router = MCPRouter(profile: "full", dbPath: dbPath) let response = router.handle([ "jsonrpc": "2.0", @@ -1744,7 +1833,7 @@ No results found. let db = BrainDatabase(path: dbPath) db.close() - let router = MCPRouter(dbPath: dbPath) + let router = MCPRouter(profile: "full", dbPath: dbPath) router.setDatabase(db) let response = router.handle([ @@ -1802,7 +1891,7 @@ No results found. defer { db.close() } db.failNextStoreWithBusyForTesting = true - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -1876,7 +1965,7 @@ No results found. } } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let started = Date() @@ -1937,7 +2026,7 @@ No results found. } } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let queuedContent = "Queued current write should drain after the writer lock clears" @@ -1994,7 +2083,7 @@ No results found. defer { db.close() } db.failNextStoreWithBusyForTesting = true - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let queuedContent = "Queued current write survives one unreadable pending store queue snapshot" @@ -2056,7 +2145,7 @@ No results found. let db = BrainDatabase(path: dbPath) defer { db.close() } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -2108,7 +2197,7 @@ No results found. defer { db.close() } db.failNextStoreWithBusyForTesting = true - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -2171,7 +2260,7 @@ No results found. let db = BrainDatabase(path: dbPath) defer { db.close() } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -2225,7 +2314,7 @@ No results found. let db = BrainDatabase(path: dbPath) defer { db.close() } - let router = MCPRouter() + let router = MCPRouter(profile: "full") router.setDatabase(db) let response = router.handle([ @@ -2505,7 +2594,7 @@ No results found. // MARK: - Notifications (no id) func testNotificationDoesNotRequireResponse() { - let router = MCPRouter() + let router = MCPRouter(profile: "full") let request: [String: Any] = [ "jsonrpc": "2.0", "method": "notifications/initialized", diff --git a/brain-bar/Tests/BrainBarTests/SocketIntegrationTests.swift b/brain-bar/Tests/BrainBarTests/SocketIntegrationTests.swift index 03963134..2929e00d 100644 --- a/brain-bar/Tests/BrainBarTests/SocketIntegrationTests.swift +++ b/brain-bar/Tests/BrainBarTests/SocketIntegrationTests.swift @@ -12,9 +12,12 @@ final class SocketIntegrationTests: XCTestCase { var server: BrainBarServer! var db: BrainDatabase! var tempDBPath: String! + var originalMCPProfile: String? override func setUp() { super.setUp() + originalMCPProfile = ProcessInfo.processInfo.environment["BRAINLAYER_MCP_PROFILE"] + setenv("BRAINLAYER_MCP_PROFILE", "full", 1) tempDBPath = NSTemporaryDirectory() + "brainbar-integration-\(UUID().uuidString).db" db = BrainDatabase(path: tempDBPath) server = BrainBarServer(socketPath: testSocketPath, dbPath: tempDBPath, database: db) @@ -28,6 +31,11 @@ final class SocketIntegrationTests: XCTestCase { try? FileManager.default.removeItem(atPath: tempDBPath) try? FileManager.default.removeItem(atPath: tempDBPath + "-wal") try? FileManager.default.removeItem(atPath: tempDBPath + "-shm") + if let originalMCPProfile { + setenv("BRAINLAYER_MCP_PROFILE", originalMCPProfile, 1) + } else { + unsetenv("BRAINLAYER_MCP_PROFILE") + } super.tearDown() } @@ -111,6 +119,74 @@ final class SocketIntegrationTests: XCTestCase { } } + func testCorePaletteExpansionIsIsolatedPerSocketClient() throws { + restartServer(profile: "core") + let firstFD = try connectClient() + defer { close(firstFD) } + let secondFD = try connectClient() + defer { close(secondFD) } + try initializeClient(fd: firstFD, name: "first-core-client") + try initializeClient(fd: secondFD, name: "second-core-client") + + let coreNames: Set = [ + "brain_search", "brain_store", "brain_recall", "brain_expand", "expand_palette", + ] + XCTAssertEqual(Set(try listedToolNames(on: firstFD, id: 2)), coreNames) + XCTAssertEqual(Set(try listedToolNames(on: secondFD, id: 2)), coreNames) + + try sendMCPRequest(on: firstFD, request: [ + "jsonrpc": "2.0", "id": 3, "method": "tools/call", + "params": ["name": "expand_palette", "arguments": [:] as [String: Any]], + ]) + let expansion = try readMCPMessage(fd: firstFD) + XCTAssertEqual((expansion["result"] as? [String: Any])?["expanded"] as? Bool, true) + + XCTAssertEqual(try listedToolNames(on: firstFD, id: 4).count, 17) + XCTAssertEqual(Set(try listedToolNames(on: secondFD, id: 3)), coreNames) + } + + func testCoreProfileRejectsServerHandledToolsUntilExpansion() throws { + restartServer(profile: "core") + let fd = try connectClient() + defer { close(fd) } + try initializeClient(fd: fd, name: "core-subscription-client") + + let serverHandledTools: [(name: String, arguments: [String: Any])] = [ + ("brain_subscribe", ["agent_id": "core-agent", "tags": ["agent-message"]]), + ("brain_unsubscribe", ["agent_id": "core-agent", "tags": ["agent-message"]]), + ("brain_ack", ["agent_id": "core-agent", "seq": 1]), + ] + for (offset, tool) in serverHandledTools.enumerated() { + try sendMCPRequest(on: fd, request: [ + "jsonrpc": "2.0", "id": 2 + offset, "method": "tools/call", + "params": ["name": tool.name, "arguments": tool.arguments] as [String: Any], + ]) + let deferred = try readMCPMessage(fd: fd) + XCTAssertEqual( + (deferred["error"] as? [String: Any])?["code"] as? Int, + -32601, + "\(tool.name) should stay deferred until this client expands its palette" + ) + } + + try sendMCPRequest(on: fd, request: [ + "jsonrpc": "2.0", "id": 5, "method": "tools/call", + "params": ["name": "expand_palette", "arguments": [:] as [String: Any]], + ]) + _ = try readMCPMessage(fd: fd) + + try sendMCPRequest(on: fd, request: [ + "jsonrpc": "2.0", "id": 6, "method": "tools/call", + "params": [ + "name": "brain_subscribe", + "arguments": ["agent_id": "core-agent", "tags": ["agent-message"]] as [String: Any], + ], + ]) + let expanded = try readMCPMessage(fd: fd) + XCTAssertNil(expanded["error"]) + XCTAssertNotNil(expanded["result"]) + } + func testRawLineToolsListCompactsForClaudeExtensionLimit() throws { let fd = try connectClient() defer { close(fd) } @@ -874,6 +950,23 @@ final class SocketIntegrationTests: XCTestCase { // MARK: - Helper + private func restartServer(profile: String) { + server.stop() + setenv("BRAINLAYER_MCP_PROFILE", profile, 1) + server = BrainBarServer(socketPath: testSocketPath, dbPath: tempDBPath, database: db) + server.start() + XCTAssertTrue(waitForSocket(at: testSocketPath), "Server should bind \(testSocketPath)") + } + + private func listedToolNames(on fd: Int32, id: Int) throws -> [String] { + try sendMCPRequest(on: fd, request: [ + "jsonrpc": "2.0", "id": id, "method": "tools/list", + ]) + let response = try readMCPMessage(fd: fd) + let tools = (response["result"] as? [String: Any])?["tools"] as? [[String: Any]] + return tools?.compactMap { $0["name"] as? String } ?? [] + } + private func waitForSocket(at path: String, timeout: TimeInterval = 3.0) -> Bool { let deadline = Date().addingTimeInterval(timeout) while Date() < deadline { diff --git a/contracts/engine-ui-contract.md b/contracts/engine-ui-contract.md index 932f3062..e3a8d4c0 100644 --- a/contracts/engine-ui-contract.md +++ b/contracts/engine-ui-contract.md @@ -17,19 +17,31 @@ Pinned shared surfaces: ## 2. BrainBar MCP Router Over `/tmp/brainbar.sock` -BrainBar is the MCP surface of record. Its router dispatches 17 tools at `brain-bar/Sources/BrainBar/MCPRouter.swift:195` and defines the tool schemas at `brain-bar/Sources/BrainBar/MCPRouter.swift:950`. +BrainBar is the MCP surface of record. Its router keeps 17 canonical definitions and dispatch handlers in `brain-bar/Sources/BrainBar/MCPRouter.swift`. Canonical BrainBar tools: `brain_search`, `brain_store`, `brain_get_person`, `brain_recall`, `brain_entity`, `brain_digest`, `brain_update`, `brain_expand`, `brain_tags`, `brain_supersede`, `brain_archive`, `brain_enrich`, `brain_subscribe`, `brain_unsubscribe`, `brain_ack`, `brain_maintenance_rebuild_trigram`, `brain_backup_vacuum_into`. -Python `src/brainlayer/mcp/` remains the secondary transport with 13 tools defined at `src/brainlayer/mcp/__init__.py:387`, `src/brainlayer/mcp/__init__.py:543`, `src/brainlayer/mcp/__init__.py:568`, `src/brainlayer/mcp/__init__.py:672`, `src/brainlayer/mcp/__init__.py:701`, `src/brainlayer/mcp/__init__.py:862`, `src/brainlayer/mcp/__init__.py:905`, `src/brainlayer/mcp/__init__.py:974`, `src/brainlayer/mcp/__init__.py:999`, `src/brainlayer/mcp/__init__.py:1042`, `src/brainlayer/mcp/__init__.py:1080`, `src/brainlayer/mcp/__init__.py:1113`, and `src/brainlayer/mcp/__init__.py:1135`. +Python `src/brainlayer/mcp/` remains the secondary transport with 13 canonical definitions built by `_full_tool_definitions()` in `src/brainlayer/mcp/__init__.py`. Differences: - BrainBar-only: `brain_subscribe`, `brain_unsubscribe`, `brain_ack`, `brain_maintenance_rebuild_trigram`, `brain_backup_vacuum_into`. - Python-only: `brain_resume`. +### MCP palette profiles + +Both transports resolve `BRAINLAYER_MCP_PROFILE` once per server session: + +- missing, blank, `core`, or an unknown value exposes `brain_search`, `brain_store`, `brain_recall`, `brain_expand`, and the `expand_palette` control tool; +- `full` and `operator` expose the transport's complete canonical inventory and omit the redundant control; +- unknown values fail closed to core rather than exposing operator tools. + +BrainBar's core declarations are compact projections of the canonical definitions: names and validation-relevant input-schema fields remain, verbose description/annotation metadata is omitted or shortened, and full/operator definitions stay byte-for-byte unchanged. Python core mode selects the existing `Tool` objects directly. Neither transport deletes or renames a canonical handler. + +Calling `expand_palette` in a core session exposes the complete inventory for subsequent `tools/list` and `tools/call` requests. The first receipt contains `expanded: true`, `already_expanded: false`, and the newly registered names; later calls are successful no-ops with `already_expanded: true`. BrainBar advertises `tools.listChanged = true`; clients refresh on the next `tools/list` because this router seam does not push a notification. + ## 3. Hybrid Helper Subprocess And Socket BrainBar starts the helper from `HybridSearchHelperClient` with `-m brainlayer.brainbar_hybrid_helper --socket-path ... --db-path ...`; see the invocation at `brain-bar/Sources/BrainBar/HybridSearchHelperClient.swift:189`. If `pythonExecutable` resolves to `/usr/bin/env`, BrainBar invokes `python3`; otherwise it invokes the resolved Python executable directly. diff --git a/docs/plans/2026-07-13-brainlayer-core4-palette-design.md b/docs/plans/2026-07-13-brainlayer-core4-palette-design.md new file mode 100644 index 00000000..8979a518 --- /dev/null +++ b/docs/plans/2026-07-13-brainlayer-core4-palette-design.md @@ -0,0 +1,91 @@ +# BrainLayer Core-4 Palette Design + +## Goal + +Reduce the default BrainLayer MCP boot surface to the four tools that cover 98.62% of observed calls while preserving every existing handler and allowing a session to expose the full palette without reconnecting. + +## Context + +The live BrainLayer surface is BrainBar's 17-tool MCP router. The Python MCP server is a secondary transport with 13 tools, including the Python-only `brain_resume`. The measured live schema is 4,038 bytes / 1,027 `o200k_base` tokens; the four canonical schemas (`brain_search`, `brain_store`, `brain_recall`, `brain_expand`) are 1,263 bytes / 326 tokens before adding the small expansion control. + +Both transports therefore need the same profile semantics, but each keeps its existing canonical inventory and dispatch implementation. + +## Considered approaches + +### 1. Stateful per-session palette with one-shot expansion (chosen) + +Each server process resolves `BRAINLAYER_MCP_PROFILE` once at construction/import time. `core` lists the Core 4 plus `expand_palette`; `full` and `operator` list the transport's complete canonical inventory. BrainBar tracks expansion per socket client, while the Python stdio server has one module session. Calling `expand_palette` changes only that MCP session to full and makes its subsequent `tools/list` calls return the full inventory. + +This directly mirrors cmuxlayer's proven seam, meets the boot budget, and supports a mid-session upgrade. The trade-off is a small amount of instance state and a `tools/list_changed` capability declaration even though BrainBar does not currently push that notification. + +### 2. Reconnect-only profile switching + +Keep `tools/list` stateless and require a reconnect with `BRAINLAYER_MCP_PROFILE=full`. This is simpler but fails the explicit self-upgrade requirement. + +### 3. Advertise all tools with deferred metadata + +Return all schemas and mark non-core tools with `defer_loading`. This preserves direct calls but does not reduce the measured `tools/list` bytes, so it fails the primary acceptance gate. + +## Profile contract + +- `core` is the default when the environment variable is missing or blank. +- `full` exposes the complete transport inventory. +- `operator` is a documented alias of `full` for operator sessions. +- Unknown profile values fail closed to `core` and log one warning; they never silently expose operator tools. +- Profile resolution is per server session. Changing the process environment after construction does not mutate an existing session. + +The resident core is exactly: + +1. `brain_search` +2. `brain_store` +3. `brain_recall` +4. `brain_expand` + +`expand_palette` is a control tool and is not counted as a BrainLayer domain tool. It appears only while the server is in an unexpanded core profile. + +## BrainBar design + +`MCPRouter` receives an optional profile override for deterministic tests and otherwise reads `BRAINLAYER_MCP_PROFILE` during initialization. It stores a boolean expanded state behind the router instance. + +`tools/list` derives compact core declarations from the immutable canonical `toolDefinitions` array and appends a minimal `expand_palette` definition. The projection keeps names and validation-relevant input-schema fields while removing verbose nested descriptions/annotations and shortening top-level descriptions. This was required because the current in-repo Core 4 serialized to 3,722 bytes even though the older live-front audit payload was 1,263 bytes. Full/operator selection returns all 17 canonical definitions byte-for-byte and omits the redundant control. + +`tools/call` validates against the currently exposed definitions. A deferred tool called before expansion receives the existing unknown-tool JSON-RPC error, including BrainBar's server-handled subscription tools. `expand_palette` is handled without touching the database, flips the client session once, and returns an idempotent structured receipt containing `expanded`, `already_expanded`, and the newly exposed names. Subsequent calls are successful no-ops. BrainBar keeps the expansion state in each `ClientState`; direct router users receive one router-local default session. + +`initialize` advertises `tools.listChanged = true` because the list can change during a session. BrainBar's request/response router has no server-push channel at this seam, so clients discover the new list on the next `tools/list`; the expansion response makes that contract explicit. + +## Python design + +The Python MCP module keeps one immutable full-tool builder and adds a small palette controller resolved from `BRAINLAYER_MCP_PROFILE`. The registered `list_tools` callback returns Core 4 plus `expand_palette` by default, or the existing 13-tool inventory for full/operator/expanded sessions. + +The existing `call_tool` branches remain intact. A profile exposure guard rejects deferred names before dispatch, while `expand_palette` updates only the module server session and returns the same idempotent receipt shape as BrainBar. No aliases are renamed or removed; Python-only `brain_resume` remains available in full/operator mode. + +## Compatibility and safety + +- No database, schema, migration, search, or write-path code changes. +- No canonical tool definition or handler is deleted. +- Core schemas are mechanically projected from the existing definitions, keeping validation-relevant fields tied to the canonical source while dropping only boot-cost metadata. +- Full/operator output is contract-compatible with today's complete inventory. +- `brain_expand` remains resident despite its Python deprecation notice because the live corpus and handoff explicitly require Core 4 for this release. +- Unknown profiles fail closed to core. + +## Verification + +Tests will prove, on both transports: + +- missing/blank/`core` profiles list exactly Core 4 plus the control tool; +- `full` and `operator` list the complete pre-change inventory; +- invalid profiles fail closed; +- deferred calls fail before expansion and dispatch after expansion; +- expansion is idempotent and round-trips from core to full; +- canonical definitions remain present and unchanged; +- the serialized default BrainBar `tools` array stays below 1,500 bytes; +- measured full and core bytes/tokens are reported with the same compact JSON and `o200k_base` method used by the audit. + +Focused Python and Swift tests run first, followed by the full non-live Python suite and the full BrainBar Swift suite. The final receipt records measured before/after bytes and token counts. + +## Non-goals + +- No Core-3 migration or folding `brain_expand` into `brain_search`. +- No launcher, daemon, database, schema, or writer changes. +- No deletion or renaming of compatibility tools. +- No attempt to make profile changes affect other sessions in the same process. diff --git a/docs/plans/2026-07-13-brainlayer-core4-palette.md b/docs/plans/2026-07-13-brainlayer-core4-palette.md new file mode 100644 index 00000000..6be6494a --- /dev/null +++ b/docs/plans/2026-07-13-brainlayer-core4-palette.md @@ -0,0 +1,304 @@ +# BrainLayer Core-4 Palette Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Make Core 4 the default BrainLayer MCP palette on BrainBar and Python, with full/operator overrides and an idempotent mid-session `expand_palette` upgrade. + +**Architecture:** Keep each transport's canonical tool definitions and dispatch handlers unchanged. Add a small stateful palette controller at each `tools/list`/`tools/call` seam; core sessions expose four domain tools plus a minimal expansion control, while full/operator sessions expose the pre-change complete inventory. + +**Tech Stack:** Swift 6/XCTest, Python 3.13/pytest, MCP JSON-RPC, compact JSON measurement, `tiktoken` `o200k_base`. + +--- + +### Task 1: Pin BrainBar profile and budget behavior + +**Files:** +- Modify: `brain-bar/Tests/BrainBarTests/MCPRouterTests.swift` +- Modify: `brain-bar/Sources/BrainBar/MCPRouter.swift` + +**Step 1: Write the failing profile tests** + +Add tests that construct `MCPRouter(profile: ...)`, call `tools/list`, and assert: + +```swift +let coreNames = ["brain_search", "brain_store", "brain_recall", "brain_expand", "expand_palette"] +XCTAssertEqual(toolNames(router: MCPRouter(profile: nil)), coreNames) +XCTAssertEqual(toolNames(router: MCPRouter(profile: "core")), coreNames) + +for profile in ["full", "operator"] { + XCTAssertEqual(toolNames(router: MCPRouter(profile: profile)).count, 17) +} +``` + +Also cover blank and invalid profiles, asserting both fail closed to the same core list. Move existing canonical schema/annotation assertions to `MCPRouter.toolDefinitions` or a full-profile router so they continue to pin all 17 definitions. + +Serialize the core `tools` array with sorted compact JSON and assert `data.count <= 1_500`. + +**Step 2: Run the focused Swift tests and verify RED** + +Run: + +```bash +cd brain-bar +swift test --filter MCPRouterTests +``` + +Expected: compilation/test failures because `MCPRouter(profile:)` and palette filtering do not exist. + +**Step 3: Implement the minimal BrainBar profile controller** + +In `MCPRouter.swift`: + +- add `BRAINLAYER_MCP_PROFILE` resolution in `init`, with an optional explicit profile for tests; +- represent `core`, `full`, and `operator` without modifying `toolDefinitions`; +- add a lock-safe `PaletteSession`, use one router-local default for direct callers, and attach a distinct session to each BrainBar socket client; +- compute compact core projections from the canonical Core 4 (preserving validation fields while dropping verbose descriptive metadata) plus a minimal control definition, or all 17 canonical definitions unchanged; +- make `handleToolsList` return `exposedToolDefinitions`; +- make the existing call guard validate against exposed definitions; +- set `tools.listChanged` to true. + +Unknown and blank values resolve to core. Log a warning only for a nonblank unknown environment value. + +**Step 4: Run focused tests and verify GREEN** + +Run `cd brain-bar && swift test --filter MCPRouterTests`. + +Expected: all `MCPRouterTests` pass with zero failures. + +**Step 5: Commit the BrainBar profile seam** + +```bash +git add brain-bar/Sources/BrainBar/MCPRouter.swift brain-bar/Tests/BrainBarTests/MCPRouterTests.swift +git commit -m "feat: add BrainBar core4 MCP profile" +``` + +### Task 2: Add BrainBar expansion round-trip + +**Files:** +- Modify: `brain-bar/Tests/BrainBarTests/MCPRouterTests.swift` +- Modify: `brain-bar/Sources/BrainBar/MCPRouter.swift` + +**Step 1: Write failing expansion tests** + +Exercise one core router instance: + +```swift +XCTAssertEqual(listedNames(router).count, 5) +let first = call(router, name: "expand_palette") +XCTAssertEqual(first["expanded"] as? Bool, true) +XCTAssertEqual(listedNames(router).count, 17) +let second = call(router, name: "expand_palette") +XCTAssertEqual(second["already_expanded"] as? Bool, true) +XCTAssertEqual(listedNames(router).count, 17) +``` + +Also prove a deferred no-database tool such as `brain_tags` receives `-32601` before expansion and reaches its existing handler after expansion (returning a normal tool result, even if that result reports missing database). + +**Step 2: Run focused tests and verify RED** + +Run `cd brain-bar && swift test --filter MCPRouterTests`. + +Expected: failure because the expansion control is listed but not dispatched and does not mutate the palette. + +**Step 3: Implement idempotent expansion** + +Handle `expand_palette` before canonical schema validation. On the first call, set `paletteExpanded = true` and return content plus top-level metadata: + +```swift +[ + "expanded": true, + "already_expanded": false, + "registered_tools": deferredNames, +] +``` + +Later calls return `expanded: false`, `already_expanded: true`, and an empty registered list. The canonical switch and all 17 handlers remain unchanged. + +**Step 4: Run focused tests and verify GREEN** + +Run `cd brain-bar && swift test --filter MCPRouterTests`. + +Expected: all focused tests pass. + +**Step 5: Commit the expansion path** + +```bash +git add brain-bar/Sources/BrainBar/MCPRouter.swift brain-bar/Tests/BrainBarTests/MCPRouterTests.swift +git commit -m "feat: expand BrainBar palette in session" +``` + +### Task 3: Pin the Python palette contract + +**Files:** +- Create: `tests/test_mcp_palette.py` +- Create: `src/brainlayer/mcp/palette.py` +- Modify: `src/brainlayer/mcp/__init__.py` + +**Step 1: Write failing controller and transport tests** + +Test a new `ToolPalette` directly and the MCP callbacks through an injected/reset palette fixture: + +```python +CORE = {"brain_search", "brain_store", "brain_recall", "brain_expand", "expand_palette"} + +@pytest.mark.parametrize("profile", [None, "", "core", "bogus"]) +def test_core_profiles_fail_closed(profile): + palette = ToolPalette(profile) + assert {tool.name for tool in palette.expose(full_tools())} == CORE + +@pytest.mark.parametrize("profile", ["full", "operator"]) +def test_full_profiles_preserve_all_python_tools(profile): + assert len(ToolPalette(profile).expose(full_tools())) == 13 +``` + +Add an async round-trip test proving the first expansion exposes all 13 tools, the second expansion is a no-op, and a deferred call is rejected before expansion. + +**Step 2: Run the focused Python tests and verify RED** + +Run: + +```bash +pytest tests/test_mcp_palette.py -q +``` + +Expected: import/attribute failures because the palette controller does not exist. + +**Step 3: Implement the Python palette controller** + +Create `palette.py` with: + +- `PROFILE_ENV = "BRAINLAYER_MCP_PROFILE"`; +- immutable Core 4 names; +- normalized `core`/`full`/`operator` resolution with invalid fail-closed behavior; +- `expose(full_tools)` that selects existing `Tool` objects and appends a minimal control in unexpanded core mode; +- `is_exposed(name)` and idempotent `expand()` returning a receipt dictionary. + +Refactor the current literal `list_tools` body into `_full_tool_definitions()`. Keep the decorated callback thin: + +```python +@server.list_tools() +async def list_tools() -> list[Tool]: + return _tool_palette.expose(_full_tool_definitions()) +``` + +At the start of `call_tool`, handle `expand_palette`, then reject other currently deferred names before entering the existing branch chain. Do not modify any existing handler body. + +**Step 4: Run focused tests and verify GREEN** + +Run `pytest tests/test_mcp_palette.py -q`. + +Expected: all palette tests pass. + +**Step 5: Run nearby Python contract tests** + +Run: + +```bash +pytest tests/test_mcp_input_schema_limits.py tests/test_mcp_labeled_field_output.py tests/test_mcp_digest_modes.py tests/test_smart_search_entity_dedup.py -q +``` + +Expected: zero failures. Update tests that intentionally inspect the complete inventory to use the full profile fixture; do not weaken canonical schema assertions. + +**Step 6: Commit the Python seam** + +```bash +git add src/brainlayer/mcp/__init__.py src/brainlayer/mcp/palette.py tests/test_mcp_palette.py tests/test_mcp_input_schema_limits.py tests/test_mcp_labeled_field_output.py tests/test_mcp_digest_modes.py tests/test_smart_search_entity_dedup.py +git commit -m "feat: add Python core4 MCP profile" +``` + +### Task 4: Document parity and measure the palette + +**Files:** +- Modify: `contracts/engine-ui-contract.md` +- Create: `scripts/measure_mcp_palette.py` +- Create: `tests/test_measure_mcp_palette.py` + +**Step 1: Write a failing measurement test** + +Add a pure helper test with a small fixed schema fixture. Assert compact JSON uses UTF-8 byte length and `o200k_base` tokens, and assert the live BrainBar core payload is at most 1,500 bytes. + +**Step 2: Run the measurement tests and verify RED** + +Run `pytest tests/test_measure_mcp_palette.py -q`. + +Expected: failure because the script/helper does not exist. + +**Step 3: Implement the measurement script** + +The script must accept JSON from a file/stdin, extract the `tools` array when given a JSON-RPC envelope, serialize with compact separators and sorted keys, and print exact bytes/tokens. Use `tiktoken.get_encoding("o200k_base")` and exit nonzero when `--max-bytes` is exceeded. + +**Step 4: Update the transport contract** + +Document: + +- `BRAINLAYER_MCP_PROFILE=core|full|operator`; +- core default and its four domain names; +- `expand_palette` semantics; +- BrainBar's 17-tool and Python's 13-tool full inventories remain canonical; +- profile state is per session and invalid values fail closed. + +**Step 5: Run tests and record measurements** + +Run: + +```bash +pytest tests/test_measure_mcp_palette.py tests/test_mcp_palette.py -q +cd brain-bar && swift test --filter MCPRouterTests +``` + +Generate compact full/core `tools` JSON for BrainBar and run the measurement script on both. Record exact before/after bytes and tokens for the receipt. + +**Step 6: Commit docs and measurement tooling** + +```bash +git add contracts/engine-ui-contract.md scripts/measure_mcp_palette.py tests/test_measure_mcp_palette.py +git commit -m "test: enforce BrainLayer palette budget" +``` + +### Task 5: Full verification and PR loop + +**Files:** +- Modify: `~/Gits/orchestrator/docs.local/collab/driver-buddy-2026-07-12.md` + +**Step 1: Run complete local verification** + +Run fresh: + +```bash +pytest tests/ -m "not integration and not live" --tb=short +cd brain-bar && swift test +git diff --check origin/main...HEAD +``` + +Read the complete summaries and report exact pass/fail/skip counts. Do not claim green from exit codes alone. + +**Step 2: Re-read the handoff and design acceptance gates** + +Verify line by line: Core 4 default, 13 deferred on BrainBar, full/operator override, mid-session expansion, no removal/rename, no DB/schema change, per-profile tests, <=1.5 KB core budget, exact measurements. + +**Step 3: Push and open a ready PR** + +```bash +git push -u origin feat/core4-palette +gh pr create --title "feat: add BrainLayer Core-4 MCP palette" --body-file +``` + +The PR body includes the design, exact before/after bytes/tokens, test receipts, and risk note that only MCP exposure changes. + +**Step 4: Request required reviews** + +Post `@codex review`, `@cursor`, and `@bugbot` on the ready PR. Inspect CI with `gh pr checks --watch` and inspect all review threads. Fix or explicitly answer every actionable/high/critical item, with no more than three review rounds. + +**Step 5: Write completion receipts** + +Append to `~/Gits/orchestrator/docs.local/collab/driver-buddy-2026-07-12.md` under `@orc-driver-v2` with: + +- PR URL and head SHA; +- exact full/core bytes and tokens; +- local and CI test counts; +- review status; +- explicit note that DB/schema were untouched. + +Store the same milestone in BrainLayer with WHAT + WHY, then search it back to verify persistence. Stop at the handoff's worker endpoint (ready PR + reviews addressed); do not merge unless separately instructed. + +TASK_DONE diff --git a/pyproject.toml b/pyproject.toml index aecda2d7..c14fa9c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,7 @@ dev = [ "ruff>=0.1.0", "scikit-learn<1.6", "setuptools>=70.0.0,<81", + "tiktoken>=0.7.0", "umap-learn>=0.5.0", ] docs = [ diff --git a/scripts/measure_mcp_palette.py b/scripts/measure_mcp_palette.py new file mode 100644 index 00000000..ae3471bf --- /dev/null +++ b/scripts/measure_mcp_palette.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Measure a compact MCP tools payload in UTF-8 bytes and o200k tokens.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import tiktoken + + +def extract_tools(payload: Any) -> list[dict[str, Any]]: + """Extract a tools array from a raw list, result object, or JSON-RPC envelope.""" + tools: Any + if isinstance(payload, list): + tools = payload + elif isinstance(payload, dict) and isinstance(payload.get("tools"), list): + tools = payload["tools"] + elif ( + isinstance(payload, dict) + and isinstance(payload.get("result"), dict) + and isinstance(payload["result"].get("tools"), list) + ): + tools = payload["result"]["tools"] + else: + raise ValueError("Expected a tools array or a JSON-RPC result containing one") + + if not all(isinstance(tool, dict) for tool in tools): + raise ValueError("Every tools entry must be a JSON object") + return tools + + +def canonical_tools_json(tools: list[dict[str, Any]]) -> str: + return json.dumps(tools, ensure_ascii=False, separators=(",", ":"), sort_keys=True) + + +def measure_tools(tools: list[dict[str, Any]]) -> dict[str, int]: + canonical = canonical_tools_json(tools) + return { + "tools": len(tools), + "bytes": len(canonical.encode("utf-8")), + "tokens": len(tiktoken.get_encoding("o200k_base").encode(canonical)), + } + + +def _load_payload(path: str) -> Any: + if path == "-": + return json.load(sys.stdin) + with Path(path).open(encoding="utf-8") as handle: + return json.load(handle) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", nargs="?", default="-", help="JSON file path, or - for stdin") + parser.add_argument("--max-bytes", type=int, help="Exit nonzero if the compact tools array exceeds this size") + args = parser.parse_args(argv) + + try: + measurement = measure_tools(extract_tools(_load_payload(args.path))) + except (OSError, json.JSONDecodeError, ValueError) as exc: + parser.error(str(exc)) + + print(json.dumps(measurement, sort_keys=True)) + if args.max_bytes is not None and measurement["bytes"] > args.max_bytes: + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/brainlayer/mcp/__init__.py b/src/brainlayer/mcp/__init__.py index 84c72aa8..342e6b4d 100644 --- a/src/brainlayer/mcp/__init__.py +++ b/src/brainlayer/mcp/__init__.py @@ -1,6 +1,7 @@ """BrainLayer MCP Server - Model Context Protocol interface for Claude Code.""" import asyncio +import json import logging import os from copy import deepcopy @@ -40,6 +41,7 @@ from .enrich_handler import _brain_enrich from .entity_handler import _brain_entity as _brain_entity from .entity_handler import _brain_get_person +from .palette import EXPAND_TOOL_NAME, ToolPalette from .search_handler import ( _brain_recall, _brain_resume, @@ -396,9 +398,8 @@ def _apply_input_limits(node: Any, field_name: str | None = None) -> None: # --- Tool registration --- -@server.list_tools() -async def list_tools() -> list[Tool]: - """List available tools — 3 primary + backward-compat aliases.""" +def _full_tool_definitions() -> list[Tool]: + """Build the complete canonical Python MCP tool inventory.""" return [ Tool( name="brain_search", @@ -1251,6 +1252,16 @@ async def list_tools() -> list[Tool]: ] +_tool_palette = ToolPalette() +_FULL_TOOL_NAMES = tuple(tool.name for tool in _full_tool_definitions()) + + +@server.list_tools() +async def list_tools() -> list[Tool]: + """List tools exposed by this server session's palette.""" + return _tool_palette.expose(_full_tool_definitions()) + + # --- Completions --- @@ -1319,6 +1330,18 @@ async def handle_completion(ref, argument) -> CompleteResult: async def call_tool(name: str, arguments: dict[str, Any]): """Handle tool calls — 3 primary tools + backward-compat aliases.""" + if name == EXPAND_TOOL_NAME: + if _tool_palette.is_full: + return _error_result(f"Unknown tool: {name}") + receipt = _tool_palette.expand(_FULL_TOOL_NAMES) + return CallToolResult( + content=[TextContent(type="text", text=json.dumps(receipt, sort_keys=True))], + structuredContent=receipt, + ) + + if name in _FULL_TOOL_NAMES and not _tool_palette.is_exposed(name): + return _error_result(f"Tool not exposed in the core profile: {name}. Call {EXPAND_TOOL_NAME} first.") + # --- New consolidated tools --- if name == "brain_search": diff --git a/src/brainlayer/mcp/palette.py b/src/brainlayer/mcp/palette.py new file mode 100644 index 00000000..88cc5f55 --- /dev/null +++ b/src/brainlayer/mcp/palette.py @@ -0,0 +1,71 @@ +"""Per-session MCP tool palette selection.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Sequence + +from mcp.types import Tool + +PROFILE_ENV = "BRAINLAYER_MCP_PROFILE" +CORE_TOOL_NAMES = ( + "brain_search", + "brain_store", + "brain_recall", + "brain_expand", +) +EXPAND_TOOL_NAME = "expand_palette" + +logger = logging.getLogger(__name__) + + +class ToolPalette: + """Resolve one MCP session's visible tools and expansion state.""" + + def __init__(self, profile: str | None = None) -> None: + environment_profile = os.environ.get(PROFILE_ENV) + raw_profile = environment_profile if profile is None else profile + normalized = (raw_profile or "").strip().lower() + + self._full = normalized in {"full", "operator"} + self._expanded = False + if normalized not in {"", "core", "full", "operator"} and profile is None: + logger.warning("Unknown %s=%r; using core profile", PROFILE_ENV, environment_profile) + + @property + def is_full(self) -> bool: + return self._full + + def expose(self, full_tools: Sequence[Tool]) -> list[Tool]: + if self._full or self._expanded: + return list(full_tools) + + tools_by_name = {tool.name: tool for tool in full_tools} + core_tools = [tools_by_name[name] for name in CORE_TOOL_NAMES if name in tools_by_name] + core_tools.append( + Tool( + name=EXPAND_TOOL_NAME, + description="Expose all tools.", + inputSchema={"type": "object"}, + ) + ) + return core_tools + + def is_exposed(self, name: str) -> bool: + return self._full or self._expanded or name in CORE_TOOL_NAMES or name == EXPAND_TOOL_NAME + + def expand(self, full_tool_names: Sequence[str]) -> dict[str, object]: + if self._expanded: + return { + "expanded": False, + "already_expanded": True, + "registered_tools": [], + } + + self._expanded = True + return { + "expanded": True, + "already_expanded": False, + "registered_tools": [name for name in full_tool_names if name not in CORE_TOOL_NAMES], + } diff --git a/tests/test_3tool_aliases.py b/tests/test_3tool_aliases.py index ccc4a2d1..c6788065 100644 --- a/tests/test_3tool_aliases.py +++ b/tests/test_3tool_aliases.py @@ -14,8 +14,20 @@ import asyncio from unittest.mock import AsyncMock, MagicMock, patch +import pytest + from brainlayer.mcp.search_handler import _brain_recall, _smart_detect_mode + +@pytest.fixture(autouse=True) +def _use_full_palette_for_legacy_contracts(monkeypatch): + """This module verifies handlers retained behind the explicit full profile.""" + import brainlayer.mcp as mcp_server + from brainlayer.mcp.palette import ToolPalette + + monkeypatch.setattr(mcp_server, "_tool_palette", ToolPalette("full")) + + # ── Smart Mode Detection ───────────────────────────────────────────────────── diff --git a/tests/test_audit_search_quality.py b/tests/test_audit_search_quality.py index 227da0eb..28157ce6 100644 --- a/tests/test_audit_search_quality.py +++ b/tests/test_audit_search_quality.py @@ -122,9 +122,9 @@ async def test_mcp_recall_enum_matches_canonical(self): async def test_mcp_entity_enum_matches_canonical(self): """brain_entity entity_type enum must match the MCP schema contract.""" - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = await list_tools() + tools = _full_tool_definitions() entity_tool = next(t for t in tools if t.name == "brain_entity") schema = entity_tool.inputSchema mcp_enum = schema["properties"]["entity_type"]["enum"] diff --git a/tests/test_entity_type_sync.py b/tests/test_entity_type_sync.py index e0ad4d36..6db6c289 100644 --- a/tests/test_entity_type_sync.py +++ b/tests/test_entity_type_sync.py @@ -28,12 +28,10 @@ def store(tmp_path): def _get_brain_entity_tool(): - """Extract the brain_entity Tool object from the MCP server tool list.""" - import asyncio + """Extract brain_entity from the explicit full-profile tool definitions.""" + from brainlayer.mcp import _full_tool_definitions - from brainlayer.mcp import list_tools - - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() for tool in tools: if tool.name == "brain_entity": return tool diff --git a/tests/test_mcp_digest_modes.py b/tests/test_mcp_digest_modes.py index 11cbf3b4..63032dba 100644 --- a/tests/test_mcp_digest_modes.py +++ b/tests/test_mcp_digest_modes.py @@ -1,6 +1,5 @@ """Tests for brain_digest MCP mode routing.""" -import asyncio from types import SimpleNamespace from unittest.mock import MagicMock @@ -80,9 +79,9 @@ async def test_brain_digest_missing_content_with_mode_digest_errors(): def test_brain_digest_input_schema_includes_mode_and_limit(): - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() digest = next(t for t in tools if t.name == "brain_digest") props = digest.inputSchema["properties"] diff --git a/tests/test_mcp_input_schema_limits.py b/tests/test_mcp_input_schema_limits.py index f554b128..967cd36c 100644 --- a/tests/test_mcp_input_schema_limits.py +++ b/tests/test_mcp_input_schema_limits.py @@ -5,11 +5,13 @@ from mcp import types -from brainlayer.mcp import list_tools, server +import brainlayer.mcp as mcp_module +from brainlayer.mcp import _full_tool_definitions, server +from brainlayer.mcp.palette import ToolPalette def _get_tools(): - return asyncio.run(list_tools()) + return _full_tool_definitions() def _iter_string_fields(schema: dict[str, Any], path: str = ""): @@ -61,7 +63,8 @@ async def _call_brain_digest(arguments: dict[str, Any]): return await handler(request) -def test_brain_digest_schema_rejects_oversized_content(): +def test_brain_digest_schema_rejects_oversized_content(monkeypatch): + monkeypatch.setattr(mcp_module, "_tool_palette", ToolPalette("full")) result = asyncio.run(_call_brain_digest({"content": "x" * 200_001})).root assert result.isError is True diff --git a/tests/test_mcp_palette.py b/tests/test_mcp_palette.py new file mode 100644 index 00000000..543af736 --- /dev/null +++ b/tests/test_mcp_palette.py @@ -0,0 +1,66 @@ +"""Tests for the stateful BrainLayer MCP tool palette.""" + +import asyncio + +import pytest + +import brainlayer.mcp as mcp_module +from brainlayer.mcp import _full_tool_definitions, call_tool, list_tools +from brainlayer.mcp.palette import CORE_TOOL_NAMES, ToolPalette + +CORE_WITH_CONTROL = (*CORE_TOOL_NAMES, "expand_palette") + + +@pytest.mark.parametrize("profile", [None, "", " ", "core", "bogus"]) +def test_core_profiles_fail_closed(monkeypatch, profile): + monkeypatch.delenv("BRAINLAYER_MCP_PROFILE", raising=False) + palette = ToolPalette(profile) + + assert tuple(tool.name for tool in palette.expose(_full_tool_definitions())) == CORE_WITH_CONTROL + + +@pytest.mark.parametrize("profile", ["full", "operator"]) +def test_full_profiles_preserve_all_python_tools(profile): + tools = ToolPalette(profile).expose(_full_tool_definitions()) + + assert len(tools) == 13 + assert tuple(tool.name for tool in tools) == tuple(tool.name for tool in _full_tool_definitions()) + + +def test_environment_profile_is_resolved_once(monkeypatch): + monkeypatch.setenv("BRAINLAYER_MCP_PROFILE", "operator") + palette = ToolPalette() + monkeypatch.setenv("BRAINLAYER_MCP_PROFILE", "core") + + assert len(palette.expose(_full_tool_definitions())) == 13 + + +def test_python_palette_expands_once_and_dispatches_deferred_tools(monkeypatch): + palette = ToolPalette("core") + monkeypatch.setattr(mcp_module, "_tool_palette", palette) + + assert tuple(tool.name for tool in asyncio.run(list_tools())) == CORE_WITH_CONTROL + + before = asyncio.run(call_tool("brain_tags", {})) + assert before.isError is True + assert "not exposed" in before.content[0].text + + first = asyncio.run(call_tool("expand_palette", {})) + assert first.isError is False + assert first.structuredContent == { + "expanded": True, + "already_expanded": False, + "registered_tools": [tool.name for tool in _full_tool_definitions() if tool.name not in CORE_TOOL_NAMES], + } + assert len(asyncio.run(list_tools())) == 13 + + after = asyncio.run(call_tool("brain_tags", {})) + assert "deprecated" in after.content[0].text + + second = asyncio.run(call_tool("expand_palette", {})) + assert second.structuredContent == { + "expanded": False, + "already_expanded": True, + "registered_tools": [], + } + assert len(asyncio.run(list_tools())) == 13 diff --git a/tests/test_measure_mcp_palette.py b/tests/test_measure_mcp_palette.py new file mode 100644 index 00000000..3fbe4a70 --- /dev/null +++ b/tests/test_measure_mcp_palette.py @@ -0,0 +1,42 @@ +"""Tests for deterministic MCP palette measurement.""" + +import json + +import tiktoken + +from scripts.measure_mcp_palette import extract_tools, main, measure_tools + + +def test_measure_tools_uses_compact_utf8_json_and_o200k_tokens(): + tools = [ + { + "name": "memory_זיכרון", + "inputSchema": {"properties": {}, "type": "object"}, + } + ] + canonical = json.dumps(tools, ensure_ascii=False, separators=(",", ":"), sort_keys=True) + + measurement = measure_tools(tools) + + assert measurement == { + "tools": 1, + "bytes": len(canonical.encode("utf-8")), + "tokens": len(tiktoken.get_encoding("o200k_base").encode(canonical)), + } + + +def test_extract_tools_accepts_jsonrpc_envelope(): + tools = [{"name": "brain_search", "inputSchema": {"type": "object"}}] + + assert extract_tools({"jsonrpc": "2.0", "id": 1, "result": {"tools": tools}}) == tools + + +def test_main_enforces_max_bytes(tmp_path, capsys): + payload_path = tmp_path / "tools.json" + payload_path.write_text(json.dumps({"tools": [{"name": "brain_search"}]}), encoding="utf-8") + + assert main([str(payload_path), "--max-bytes", "1"]) == 1 + output = json.loads(capsys.readouterr().out) + assert output["bytes"] > 1 + + assert main([str(payload_path), "--max-bytes", str(output["bytes"])]) == 0 diff --git a/tests/test_phase3_digest.py b/tests/test_phase3_digest.py index b7f36c90..1924bcdd 100644 --- a/tests/test_phase3_digest.py +++ b/tests/test_phase3_digest.py @@ -287,23 +287,20 @@ def fake_faceted_enrich(*, content, project, title, participants): # noqa: ARG0 def test_brain_digest_tool_exists(): - """brain_digest tool is registered in MCP server.""" - import asyncio + """brain_digest tool is registered in the full MCP profile.""" - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() tool_names = [t.name for t in tools] assert "brain_digest" in tool_names def test_brain_digest_schema_has_required_fields(): """brain_digest tool exposes digest fields and mode-based enrich controls.""" - import asyncio + from brainlayer.mcp import _full_tool_definitions - from brainlayer.mcp import list_tools - - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() digest = next(t for t in tools if t.name == "brain_digest") props = digest.inputSchema.get("properties", {}) assert "content" in props @@ -315,11 +312,9 @@ def test_brain_digest_schema_has_required_fields(): def test_brain_digest_description_teaches_routing(): """brain_digest description explains when to use it and how it differs from brain_store.""" - import asyncio + from brainlayer.mcp import _full_tool_definitions - from brainlayer.mcp import list_tools - - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() digest = next(t for t in tools if t.name == "brain_digest") desc = digest.description.lower() @@ -336,23 +331,20 @@ def test_brain_digest_description_teaches_routing(): def test_brain_entity_tool_exists(): - """brain_entity tool is registered in MCP server.""" - import asyncio + """brain_entity tool is registered in the full MCP profile.""" - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() tool_names = [t.name for t in tools] assert "brain_entity" in tool_names def test_brain_entity_schema(): """brain_entity keeps query in properties even when optional.""" - import asyncio - - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() entity_tool = next(t for t in tools if t.name == "brain_entity") props = entity_tool.inputSchema.get("properties", {}) assert "query" in props diff --git a/tests/test_search_filter_params.py b/tests/test_search_filter_params.py index 74d95c70..721c3647 100644 --- a/tests/test_search_filter_params.py +++ b/tests/test_search_filter_params.py @@ -803,9 +803,9 @@ class TestBrainResumeSchema: """Verify the explicit checkpoint resume tool is exposed.""" def test_brain_resume_tool_schema(self): - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() brain_resume = next(tool for tool in tools if tool.name == "brain_resume") props = brain_resume.inputSchema["properties"] diff --git a/tests/test_think_recall_integration.py b/tests/test_think_recall_integration.py index 038e8785..a02d7897 100644 --- a/tests/test_think_recall_integration.py +++ b/tests/test_think_recall_integration.py @@ -246,12 +246,11 @@ class TestMCPToolCount: """Verify MCP server has correct tool count.""" def test_tool_count(self): - """MCP server should have 13 tools including explicit checkpoint resume.""" - import asyncio + """Full MCP profile has 13 tools including explicit checkpoint resume.""" - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() assert len(tools) == 13 def test_consolidated_tools_registered(self): diff --git a/tests/test_tool_annotations.py b/tests/test_tool_annotations.py index e192fc66..4f0d9960 100644 --- a/tests/test_tool_annotations.py +++ b/tests/test_tool_annotations.py @@ -67,9 +67,9 @@ class TestToolAnnotationsPresent: """Every tool MUST have ToolAnnotations with readOnlyHint, destructiveHint, idempotentHint.""" def _get_tools(self): - from brainlayer.mcp import list_tools + from brainlayer.mcp import _full_tool_definitions - return asyncio.run(list_tools()) + return _full_tool_definitions() def test_all_tools_have_annotations(self): """Every expected tool must have a non-None annotations field.""" diff --git a/tests/test_tool_description_quality.py b/tests/test_tool_description_quality.py index 83cfc21b..8fec35e0 100644 --- a/tests/test_tool_description_quality.py +++ b/tests/test_tool_description_quality.py @@ -1,6 +1,4 @@ -import asyncio - -from brainlayer.mcp import list_tools +from brainlayer.mcp import _full_tool_definitions ACTIVE_TOOL_NAMES = { "brain_search", @@ -13,7 +11,7 @@ def _tool_descriptions() -> dict[str, str]: - tools = asyncio.run(list_tools()) + tools = _full_tool_definitions() return {tool.name: tool.description for tool in tools}