Skip to content

Commit ddc8ee9

Browse files
macOS26macOS26
committed
O(1) lookups: dictionary-backed ServerManager + toolsByID cache
ServerManager: replaced 5 linear scans (contains, firstIndex, removeAll, first) with UUID→index dictionary. All ID-based operations now O(1). MCPClient: added toolsByID dictionary for O(1) tool lookup by UUID in callTool(toolId:). Synced on discovery and disconnect.
1 parent 9337382 commit ddc8ee9

2 files changed

Lines changed: 60 additions & 57 deletions

File tree

‎Sources/AgentMCP/MCPClient.swift‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,8 @@ public actor MCPClient {
173173
private var connections: [UUID: any MCPConnection] = [:]
174174
private var configs: [UUID: ServerConfig] = [:]
175175
private var discoveredTools: [UUID: [DiscoveredTool]] = [:]
176+
/// O(1) tool lookup by tool UUID
177+
private var toolsByID: [UUID: DiscoveredTool] = [:]
176178
private var discoveredResources: [UUID: [DiscoveredResource]] = [:]
177179
private var errors: [UUID: String] = [:]
178180

@@ -245,6 +247,8 @@ public actor MCPClient {
245247
connections[serverId]?.disconnect()
246248
connections.removeValue(forKey: serverId)
247249
configs.removeValue(forKey: serverId)
250+
// Remove tools from O(1) cache before removing from per-server list
251+
for tool in discoveredTools[serverId] ?? [] { toolsByID.removeValue(forKey: tool.id) }
248252
discoveredTools.removeValue(forKey: serverId)
249253
discoveredResources.removeValue(forKey: serverId)
250254
errors.removeValue(forKey: serverId)
@@ -278,6 +282,7 @@ public actor MCPClient {
278282
guard connection.isAlive else {
279283
connections.removeValue(forKey: serverId)
280284
configs.removeValue(forKey: serverId)
285+
for tool in discoveredTools[serverId] ?? [] { toolsByID.removeValue(forKey: tool.id) }
281286
discoveredTools.removeValue(forKey: serverId)
282287
throw MCPClientError.connectionFailed("Server process is no longer running")
283288
}
@@ -334,7 +339,7 @@ public actor MCPClient {
334339
}
335340

336341
public func callTool(toolId: UUID, arguments: [String: JSONValue] = [:]) async throws -> ToolResult {
337-
guard let tool = getAllTools().first(where: { $0.id == toolId }) else {
342+
guard let tool = toolsByID[toolId] else {
338343
throw MCPClientError.toolNotFound(toolId)
339344
}
340345
return try await callTool(serverId: tool.serverId, name: tool.name, arguments: arguments)
@@ -470,6 +475,8 @@ public actor MCPClient {
470475
return DiscoveredTool(serverId: serverId, serverName: serverName, name: name,
471476
description: String(description.prefix(2048)), inputSchemaJSON: schemaJSON)
472477
}
478+
// Sync O(1) lookup cache
479+
for tool in discoveredTools[serverId] ?? [] { toolsByID[tool.id] = tool }
473480
}
474481
} catch { discoveredTools[serverId] = [] }
475482
} else {

‎Sources/AgentMCP/ServerManager.swift‎

Lines changed: 52 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -2,183 +2,179 @@ import Foundation
22

33
/// Manages persistent MCP server configurations
44
/// Stores server configs in ~/Library/Application Support/Agent!/MCPServers/
5+
/// Uses dictionary-backed O(1) lookups for all ID-based operations.
56
public actor ServerManager {
6-
7+
78
// MARK: - Properties
8-
9+
910
public static let shared = ServerManager()
10-
11+
1112
private let configDirectory: URL
1213
private let configFile: URL
14+
/// Ordered list for serialization and display
1315
private var servers: [MCPClient.ServerConfig] = []
14-
16+
/// O(1) lookup by UUID
17+
private var serverIndex: [UUID: Int] = [:]
18+
1519
// MARK: - Initialization
16-
20+
1721
public init() {
1822
let appSupport = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
1923
?? FileManager.default.homeDirectoryForCurrentUser.appendingPathComponent("Library/Application Support")
2024
configDirectory = appSupport.appendingPathComponent("Agent!/MCPServers")
2125
configFile = configDirectory.appendingPathComponent("servers.json")
2226
}
23-
27+
28+
/// Rebuild the UUID → index map from the servers array
29+
private func rebuildIndex() {
30+
serverIndex.removeAll(keepingCapacity: true)
31+
for (i, s) in servers.enumerated() {
32+
serverIndex[s.id] = i
33+
}
34+
}
35+
2436
// MARK: - Server Management
25-
37+
2638
/// Load saved server configurations
2739
public func loadServers() throws -> [MCPClient.ServerConfig] {
28-
// Create directory if needed
2940
try FileManager.default.createDirectory(at: configDirectory, withIntermediateDirectories: true)
30-
31-
// Load from file if exists
41+
3242
guard FileManager.default.fileExists(atPath: configFile.path),
3343
let data = FileManager.default.contents(atPath: configFile.path),
3444
let decoded = try? JSONDecoder().decode([MCPClient.ServerConfig].self, from: data) else {
3545
servers = []
46+
serverIndex = [:]
3647
return servers
3748
}
38-
49+
3950
servers = decoded
51+
rebuildIndex()
4052
return servers
4153
}
42-
54+
4355
/// Save server configurations to disk
4456
public func saveServers() throws {
45-
// Create directory if needed
4657
try FileManager.default.createDirectory(at: configDirectory, withIntermediateDirectories: true)
47-
4858
let data = try JSONEncoder().encode(servers)
4959
try data.write(to: configFile)
5060
}
51-
61+
5262
/// Add a new server configuration
5363
public func addServer(_ config: MCPClient.ServerConfig) throws {
54-
// Check for duplicate
55-
if servers.contains(where: { $0.id == config.id }) {
64+
guard serverIndex[config.id] == nil else {
5665
throw ServerManagerError.duplicateServer
5766
}
58-
67+
serverIndex[config.id] = servers.count
5968
servers.append(config)
6069
try saveServers()
6170
}
62-
71+
6372
/// Remove a server configuration
6473
public func removeServer(_ serverId: UUID) throws {
65-
servers.removeAll { $0.id == serverId }
74+
guard let index = serverIndex[serverId] else {
75+
throw ServerManagerError.serverNotFound
76+
}
77+
servers.remove(at: index)
78+
rebuildIndex()
6679
try saveServers()
6780
}
68-
81+
6982
/// Update a server configuration
7083
public func updateServer(_ config: MCPClient.ServerConfig) throws {
71-
guard let index = servers.firstIndex(where: { $0.id == config.id }) else {
84+
guard let index = serverIndex[config.id] else {
7285
throw ServerManagerError.serverNotFound
7386
}
74-
7587
servers[index] = config
7688
try saveServers()
7789
}
78-
90+
7991
/// Get all server configurations
8092
public func getServers() -> [MCPClient.ServerConfig] {
8193
servers
8294
}
83-
84-
/// Get a server by ID
95+
96+
/// Get a server by ID — O(1)
8597
public func getServer(_ id: UUID) -> MCPClient.ServerConfig? {
86-
servers.first { $0.id == id }
98+
guard let index = serverIndex[id] else { return nil }
99+
return servers[index]
87100
}
88-
101+
89102
/// Enable a server
90103
public func enableServer(_ serverId: UUID) throws {
91-
guard let index = servers.firstIndex(where: { $0.id == serverId }) else {
104+
guard let index = serverIndex[serverId] else {
92105
throw ServerManagerError.serverNotFound
93106
}
94-
95-
var config = servers[index]
96-
config = MCPClient.ServerConfig(
107+
let config = servers[index]
108+
servers[index] = MCPClient.ServerConfig(
97109
id: config.id,
98110
name: config.name,
99111
command: config.command,
100112
arguments: config.arguments,
101113
env: config.env,
102114
enabled: true
103115
)
104-
servers[index] = config
105116
try saveServers()
106117
}
107-
118+
108119
/// Disable a server
109120
public func disableServer(_ serverId: UUID) throws {
110-
guard let index = servers.firstIndex(where: { $0.id == serverId }) else {
121+
guard let index = serverIndex[serverId] else {
111122
throw ServerManagerError.serverNotFound
112123
}
113-
114-
var config = servers[index]
115-
config = MCPClient.ServerConfig(
124+
let config = servers[index]
125+
servers[index] = MCPClient.ServerConfig(
116126
id: config.id,
117127
name: config.name,
118128
command: config.command,
119129
arguments: config.arguments,
120130
env: config.env,
121131
enabled: false
122132
)
123-
servers[index] = config
124133
try saveServers()
125134
}
126-
135+
127136
// MARK: - Presets
128-
137+
129138
/// Common MCP server presets for quick setup
130139
public static let presets: [MCPClient.ServerConfig] = [
131-
// Filesystem server (from official MCP servers)
132140
MCPClient.ServerConfig(
133141
name: "Filesystem",
134142
command: "/usr/local/bin/mcp-server-filesystem",
135143
arguments: ["/Users/\(NSUserName())/Documents"],
136144
enabled: false
137145
),
138-
139-
// GitHub server
140146
MCPClient.ServerConfig(
141147
name: "GitHub",
142148
command: "/usr/local/bin/mcp-server-github",
143149
arguments: [],
144150
env: ["GITHUB_TOKEN": ""],
145151
enabled: false
146152
),
147-
148-
// Puppeteer/Playwright browser automation
149153
MCPClient.ServerConfig(
150154
name: "Puppeteer",
151155
command: "/usr/local/bin/mcp-server-puppeteer",
152156
arguments: [],
153157
enabled: false
154158
),
155-
156-
// SQLite database
157159
MCPClient.ServerConfig(
158160
name: "SQLite",
159161
command: "/usr/local/bin/mcp-server-sqlite",
160162
arguments: [],
161163
enabled: false
162164
),
163-
164-
// Brave Search
165165
MCPClient.ServerConfig(
166166
name: "Brave Search",
167167
command: "/usr/local/bin/mcp-server-brave-search",
168168
arguments: [],
169169
env: ["BRAVE_API_KEY": ""],
170170
enabled: false
171171
),
172-
173-
// Memory/Knowledge graph
174172
MCPClient.ServerConfig(
175173
name: "Memory",
176174
command: "/usr/local/bin/mcp-server-memory",
177175
arguments: [],
178176
enabled: false
179177
),
180-
181-
// Sequential Thinking
182178
MCPClient.ServerConfig(
183179
name: "Sequential Thinking",
184180
command: "/usr/local/bin/mcp-server-sequential-thinking",
@@ -195,7 +191,7 @@ public enum ServerManagerError: LocalizedError {
195191
case serverNotFound
196192
case saveFailed(String)
197193
case loadFailed(String)
198-
194+
199195
public var errorDescription: String? {
200196
switch self {
201197
case .duplicateServer:
@@ -208,4 +204,4 @@ public enum ServerManagerError: LocalizedError {
208204
return "Failed to load servers: \(reason)"
209205
}
210206
}
211-
}
207+
}

0 commit comments

Comments
 (0)