diff --git a/Packages/OsaurusCore/Managers/Documents/DocumentFormatRegistry.swift b/Packages/OsaurusCore/Managers/Documents/DocumentFormatRegistry.swift new file mode 100644 index 000000000..eb77af7e8 --- /dev/null +++ b/Packages/OsaurusCore/Managers/Documents/DocumentFormatRegistry.swift @@ -0,0 +1,108 @@ +// +// DocumentFormatRegistry.swift +// osaurus +// +// Process-wide routing table from a URL (or a `StructuredDocument`) to the +// adapter / emitter / streamer responsible for it. Adapters are registered +// once — at app launch for in-tree formats, at plugin load for plugin- +// provided formats — and looked up every time a file is ingested or an +// artifact is emitted, so the hot path here is the lookup, not registration. +// +// Thread safety: the registry guards its internal state with an `NSLock` +// rather than `@MainActor` isolation. Attachment ingress happens on the +// main actor today, but the agent tool surface (PR 7 in the stage-4 +// roadmap) runs off the main actor, and we don't want every tool call to +// pay an `await`-hop just to look up an adapter. +// + +import Foundation + +public final class DocumentFormatRegistry: @unchecked Sendable { + public static let shared = DocumentFormatRegistry() + + private let lock = NSLock() + + // Insertion order preserved; lookup walks in reverse so the most + // recently-registered claimant wins ties. That lets a plugin override + // a built-in for a specific URL without having to unregister first. + private var adapters: [any DocumentFormatAdapter] = [] + private var emitters: [any DocumentFormatEmitter] = [] + private var streamersByFormatId: [String: any DocumentFormatStreamer] = [:] + + /// `public` so tests can spin up an isolated registry without touching + /// `shared`. Production code should always use `shared`. + public init() {} + + // MARK: - Registration + + public func register(adapter: any DocumentFormatAdapter) { + lock.lock() + defer { lock.unlock() } + adapters.append(adapter) + } + + public func register(emitter: any DocumentFormatEmitter) { + lock.lock() + defer { lock.unlock() } + emitters.append(emitter) + } + + public func register(streamer: any DocumentFormatStreamer) { + lock.lock() + defer { lock.unlock() } + streamersByFormatId[streamer.formatId] = streamer + } + + /// Removes every registration (adapter, emitter, streamer) whose + /// `formatId` matches. Returns `true` if anything was actually removed. + /// Used by plugin unload and by tests that want a clean slate. + @discardableResult + public func unregisterAll(formatId: String) -> Bool { + lock.lock() + defer { lock.unlock() } + let before = adapters.count + emitters.count + streamersByFormatId.count + adapters.removeAll { $0.formatId == formatId } + emitters.removeAll { $0.formatId == formatId } + streamersByFormatId.removeValue(forKey: formatId) + let after = adapters.count + emitters.count + streamersByFormatId.count + return before != after + } + + // MARK: - Lookup + + /// Returns the most-recently-registered adapter whose `canHandle` + /// accepts the URL. `nil` when nothing claims it — callers can + /// decide whether to fall through to a legacy path or throw + /// `DocumentAdapterError.unsupportedFormat`. + public func adapter(for url: URL, uti: String? = nil) -> (any DocumentFormatAdapter)? { + lock.lock() + defer { lock.unlock() } + return adapters.reversed().first(where: { $0.canHandle(url: url, uti: uti) }) + } + + public func emitter(for document: StructuredDocument) -> (any DocumentFormatEmitter)? { + lock.lock() + defer { lock.unlock() } + return emitters.reversed().first(where: { $0.canEmit(document) }) + } + + public func streamer(forFormatId id: String) -> (any DocumentFormatStreamer)? { + lock.lock() + defer { lock.unlock() } + return streamersByFormatId[id] + } + + // MARK: - Introspection + + /// Union of format ids currently registered across adapters, emitters, + /// and streamers. Useful for plugin-host diagnostics and for tests. + public func registeredFormatIds() -> Set { + lock.lock() + defer { lock.unlock() } + var ids: Set = [] + for adapter in adapters { ids.insert(adapter.formatId) } + for emitter in emitters { ids.insert(emitter.formatId) } + for id in streamersByFormatId.keys { ids.insert(id) } + return ids + } +} diff --git a/Packages/OsaurusCore/Models/Documents/DocumentAdapterError.swift b/Packages/OsaurusCore/Models/Documents/DocumentAdapterError.swift new file mode 100644 index 000000000..69924c21f --- /dev/null +++ b/Packages/OsaurusCore/Models/Documents/DocumentAdapterError.swift @@ -0,0 +1,37 @@ +// +// DocumentAdapterError.swift +// osaurus +// +// Shared error surface for every document format adapter, emitter, and +// streamer. Callers that don't care about the specific failure can still +// catch the common cases (size, cancellation) at the protocol boundary +// without per-format knowledge. +// + +import Foundation + +public enum DocumentAdapterError: LocalizedError, Sendable { + case unsupportedFormat(formatId: String) + case sizeLimitExceeded(actual: Int64, limit: Int64) + case readFailed(underlying: String) + case writeFailed(underlying: String) + case emptyContent + case cancelled + + public var errorDescription: String? { + switch self { + case .unsupportedFormat(let id): + return "No registered adapter for format '\(id)'" + case .sizeLimitExceeded(let actual, let limit): + return "File exceeds size limit (\(actual) bytes > \(limit) bytes)" + case .readFailed(let reason): + return "Document read failed: \(reason)" + case .writeFailed(let reason): + return "Document write failed: \(reason)" + case .emptyContent: + return "Document contains no readable content" + case .cancelled: + return "Document parse was cancelled" + } + } +} diff --git a/Packages/OsaurusCore/Models/Documents/DocumentFormatAdapter.swift b/Packages/OsaurusCore/Models/Documents/DocumentFormatAdapter.swift new file mode 100644 index 000000000..0ee12e5f2 --- /dev/null +++ b/Packages/OsaurusCore/Models/Documents/DocumentFormatAdapter.swift @@ -0,0 +1,34 @@ +// +// DocumentFormatAdapter.swift +// osaurus +// +// Read-side handler for a single file format. Adapters are registered +// with `DocumentFormatRegistry` at app launch (for in-tree adapters) or +// at plugin load (for plugin-provided adapters), then looked up every +// time a file is ingested. `canHandle` is kept separate from `parse` so +// the registry can iterate candidates without paying parse cost on every +// file type that doesn't match. +// + +import Foundation + +public protocol DocumentFormatAdapter: Sendable { + /// Stable identifier used for logging, registry tie-breaks, and as the + /// plugin registration key. Examples: "xlsx", "docx", "pdf", "csv". + var formatId: String { get } + + /// Lightweight precondition check. Must NOT open the file; it is called + /// before `parse` as the registry enumerates adapters. An adapter that + /// narrows (e.g. "I only handle PDFs with a tagged text layer") uses + /// this hook to defer to a more permissive adapter registered later. + func canHandle(url: URL, uti: String?) -> Bool + + /// Parse a file into its typed representation. Adapters that read the + /// whole file into memory must throw + /// `DocumentAdapterError.sizeLimitExceeded` when the file exceeds + /// `sizeLimit`. Streaming adapters may return a `StructuredDocument` + /// whose representation carries an async stream (see CSVTable in + /// stage-4 PR 4). The registry supplies a per-format cap from + /// `DocumentLimits`. + func parse(url: URL, sizeLimit: Int64) async throws -> StructuredDocument +} diff --git a/Packages/OsaurusCore/Models/Documents/DocumentFormatEmitter.swift b/Packages/OsaurusCore/Models/Documents/DocumentFormatEmitter.swift new file mode 100644 index 000000000..22f280d44 --- /dev/null +++ b/Packages/OsaurusCore/Models/Documents/DocumentFormatEmitter.swift @@ -0,0 +1,25 @@ +// +// DocumentFormatEmitter.swift +// osaurus +// +// Write-side handler for a single file format. Deliberately split from +// `DocumentFormatAdapter` so read-only formats (XLS, OFX 1.x, PPTX in +// the stage-2 priority list) don't have to carry a stub writer. Sandbox +// containment and destination checking live in the caller +// (`ShareArtifactTool` today); emitters are just byte producers. +// + +import Foundation + +public protocol DocumentFormatEmitter: Sendable { + var formatId: String { get } + + /// Keyed on the concrete shape of `document.representation`, not on + /// the file extension — an emitter produces exactly one representation + /// shape, so the registry uses this hook to pick the right writer. + func canEmit(_ document: StructuredDocument) -> Bool + + /// Write the document to `url`. The caller is responsible for having + /// already resolved and contained `url`; the emitter writes raw bytes. + func emit(_ document: StructuredDocument, to url: URL) async throws +} diff --git a/Packages/OsaurusCore/Models/Documents/DocumentFormatStreamer.swift b/Packages/OsaurusCore/Models/Documents/DocumentFormatStreamer.swift new file mode 100644 index 000000000..159f6e2e2 --- /dev/null +++ b/Packages/OsaurusCore/Models/Documents/DocumentFormatStreamer.swift @@ -0,0 +1,25 @@ +// +// DocumentFormatStreamer.swift +// osaurus +// +// Optional streaming read for formats where whole-file-into-memory is +// not viable — multi-GB CSVs, page-by-page PDFs, row-by-row XLSX. +// Orthogonal to `DocumentFormatAdapter` so small formats (QIF, MT940) +// can skip the streaming surface entirely. Callers that can back- +// pressure (the agent tool surface in particular) prefer streaming when +// both an adapter and a streamer are registered for the same format. +// + +import Foundation + +public protocol DocumentFormatStreamer: Sendable { + associatedtype Element: Sendable + + var formatId: String { get } + + /// Stream format-native records out of the file. `AsyncThrowingStream` + /// gives callers cancellation and back-pressure without bespoke + /// plumbing; adapters that can't produce records incrementally should + /// not conform to this protocol at all. + func stream(url: URL) -> AsyncThrowingStream +} diff --git a/Packages/OsaurusCore/Models/Documents/DocumentLimits.swift b/Packages/OsaurusCore/Models/Documents/DocumentLimits.swift new file mode 100644 index 000000000..4dbe3cd57 --- /dev/null +++ b/Packages/OsaurusCore/Models/Documents/DocumentLimits.swift @@ -0,0 +1,38 @@ +// +// DocumentLimits.swift +// osaurus +// +// Per-format byte ceilings applied to in-memory parsing. Streaming +// adapters are not bound by these caps — they negotiate back-pressure +// with their caller — but any adapter that reads the whole file into +// memory must honour the limit returned by `limit(forFormatId:)`. +// +// The numeric defaults here are intentionally generous compared to the +// 500 KB text cap on the legacy `DocumentParser`. They exist to prevent +// OOM under adversarial input, not to shape the user-facing attachment +// experience; the chat attachment flow keeps its own smaller caps. +// + +import Foundation + +public enum DocumentLimits { + public static let plainText: Int64 = 5 * 1024 * 1024 + public static let csv: Int64 = 25 * 1024 * 1024 + public static let xlsx: Int64 = 50 * 1024 * 1024 + public static let pdf: Int64 = 100 * 1024 * 1024 + public static let docx: Int64 = 50 * 1024 * 1024 + + /// Fallback for formats that haven't been assigned a tuned cap. + public static let defaultLimit: Int64 = 10 * 1024 * 1024 + + public static func limit(forFormatId id: String) -> Int64 { + switch id.lowercased() { + case "plaintext", "text", "txt", "md", "markdown": return plainText + case "csv", "tsv": return csv + case "xlsx", "xls", "ods": return xlsx + case "pdf": return pdf + case "docx", "doc", "rtf": return docx + default: return defaultLimit + } + } +} diff --git a/Packages/OsaurusCore/Models/Documents/StructuredDocument.swift b/Packages/OsaurusCore/Models/Documents/StructuredDocument.swift new file mode 100644 index 000000000..d8f7ef519 --- /dev/null +++ b/Packages/OsaurusCore/Models/Documents/StructuredDocument.swift @@ -0,0 +1,56 @@ +// +// StructuredDocument.swift +// osaurus +// +// Typed parse result that carries both a format-native representation +// AND a plain-text fallback. The fallback is load-bearing because the +// existing chat attachment flow consumes +// `Attachment.Kind.document(content: String, …)`; keeping a text view +// on every parsed document lets adapters migrate onto the typed surface +// one at a time without breaking that contract. +// + +import Foundation + +/// Marker protocol for per-format typed representations (`Workbook`, +/// `WordDocument`, `PDFDocument`, …). Concrete types live next to their +/// adapter under `Packages/OsaurusCore/Models/Documents//`. +public protocol StructuredRepresentation: Sendable {} + +/// Type-erasing container so a `StructuredDocument` can cross layers +/// (registry, tool surface, artifact pipeline) without leaking the +/// concrete representation type into every caller. +public struct AnyStructuredRepresentation: @unchecked Sendable { + public let formatId: String + public let underlying: any StructuredRepresentation + + public init(formatId: String, underlying: any StructuredRepresentation) { + self.formatId = formatId + self.underlying = underlying + } +} + +public struct StructuredDocument: @unchecked Sendable { + public let formatId: String + public let filename: String + public let fileSize: Int64 + public let representation: AnyStructuredRepresentation + public let textFallback: String + public let createdAt: Date + + public init( + formatId: String, + filename: String, + fileSize: Int64, + representation: AnyStructuredRepresentation, + textFallback: String, + createdAt: Date = Date() + ) { + self.formatId = formatId + self.filename = filename + self.fileSize = fileSize + self.representation = representation + self.textFallback = textFallback + self.createdAt = createdAt + } +} diff --git a/Packages/OsaurusCore/Services/Keychain/AgentSecretsKeychain.swift b/Packages/OsaurusCore/Services/Keychain/AgentSecretsKeychain.swift index 1dc88055a..9803ef988 100644 --- a/Packages/OsaurusCore/Services/Keychain/AgentSecretsKeychain.swift +++ b/Packages/OsaurusCore/Services/Keychain/AgentSecretsKeychain.swift @@ -15,11 +15,76 @@ import Security public enum AgentSecretsKeychain { private static let service = "ai.osaurus.agent-secrets" + #if DEBUG + private static let inMemoryStoreLock = NSLock() + nonisolated(unsafe) private static var inMemoryStoreForTesting: [String: String]? + + static func _withInMemoryStoreForTesting( + _ body: () throws -> T + ) rethrows -> T { + inMemoryStoreLock.lock() + let previous = inMemoryStoreForTesting + inMemoryStoreForTesting = [:] + inMemoryStoreLock.unlock() + + defer { + inMemoryStoreLock.lock() + inMemoryStoreForTesting = previous + inMemoryStoreLock.unlock() + } + + return try body() + } + + private static func testingSave(_ value: String, account: String) -> (enabled: Bool, saved: Bool) { + inMemoryStoreLock.lock() + defer { inMemoryStoreLock.unlock() } + guard inMemoryStoreForTesting != nil else { + return (enabled: false, saved: false) + } + inMemoryStoreForTesting?[account] = value + return (enabled: true, saved: true) + } + + private static func testingGet(account: String) -> (enabled: Bool, value: String?) { + inMemoryStoreLock.lock() + defer { inMemoryStoreLock.unlock() } + guard let store = inMemoryStoreForTesting else { + return (enabled: false, value: nil) + } + return (enabled: true, value: store[account]) + } + + private static func testingDelete(account: String) -> (enabled: Bool, deleted: Bool) { + inMemoryStoreLock.lock() + defer { inMemoryStoreLock.unlock() } + guard inMemoryStoreForTesting != nil else { + return (enabled: false, deleted: false) + } + inMemoryStoreForTesting?[account] = nil + return (enabled: true, deleted: true) + } + + private static func testingAllAccounts() -> [String]? { + inMemoryStoreLock.lock() + defer { inMemoryStoreLock.unlock() } + guard let store = inMemoryStoreForTesting else { return nil } + return Array(store.keys) + } + #endif + @discardableResult public static func saveSecret(_ value: String, id: String, agentId: UUID) -> Bool { let account = "\(agentId.uuidString).\(id)" guard let valueData = value.data(using: .utf8) else { return false } + #if DEBUG + let testing = testingSave(value, account: account) + if testing.enabled { + return testing.saved + } + #endif + deleteSecret(id: id, agentId: agentId) let query: [String: Any] = [ @@ -35,6 +100,13 @@ public enum AgentSecretsKeychain { public static func getSecret(id: String, agentId: UUID) -> String? { let account = "\(agentId.uuidString).\(id)" + #if DEBUG + let testing = testingGet(account: account) + if testing.enabled { + return testing.value + } + #endif + let query: [String: Any] = [ kSecClass as String: kSecClassGenericPassword, kSecAttrService as String: service, @@ -55,6 +127,13 @@ public enum AgentSecretsKeychain { public static func deleteSecret(id: String, agentId: UUID) -> Bool { let account = "\(agentId.uuidString).\(id)" + #if DEBUG + let testing = testingDelete(account: account) + if testing.enabled { + return testing.deleted + } + #endif + let query: [String: Any] = [ kSecClass as String: kSecClassGenericPassword, kSecAttrService as String: service, @@ -119,6 +198,12 @@ public enum AgentSecretsKeychain { // MARK: - Private private static func allAccounts() -> [String] { + #if DEBUG + if let accounts = testingAllAccounts() { + return accounts + } + #endif + let query: [String: Any] = [ kSecClass as String: kSecClassGenericPassword, kSecAttrService as String: service, diff --git a/Packages/OsaurusCore/Tests/Chat/ChatViewSandboxTests.swift b/Packages/OsaurusCore/Tests/Chat/ChatViewSandboxTests.swift index 1f5686d86..8d894ded4 100644 --- a/Packages/OsaurusCore/Tests/Chat/ChatViewSandboxTests.swift +++ b/Packages/OsaurusCore/Tests/Chat/ChatViewSandboxTests.swift @@ -46,29 +46,32 @@ struct ChatViewSandboxTests { @Test func estimatedContextBreakdown_includesSandboxPromptAndToolsWhenEnabled() async { - let manager = AgentManager.shared - let originalActiveAgentId = manager.activeAgentId - let inactiveAgent = Agent(name: "Chat Estimate Off") - let sandboxAgent = Agent( - name: "Chat Estimate On", - autonomousExec: AutonomousExecConfig(enabled: true) - ) - manager.add(inactiveAgent) - manager.add(sandboxAgent) - defer { - manager.setActiveAgent(originalActiveAgentId) - Task { - _ = await manager.delete(id: inactiveAgent.id) - _ = await manager.delete(id: sandboxAgent.id) - } - } + await SandboxTestLock.shared.run { + let manager = AgentManager.shared + let originalActiveAgentId = manager.activeAgentId + let inactiveAgent = Agent( + name: "Chat Estimate Off", + agentAddress: "test-chat-estimate-off" + ) + let sandboxAgent = Agent( + name: "Chat Estimate On", + agentAddress: "test-chat-estimate-on", + autonomousExec: AutonomousExecConfig(enabled: true) + ) + manager.add(inactiveAgent) + manager.add(sandboxAgent) - let inactiveSession = ChatSession() - inactiveSession.agentId = inactiveAgent.id - let sandboxSession = ChatSession() - sandboxSession.agentId = sandboxAgent.id + let inactiveSession = ChatSession() + inactiveSession.agentId = inactiveAgent.id + let sandboxSession = ChatSession() + sandboxSession.agentId = sandboxAgent.id + + BuiltinSandboxTools.register( + agentId: sandboxAgent.id.uuidString, + agentName: sandboxAgent.name, + config: AutonomousExecConfig(enabled: true) + ) - await withRegisteredSandboxBuiltins { let inactiveBreakdown = inactiveSession.estimatedContextBreakdown let sandboxBreakdown = sandboxSession.estimatedContextBreakdown @@ -80,6 +83,11 @@ struct ChatViewSandboxTests { let inactiveToolTokens = inactiveBreakdown.context.first { $0.id == "tools" }?.tokens ?? 0 #expect(sandboxToolTokens > inactiveToolTokens) #expect(sandboxToolTokens >= ToolRegistry.shared.estimatedTokens(for: "sandbox_exec")) + + ToolRegistry.shared.unregisterAllSandboxTools() + manager.setActiveAgent(originalActiveAgentId) + _ = await manager.delete(id: inactiveAgent.id) + _ = await manager.delete(id: sandboxAgent.id) } } @@ -127,9 +135,13 @@ struct ChatViewSandboxTests { let originalStatus = SandboxManager.State.shared.status let originalProvisionOverride = registrar.provisionAgentOverride - let inactiveAgent = Agent(name: "Chat Sandbox Off") + let inactiveAgent = Agent( + name: "Chat Sandbox Off", + agentAddress: "test-chat-sandbox-off" + ) let sandboxAgent = Agent( name: "Chat Sandbox On", + agentAddress: "test-chat-sandbox-on", autonomousExec: AutonomousExecConfig(enabled: true) ) manager.add(inactiveAgent) @@ -138,6 +150,11 @@ struct ChatViewSandboxTests { SandboxManager.State.shared.status = .running registrar.provisionAgentOverride = { _ in } + BuiltinSandboxTools.register( + agentId: sandboxAgent.id.uuidString, + agentName: sandboxAgent.name, + config: AutonomousExecConfig(enabled: true) + ) let session = ChatSession() let inactiveMode = await session.prepareChatExecutionMode(agentId: inactiveAgent.id) diff --git a/Packages/OsaurusCore/Tests/Chat/SessionPreflightCacheTests.swift b/Packages/OsaurusCore/Tests/Chat/SessionPreflightCacheTests.swift index bc4a00456..6a6fde92d 100644 --- a/Packages/OsaurusCore/Tests/Chat/SessionPreflightCacheTests.swift +++ b/Packages/OsaurusCore/Tests/Chat/SessionPreflightCacheTests.swift @@ -32,79 +32,82 @@ struct SessionPreflightCacheTests { @Test func resolveTools_includesAdditionalToolNamesEvenWithEmptyPreflight() async { - let manager = AgentManager.shared - let agent = Agent( - name: "SessionPreflightCacheTestAgent-\(UUID().uuidString.prefix(6))" - ) - manager.add(agent) - defer { Task { _ = await manager.delete(id: agent.id) } } - - // Empty preflight (mirrors a cached session that captured "no LLM - // additions") should still inflate to include the agent's - // capabilities_load union. - let tools = SystemPromptComposer.resolveTools( - agentId: agent.id, - executionMode: .none, - preflight: PreflightResult(toolSpecs: [], items: []), - additionalToolNames: ["search_memory"] - ) - let names = tools.map { $0.function.name } - #expect(names.contains("search_memory")) + await withSessionPreflightAgent { agentId in + + // Empty preflight (mirrors a cached session that captured "no LLM + // additions") should still inflate to include the agent's + // capabilities_load union. + let tools = SystemPromptComposer.resolveTools( + agentId: agentId, + executionMode: .none, + preflight: PreflightResult(toolSpecs: [], items: []), + additionalToolNames: ["search_memory"] + ) + let names = tools.map { $0.function.name } + #expect(names.contains("search_memory")) + } } @Test func composeChatContext_doesNotRunFreshPreflightWhenCached() async { - let manager = AgentManager.shared - let agent = Agent( - name: "SessionPreflightCacheTestAgent-\(UUID().uuidString.prefix(6))" - ) - manager.add(agent) - defer { Task { _ = await manager.delete(id: agent.id) } } - - // Seed cache with a known PreflightResult that includes a specific - // tool we can fingerprint in the rendered output. - let memorySpec = ToolRegistry.shared.specs(forTools: ["search_memory"]).first - guard let memorySpec else { - // search_memory isn't registered in this test environment — skip - // (the property under test is exercised by other tests anyway). - return - } - let cached = PreflightResult(toolSpecs: [memorySpec], items: []) + await withSessionPreflightAgent { agentId in - let ctx = await SystemPromptComposer.composeChatContext( - agentId: agent.id, - executionMode: .none, - query: "this query would normally trigger a fresh LLM preflight", - cachedPreflight: cached - ) + // Seed cache with a known PreflightResult that includes a specific + // tool we can fingerprint in the rendered output. + let memorySpec = ToolRegistry.shared.specs(forTools: ["search_memory"]).first + guard let memorySpec else { + // search_memory isn't registered in this test environment — skip + // (the property under test is exercised by other tests anyway). + return + } + let cached = PreflightResult(toolSpecs: [memorySpec], items: []) + + let ctx = await SystemPromptComposer.composeChatContext( + agentId: agentId, + executionMode: .none, + query: "this query would normally trigger a fresh LLM preflight", + cachedPreflight: cached + ) - // The cached preflight must echo back through ComposedContext.preflight - // so the caller can re-stash it. - let cachedNames = Set(ctx.preflight.toolSpecs.map { $0.function.name }) - #expect(cachedNames == ["search_memory"]) - // And the resolved tool union must contain the cached preflight tool. - let resolvedNames = ctx.tools.map { $0.function.name } - #expect(resolvedNames.contains("search_memory")) + // The cached preflight must echo back through ComposedContext.preflight + // so the caller can re-stash it. + let cachedNames = Set(ctx.preflight.toolSpecs.map { $0.function.name }) + #expect(cachedNames == ["search_memory"]) + // And the resolved tool union must contain the cached preflight tool. + let resolvedNames = ctx.tools.map { $0.function.name } + #expect(resolvedNames.contains("search_memory")) + } } @Test func composeChatContext_returnsMemorySectionSeparately() async { - let manager = AgentManager.shared - let agent = Agent( - name: "SessionPreflightCacheTestAgent-\(UUID().uuidString.prefix(6))" - ) - manager.add(agent) - defer { Task { _ = await manager.delete(id: agent.id) } } + await withSessionPreflightAgent { agentId in - let ctx = await SystemPromptComposer.composeChatContext( - agentId: agent.id, - executionMode: .none - ) + let ctx = await SystemPromptComposer.composeChatContext( + agentId: agentId, + executionMode: .none + ) - // Even when memory has no content for a brand-new agent, the - // rendered system prompt must NOT contain a [Memory] block — the - // helper is the only writer of that marker, and it goes onto the - // user message instead. - #expect(ctx.prompt.contains("[Memory]") == false) + // Even when memory has no content for a brand-new agent, the + // rendered system prompt must NOT contain a [Memory] block — the + // helper is the only writer of that marker, and it goes onto the + // user message instead. + #expect(ctx.prompt.contains("[Memory]") == false) + } + } + + private func withSessionPreflightAgent( + _ body: @MainActor @Sendable (UUID) async -> Void + ) async { + await SandboxTestLock.shared.run { + let manager = AgentManager.shared + let agent = Agent( + name: "SessionPreflightCacheTestAgent-\(UUID().uuidString.prefix(6))", + agentAddress: "test-session-preflight-\(UUID().uuidString)" + ) + manager.add(agent) + await body(agent.id) + _ = await manager.delete(id: agent.id) + } } } diff --git a/Packages/OsaurusCore/Tests/Chat/SystemPromptComposerToolResolutionTests.swift b/Packages/OsaurusCore/Tests/Chat/SystemPromptComposerToolResolutionTests.swift index 418e75ed3..e2e08a316 100644 --- a/Packages/OsaurusCore/Tests/Chat/SystemPromptComposerToolResolutionTests.swift +++ b/Packages/OsaurusCore/Tests/Chat/SystemPromptComposerToolResolutionTests.swift @@ -25,36 +25,40 @@ struct SystemPromptComposerToolResolutionTests { private func withSandboxAgent( autonomous: Bool, manualToolNames: [String]? = nil, - body: (UUID) async -> Void + body: @MainActor @Sendable (UUID) async -> Void ) async { - let manager = AgentManager.shared - let agent: Agent - if let names = manualToolNames { - agent = Agent( - name: "ToolResolutionTestAgent-\(UUID().uuidString.prefix(6))", - autonomousExec: autonomous ? AutonomousExecConfig(enabled: true) : nil, - toolSelectionMode: .manual, - manualToolNames: names - ) - } else { - agent = Agent( - name: "ToolResolutionTestAgent-\(UUID().uuidString.prefix(6))", - autonomousExec: autonomous ? AutonomousExecConfig(enabled: true) : nil - ) + await SandboxTestLock.shared.run { + let manager = AgentManager.shared + let agent: Agent + if let names = manualToolNames { + agent = Agent( + name: "ToolResolutionTestAgent-\(UUID().uuidString.prefix(6))", + agentAddress: "test-tool-resolution-\(UUID().uuidString)", + autonomousExec: autonomous ? AutonomousExecConfig(enabled: true) : nil, + toolSelectionMode: .manual, + manualToolNames: names + ) + } else { + agent = Agent( + name: "ToolResolutionTestAgent-\(UUID().uuidString.prefix(6))", + agentAddress: "test-tool-resolution-\(UUID().uuidString)", + autonomousExec: autonomous ? AutonomousExecConfig(enabled: true) : nil + ) + } + manager.add(agent) + await body(agent.id) + _ = await manager.delete(id: agent.id) } - manager.add(agent) - defer { Task { _ = await manager.delete(id: agent.id) } } - await body(agent.id) } - private func registerSandboxBuiltins(_ body: () -> Void) { + private func withRegisteredSandboxBuiltins(_ body: @MainActor @Sendable () -> Void) { BuiltinSandboxTools.register( agentId: "tool-resolution-test", agentName: "tool-resolution-test", config: AutonomousExecConfig(enabled: true) ) - defer { ToolRegistry.shared.unregisterAllSandboxTools() } body() + ToolRegistry.shared.unregisterAllSandboxTools() } // MARK: - Auto mode @@ -101,7 +105,7 @@ struct SystemPromptComposerToolResolutionTests { @Test func manualMode_includesSandboxBuiltinsWhenSandboxActive() async { await withSandboxAgent(autonomous: true, manualToolNames: ["render_chart"]) { agentId in - registerSandboxBuiltins { + withRegisteredSandboxBuiltins { let tools = SystemPromptComposer.resolveTools( agentId: agentId, executionMode: .sandbox @@ -120,7 +124,7 @@ struct SystemPromptComposerToolResolutionTests { @Test func manualMode_emptyManualNames_stillIncludesAlwaysLoaded() async { await withSandboxAgent(autonomous: true, manualToolNames: []) { agentId in - registerSandboxBuiltins { + withRegisteredSandboxBuiltins { let tools = SystemPromptComposer.resolveTools( agentId: agentId, executionMode: .sandbox @@ -155,7 +159,7 @@ struct SystemPromptComposerToolResolutionTests { } await withSandboxAgent(autonomous: true) { agentId in - registerSandboxBuiltins { + withRegisteredSandboxBuiltins { let names = Set( SystemPromptComposer.resolveTools(agentId: agentId, executionMode: .sandbox) .map { $0.function.name } @@ -171,7 +175,7 @@ struct SystemPromptComposerToolResolutionTests { @Test func canonicalToolOrder_pinsLoopToolsToTheTop() async { await withSandboxAgent(autonomous: true) { agentId in - registerSandboxBuiltins { + withRegisteredSandboxBuiltins { let names = SystemPromptComposer.resolveTools( agentId: agentId, executionMode: .sandbox @@ -190,7 +194,7 @@ struct SystemPromptComposerToolResolutionTests { @Test func toolsDisabled_returnsEmpty() async { await withSandboxAgent(autonomous: true) { agentId in - registerSandboxBuiltins { + withRegisteredSandboxBuiltins { let tools = SystemPromptComposer.resolveTools( agentId: agentId, executionMode: .sandbox, @@ -263,7 +267,7 @@ struct SystemPromptComposerToolResolutionTests { @Test func canonicalToolOrder_isStableAcrossInvocations() async { await withSandboxAgent(autonomous: true) { agentId in - registerSandboxBuiltins { + withRegisteredSandboxBuiltins { // Two compositions with identical inputs must return the // exact same tool ordering — that's what makes the rendered // block byte-stable across sends. diff --git a/Packages/OsaurusCore/Tests/Context/PluginCreatorInjectionTests.swift b/Packages/OsaurusCore/Tests/Context/PluginCreatorInjectionTests.swift index 77942e4d0..49f5a733a 100644 --- a/Packages/OsaurusCore/Tests/Context/PluginCreatorInjectionTests.swift +++ b/Packages/OsaurusCore/Tests/Context/PluginCreatorInjectionTests.swift @@ -137,18 +137,22 @@ struct PluginCreatorComposerWiringTests { @Test func composeChatContext_skipsPluginCreatorOutsideSandbox() async { - let agent = Agent( - name: "Plugin Creator Non-Sandbox Agent", - autonomousExec: AutonomousExecConfig(enabled: false, pluginCreate: true) - ) - AgentManager.shared.add(agent) - defer { Task { _ = await AgentManager.shared.delete(id: agent.id) } } - - let context = await SystemPromptComposer.composeChatContext( - agentId: agent.id, - executionMode: .none - ) - let labels = context.manifest.sections.map(\.label) - #expect(labels.contains("Plugin Creator") == false) + await SandboxTestLock.shared.run { + let agent = Agent( + name: "Plugin Creator Non-Sandbox Agent", + agentAddress: "test-plugin-creator-\(UUID().uuidString)", + autonomousExec: AutonomousExecConfig(enabled: false, pluginCreate: true) + ) + AgentManager.shared.add(agent) + + let context = await SystemPromptComposer.composeChatContext( + agentId: agent.id, + executionMode: .none + ) + let labels = context.manifest.sections.map(\.label) + #expect(labels.contains("Plugin Creator") == false) + + _ = await AgentManager.shared.delete(id: agent.id) + } } } diff --git a/Packages/OsaurusCore/Tests/Documents/DocumentFormatRegistryTests.swift b/Packages/OsaurusCore/Tests/Documents/DocumentFormatRegistryTests.swift new file mode 100644 index 000000000..0eb1ab27a --- /dev/null +++ b/Packages/OsaurusCore/Tests/Documents/DocumentFormatRegistryTests.swift @@ -0,0 +1,143 @@ +// +// DocumentFormatRegistryTests.swift +// osaurusTests +// +// Contract tests for the registry. Covers the three invariants we care +// about for PR 1: adapters are routed by `canHandle`, later registrations +// win ties, and `unregisterAll` is a no-op when nothing matches. The +// thread-safety test just asserts that the internal lock serialises +// concurrent registrations without losing entries. +// + +import Foundation +import Testing + +@testable import OsaurusCore + +@Suite("DocumentFormatRegistry") +struct DocumentFormatRegistryTests { + + // MARK: - Adapter routing + + @Test func adapter_returnsNilWhenNoneRegistered() { + let registry = DocumentFormatRegistry() + #expect(registry.adapter(for: URL(fileURLWithPath: "/tmp/unknown.xyz")) == nil) + } + + @Test func adapter_routesByCanHandle() { + let registry = DocumentFormatRegistry() + registry.register(adapter: FakeAdapter(formatId: "pdf", extensions: ["pdf"])) + registry.register(adapter: FakeAdapter(formatId: "docx", extensions: ["docx"])) + + #expect(registry.adapter(for: URL(fileURLWithPath: "/tmp/a.pdf"))?.formatId == "pdf") + #expect(registry.adapter(for: URL(fileURLWithPath: "/tmp/a.docx"))?.formatId == "docx") + #expect(registry.adapter(for: URL(fileURLWithPath: "/tmp/a.rtf")) == nil) + } + + @Test func adapter_laterRegistrationWinsTie() { + let registry = DocumentFormatRegistry() + registry.register(adapter: FakeAdapter(formatId: "pdf-builtin", extensions: ["pdf"])) + registry.register(adapter: FakeAdapter(formatId: "pdf-plugin", extensions: ["pdf"])) + + #expect(registry.adapter(for: URL(fileURLWithPath: "/tmp/a.pdf"))?.formatId == "pdf-plugin") + } + + // MARK: - Emitter and streamer + + @Test func emitter_pickedByCanEmit() { + let registry = DocumentFormatRegistry() + registry.register(emitter: FakeEmitter(formatId: "xlsx")) + registry.register(emitter: FakeEmitter(formatId: "docx")) + + #expect(registry.emitter(for: Self.fakeDocument(formatId: "docx"))?.formatId == "docx") + #expect(registry.emitter(for: Self.fakeDocument(formatId: "xlsx"))?.formatId == "xlsx") + #expect(registry.emitter(for: Self.fakeDocument(formatId: "pdf")) == nil) + } + + @Test func streamer_pickedByFormatId() { + let registry = DocumentFormatRegistry() + registry.register(streamer: FakeStreamer(formatId: "csv")) + + #expect(registry.streamer(forFormatId: "csv")?.formatId == "csv") + #expect(registry.streamer(forFormatId: "xlsx") == nil) + } + + // MARK: - Unregistration + + @Test func unregisterAll_dropsAllRegistrationsForFormatId() { + let registry = DocumentFormatRegistry() + registry.register(adapter: FakeAdapter(formatId: "xlsx", extensions: ["xlsx"])) + registry.register(emitter: FakeEmitter(formatId: "xlsx")) + registry.register(streamer: FakeStreamer(formatId: "xlsx")) + + #expect(registry.registeredFormatIds() == ["xlsx"]) + #expect(registry.unregisterAll(formatId: "xlsx")) + #expect(registry.registeredFormatIds().isEmpty) + #expect(registry.unregisterAll(formatId: "xlsx") == false) + } + + // MARK: - Thread safety + + @Test func register_isThreadSafe() async { + let registry = DocumentFormatRegistry() + await withTaskGroup(of: Void.self) { group in + for index in 0 ..< 200 { + group.addTask { + registry.register(adapter: FakeAdapter(formatId: "fmt-\(index)", extensions: ["x"])) + } + } + } + #expect(registry.registeredFormatIds().count == 200) + } + + // MARK: - Fixtures + + private static func fakeDocument(formatId: String) -> StructuredDocument { + StructuredDocument( + formatId: formatId, + filename: "fixture.\(formatId)", + fileSize: 0, + representation: AnyStructuredRepresentation( + formatId: formatId, + underlying: EmptyRepresentation() + ), + textFallback: "" + ) + } + + private struct EmptyRepresentation: StructuredRepresentation {} + + private struct FakeAdapter: DocumentFormatAdapter { + let formatId: String + let extensions: Set + + func canHandle(url: URL, uti: String?) -> Bool { + extensions.contains(url.pathExtension.lowercased()) + } + + func parse(url: URL, sizeLimit: Int64) async throws -> StructuredDocument { + DocumentFormatRegistryTests.fakeDocument(formatId: formatId) + } + } + + private struct FakeEmitter: DocumentFormatEmitter { + let formatId: String + + func canEmit(_ document: StructuredDocument) -> Bool { + document.formatId == formatId + } + + func emit(_ document: StructuredDocument, to url: URL) async throws {} + } + + private struct FakeStreamer: DocumentFormatStreamer { + typealias Element = String + let formatId: String + + func stream(url: URL) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + continuation.finish() + } + } + } +} diff --git a/Packages/OsaurusCore/Tests/Helpers/SandboxTestLock.swift b/Packages/OsaurusCore/Tests/Helpers/SandboxTestLock.swift index 15cc44a17..adff07124 100644 --- a/Packages/OsaurusCore/Tests/Helpers/SandboxTestLock.swift +++ b/Packages/OsaurusCore/Tests/Helpers/SandboxTestLock.swift @@ -2,9 +2,10 @@ // SandboxTestLock.swift // OsaurusCoreTests // -// Process-wide serialization for tests that mutate sandbox globals: +// Process-wide serialization for tests that mutate sandbox-adjacent globals: // `ToolRegistry` sandbox tools, `SandboxManager.State`, -// `SandboxToolRegistrar` overrides, or `HostAPIBridgeServer.shared`. +// `SandboxToolRegistrar` overrides, `HostAPIBridgeServer.shared`, +// or synthetic `AgentManager.shared` agents used to resolve sandbox modes. // import Foundation diff --git a/Packages/OsaurusCore/Tests/Sandbox/SandboxPluginRegistrationTests.swift b/Packages/OsaurusCore/Tests/Sandbox/SandboxPluginRegistrationTests.swift index 0f30f0a66..9c611f93d 100644 --- a/Packages/OsaurusCore/Tests/Sandbox/SandboxPluginRegistrationTests.swift +++ b/Packages/OsaurusCore/Tests/Sandbox/SandboxPluginRegistrationTests.swift @@ -104,37 +104,38 @@ struct SandboxPluginRegistrationTests { @Test func validateAndStage_rejectsMissingSecrets() { - // Use a fresh random agent UUID so the keychain lookup definitely - // misses (the keychain is shared across the test runner). - let agentId = UUID() - var plugin = SandboxPlugin( - name: "Needs Secrets", - description: "Declares an API key it never received", - secrets: ["__OSAURUS_TEST_SECRET_THAT_DOES_NOT_EXIST__"] - ) - let message = expectInvalidArgs { - try SandboxPluginRegistration.validateAndStage(&plugin, agentId: agentId.uuidString) + AgentSecretsKeychain._withInMemoryStoreForTesting { + let agentId = UUID() + var plugin = SandboxPlugin( + name: "Needs Secrets", + description: "Declares an API key it never received", + secrets: ["__OSAURUS_TEST_SECRET_THAT_DOES_NOT_EXIST__"] + ) + let message = expectInvalidArgs { + try SandboxPluginRegistration.validateAndStage(&plugin, agentId: agentId.uuidString) + } + #expect(message?.contains("Missing secrets") == true) + #expect(message?.contains("__OSAURUS_TEST_SECRET_THAT_DOES_NOT_EXIST__") == true) } - #expect(message?.contains("Missing secrets") == true) - #expect(message?.contains("__OSAURUS_TEST_SECRET_THAT_DOES_NOT_EXIST__") == true) } @Test func validateAndStage_acceptsSecretsAfterUserStores() throws { - let agentId = UUID() - let key = "OSAURUS_TEST_REGISTRATION_SECRET_\(UUID().uuidString.prefix(8))" - AgentSecretsKeychain.saveSecret("value", id: key, agentId: agentId) - defer { AgentSecretsKeychain.deleteSecret(id: key, agentId: agentId) } + try AgentSecretsKeychain._withInMemoryStoreForTesting { + let agentId = UUID() + let key = "OSAURUS_TEST_REGISTRATION_SECRET_\(UUID().uuidString.prefix(8))" + AgentSecretsKeychain.saveSecret("value", id: key, agentId: agentId) - var plugin = SandboxPlugin( - name: "Has Secret", - description: "Reads an API key set by the user", - secrets: [key] - ) - try SandboxPluginRegistration.validateAndStage( - &plugin, - agentId: agentId.uuidString - ) + var plugin = SandboxPlugin( + name: "Has Secret", + description: "Reads an API key set by the user", + secrets: [key] + ) + try SandboxPluginRegistration.validateAndStage( + &plugin, + agentId: agentId.uuidString + ) + } } // MARK: - SandboxPluginRegisterTool early-failure paths