diff --git a/Sources/Aria/Agent/AgentToolExecution.swift b/Sources/Aria/Agent/AgentToolExecution.swift index 8a8b4ea..ef78ad4 100644 --- a/Sources/Aria/Agent/AgentToolExecution.swift +++ b/Sources/Aria/Agent/AgentToolExecution.swift @@ -105,10 +105,22 @@ extension Agent { ) async -> ToolExecutionResult { let context = ToolContext(runId: UUID()) let started = ContinuousClock.now + // Reconcile the model's arguments with the schema it was given. + // + // Small models quote their numbers. A 0.8B model called a + // weather tool with `{"days": "1", "latitude": "56.35"}` and + // the server refused it — right tool, right intent, right + // values, defeated by quotation marks. The schema says what + // each field is, so this is fixable here rather than by asking + // the model to try again. + let arguments = ToolArgumentCoercion.coerce( + call.arguments, + to: tool.definition.inputSchema + ) do { let output = try await Self.runWithTimeout( timeout: self.config.toolTimeout, - operation: { try await tool.invoke(call.arguments, context) } + operation: { try await tool.invoke(arguments, context) } ) return ToolExecutionResult( output: output, diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift index bdd9160..d0ec9c2 100644 --- a/Sources/Aria/Context/ContextAssembler.swift +++ b/Sources/Aria/Context/ContextAssembler.swift @@ -115,11 +115,28 @@ public struct DefaultContextAssembler: ContextAssembler { state: AgentState, budget: ContextBudget ) async -> AssembledContext { - let selectedTools = await self.selectTools( + let (ranked, rankedNames) = await self.selectTools( from: tools, state: state, budget: budget ) + // Last line of defence: tools must leave room for the message + // they are meant to answer. + // + // `toolTokenLimit` is a *share* and is deliberately breakable — + // sending no usable tool is worse than overspending the share. + // `available` is not. Without a hard bound the two rules + // compose into the worst outcome: a run selected three MCP + // tools worth 4,533 tokens against a 1,198 share and a 2,996 + // budget, history collapsed to nothing to make room, and the + // provider refused the request at 4,727 of 4,096 — so the turn + // failed anyway, having thrown away the conversation first. + let selectedTools = self.withinHardLimit( + ranked, + systemPrompt: systemPrompt, + budget: budget, + state: state + ) // Guidance rides with the tools that survived, so instructions // can never describe a tool the model wasn't given. Consumers @@ -161,6 +178,18 @@ public struct DefaultContextAssembler: ContextAssembler { toolsOffered: tools.count, toolsSelected: selectedTools.count, selectedToolNames: selectedTools.map(\.name), + offeredToolNames: tools.map(\.name), + rankedToolNames: rankedNames, + maxTools: budget.maxTools, + toolTokenLimit: budget.toolTokenLimit, + imageCount: messages.reduce(0) { running, message in + running + message.content.reduce(0) { count, part in + if case .image = part { + return count + 1 + } + return count + } + }, messagesDropped: history.dropped, memoriesDropped: 0, toolResultsTruncated: truncated @@ -327,6 +356,119 @@ public struct DefaultContextAssembler: ContextAssembler { return out.trimmingCharacters(in: .whitespacesAndNewlines) } + /// Drop ranked tools that do not fit the token ceiling. + /// + /// `maxTools` bounds the *count* of tools; this bounds their + /// *cost*. Both are needed, and only the count was enforced on the + /// ranked path — the ceiling was consulted for the "everything + /// already fits" early exit and for filler, then skipped entirely + /// once ranking returned something. Six tools is a fine cap until + /// the six are MCP schemas: a connected weather server put 5,845 + /// tokens into a 4,096-token window and the whole turn was + /// refused, with the budget reporting itself satisfied. + /// + /// Required tools are kept regardless. They are pinned or already + /// invoked this turn, and dropping one breaks the conversation + /// outright, where overshooting only risks a refusal that trimming + /// the rest may still avoid. + /// + /// Ranked tools are taken in order and the first that does not fit + /// stops the loop rather than being skipped over. Skipping to a + /// smaller lower-ranked tool would quietly prefer cheap tools to + /// relevant ones, which is the opposite of what ranking is for. + private static func fitting( + _ selected: [AnyTool], + required: [AnyTool], + ceiling: Int, + cost: (AnyTool) -> Int + ) -> [AnyTool] { + let requiredNames = Set(required.map(\.name)) + var kept = required + var spent = required.reduce(0) { $0 + cost($1) } + // Compaction is deliberately *not* applied here. + // + // `ToolDefinitionCompactor` shrinks a definition safely, and + // shrinking one changes what the assembler counts. It does not + // change what every provider sends: `FoundationModelsProvider` + // builds its typed tools from factory closures supplied at + // construction, which captured the original description and + // cannot see a definition swapped afterwards. + // + // So compacting here made the budget lie. Measured: a turn + // priced at 1,366 tokens was refused at 5,362 — the assembler + // had costed six compacted tools and the provider had sent six + // full ones. An accurate budget with fewer tools beats an + // inaccurate budget with more, and a refused turn sends none at + // all. + // + // The compactor keeps its tests and stays ready for the path + // that can honour it: a provider that builds its tool list from + // `AssembledContext.tools` rather than from closures fixed at + // construction. + for tool in selected where !requiredNames.contains(tool.name) { + let next = cost(tool) + guard spent + next <= ceiling else { + break + } + kept.append(tool) + spent += next + } + + // If the share admitted nothing, admit the best one anyway. + // + // `ceiling` is a *share* of the budget — 40% by default — not + // the window. Honouring it strictly is right when it costs a + // fourth tool and wrong when it costs the only one: a turn with + // no actionable tool cannot do the thing it was asked to do, + // and no amount of budget discipline redeems that. + // + // Observed: a weather question ranked `get_weather_summary` + // second out of six, and every MCP schema on that server runs + // to roughly a thousand tokens — over the 1,080 left after the + // pinned tool. The trim was correct by its own rule and + // returned a request that could only answer from imagination, + // which is exactly what the model then did. + // + // The one admitted here is the top-ranked, not the smallest + // that fits. Preferring a cheap tool to a relevant one inverts + // ranking, and this is a last resort rather than a second + // policy. + if kept.count == required.count, + let best = selected.first(where: { !requiredNames.contains($0.name) }) { + kept.append(best) + } + return kept + } + + /// Drop tools until the request can physically fit, whatever the + /// share allowed. + /// + /// Tools are dropped from the end — lowest-ranked first — and the + /// prompt is recomposed each time, because guidance rides with the + /// tools that survived. The newest message is reserved for: a + /// request that cannot carry the user's turn is not a smaller + /// request, it is a broken one. + private func withinHardLimit( + _ selected: [AnyTool], + systemPrompt: String?, + budget: ContextBudget, + state: AgentState + ) -> [AnyTool] { + let newestTokens = state.messages.last + .map { self.tokenCounter.count(message: $0) } ?? 0 + var kept = selected + while !kept.isEmpty { + let prompt = Self.compose(systemPrompt: systemPrompt, guidanceFrom: kept) + let promptTokens = prompt.map { self.tokenCounter.count(text: $0) } ?? 0 + let toolTokens = kept.reduce(0) { $0 + self.tokenCounter.count(tool: $1.definition) } + if promptTokens + toolTokens + newestTokens <= budget.available { + return kept + } + kept.removeLast() + } + return kept + } + /// Cost of the recalled-memory block, if the caller named it. private func memoryTokens(in messages: [Message]) -> Int { guard let prefix = memoryMessagePrefix else { @@ -337,13 +479,19 @@ public struct DefaultContextAssembler: ContextAssembler { .reduce(0) { $0 + self.tokenCounter.count(message: $1) } } + /// - Returns: The tools to send, and — separately — what ranking + /// chose *before* the token ceiling trimmed it. The difference + /// between those two lists is the whole diagnosis when a tool + /// goes missing: empty `ranked` means the ranker found nothing, + /// while a full `ranked` and a short `selected` means the budget + /// did the cutting. Both look identical from the outside. private func selectTools( from tools: [AnyTool], state: AgentState, budget: ContextBudget - ) async -> [AnyTool] { + ) async -> (selected: [AnyTool], ranked: [String]) { guard !tools.isEmpty else { - return [] + return ([], []) } let maxTools = budget.maxTools ?? Int.max @@ -359,7 +507,7 @@ public struct DefaultContextAssembler: ContextAssembler { // where it was absent. let totalCost = tools.reduce(0) { $0 + self.tokenCounter.count(tool: $1.definition) } if tools.count <= maxTools, totalCost <= ceiling { - return tools + return (tools, tools.map(\.name)) } let required = self.pinnedToolNames @@ -373,16 +521,37 @@ public struct DefaultContextAssembler: ContextAssembler { // only costs tokens. let room = maxTools == Int.max ? candidates.count : max(0, maxTools - requiredTools.count) guard room > 0, !candidates.isEmpty else { - return requiredTools + return (requiredTools, []) } + let query = Self.latestUserText(in: state.messages) + // An image-only turn has no text to rank against. + // + // That is not the same as "nothing matched", and conflating + // them sent zero tools: the selector returns nothing for an + // empty query, `unrankedFillLimit: 0` reads that as a decision, + // and a photo with no caption arrived at the model with no way + // to act on it. The ranker did not decide anything here — it + // was never given a question. + guard !query.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + return ( + Self.fitting( + requiredTools + candidates, + required: requiredTools, + ceiling: ceiling, + cost: { self.tokenCounter.count(tool: $0.definition) } + ), + [] + ) + } let ranked = await self.selector.select( from: candidates.map(\.definition), - query: Self.latestUserText(in: state.messages), + query: query, limit: room ) let byName = Dictionary(candidates.map { ($0.name, $0) }, uniquingKeysWith: { first, _ in first }) var selected = requiredTools + ranked.compactMap { byName[$0.name] } + let rankedNames = selected.map(\.name) // Fill only when ranking had nothing to say. // @@ -406,14 +575,22 @@ public struct DefaultContextAssembler: ContextAssembler { // and a half-empty context window is only a problem if the tool // the user needed is the part that's missing. guard selected.count == requiredTools.count else { - return selected + return ( + Self.fitting( + selected, + required: requiredTools, + ceiling: ceiling, + cost: { self.tokenCounter.count(tool: $0.definition) } + ), + rankedNames + ) } // Nothing scored. How much to read into that depends on the // ranker, so the caller decides — see `unrankedFillLimit`. let fillCeiling = self.unrankedFillLimit ?? maxTools guard fillCeiling > 0 else { - return selected + return (selected, rankedNames) } var chosen = Set(selected.map(\.name)) var spent = selected.reduce(0) { $0 + self.tokenCounter.count(tool: $1.definition) } @@ -433,7 +610,7 @@ public struct DefaultContextAssembler: ContextAssembler { filled += 1 } - return selected + return (selected, rankedNames) } /// Cap each tool result at `limit` tokens, returning the bounded diff --git a/Sources/Aria/Context/ContextBudget.swift b/Sources/Aria/Context/ContextBudget.swift index 127f5a9..8580b08 100644 --- a/Sources/Aria/Context/ContextBudget.swift +++ b/Sources/Aria/Context/ContextBudget.swift @@ -152,6 +152,11 @@ public struct ContextAllocation: Sendable, Equatable, Codable { toolsOffered: Int = 0, toolsSelected: Int = 0, selectedToolNames: [String] = [], + offeredToolNames: [String] = [], + rankedToolNames: [String] = [], + maxTools: Int? = nil, + toolTokenLimit: Int = 0, + imageCount: Int = 0, messagesDropped: Int = 0, memoriesDropped: Int = 0, toolResultsTruncated: Int = 0 @@ -164,6 +169,11 @@ public struct ContextAllocation: Sendable, Equatable, Codable { self.toolsOffered = toolsOffered self.toolsSelected = toolsSelected self.selectedToolNames = selectedToolNames + self.offeredToolNames = offeredToolNames + self.rankedToolNames = rankedToolNames + self.maxTools = maxTools + self.toolTokenLimit = toolTokenLimit + self.imageCount = imageCount self.messagesDropped = messagesDropped self.memoriesDropped = memoriesDropped self.toolResultsTruncated = toolResultsTruncated @@ -192,6 +202,37 @@ public struct ContextAllocation: Sendable, Equatable, Codable { /// is "was the tool I expected even offered?", and only the list /// answers it. public let selectedToolNames: [String] + /// Every tool the selector could have chosen from. + /// + /// `toolsOffered` was a count, which is enough to notice a surprise + /// and never enough to explain one. Asked for a weather summary + /// with a weather server connected, a run sent exactly one tool + /// (`load_skill`) out of twenty-eight — and the count could not + /// distinguish "the ranker buried the weather tools" from "the + /// weather tools were never candidates", which have nothing in + /// common except the symptom. + public let offeredToolNames: [String] + /// What ranking chose, *before* the token ceiling trimmed it. + /// + /// The distinction is the whole diagnosis. If this is empty the + /// ranker found nothing and the query or the corpus is the + /// problem; if it is full and `selectedToolNames` is short, the + /// budget did the cutting. Both end with a tool missing from the + /// request and there is otherwise no way to tell them apart. + public let rankedToolNames: [String] + /// Count cap in force for this turn. + public let maxTools: Int? + /// Token ceiling in force for this turn. + public let toolTokenLimit: Int + /// Images that survived assembly into the request. + /// + /// Asked "What's the tracking number" over a receipt photo, a model + /// answered from the caption alone and invented coordinates. The + /// diagnostic could not say whether the image had reached it — the + /// allocation described tokens and tools and was silent on the one + /// thing that mattered. Zero here with an attachment on screen + /// localises the loss to assembly; non-zero moves it downstream. + public let imageCount: Int /// History messages dropped to fit. public let messagesDropped: Int diff --git a/Sources/Aria/Context/TokenCounter.swift b/Sources/Aria/Context/TokenCounter.swift index f1e0368..36d0cbf 100644 --- a/Sources/Aria/Context/TokenCounter.swift +++ b/Sources/Aria/Context/TokenCounter.swift @@ -41,8 +41,30 @@ extension TokenCounter { 4 } + /// What one image costs in the model's context. + /// + /// Vision models price an image by tile count — a 1024×1024 image + /// runs to roughly a thousand tokens on the Qwen-VL family — and + /// nothing about that is derivable from the bytes we hold. So this + /// is a constant, and a deliberately large one: an image counted at + /// zero is invisible to every budget decision downstream, which is + /// what happened before this existed. History windowing never + /// dropped an image message because it looked free, and the request + /// overflowed with the budget reporting itself satisfied. + /// + /// Over-charging costs a few tools. Under-charging costs the turn. + public var tokensPerImage: Int { + 1024 + } + public func count(message: Message) -> Int { var total = self.count(text: message.textContent) + total += message.content.reduce(0) { running, part in + if case .image = part { + return running + self.tokensPerImage + } + return running + } for call in message.toolCalls { total += self.count(text: call.name) if let data = try? call.arguments.canonicalData() { diff --git a/Sources/Aria/Context/ToolArgumentCoercion.swift b/Sources/Aria/Context/ToolArgumentCoercion.swift new file mode 100644 index 0000000..e9fef4f --- /dev/null +++ b/Sources/Aria/Context/ToolArgumentCoercion.swift @@ -0,0 +1,147 @@ +import Foundation + +// MARK: - ToolArgumentCoercion + +/// Reconciles a model's tool arguments with the schema it was given. +/// +/// Small models quote their numbers. Asked for a weather summary, a +/// 0.8B model emitted `{"days": "1", "latitude": "56.35"}` — every +/// value a string — and the server rejected the call outright: +/// *"Invalid days: must be a finite number, received string"*. The +/// model had the right tool, the right intent and the right values, and +/// the turn failed on quotation marks. +/// +/// The schema says what each field is, so the host can fix that without +/// asking the model to try again. This is the same principle the rest +/// of the context layer runs on: the model decides *what*, the host is +/// responsible for the shape. +/// +/// **Only lossless conversions.** `"1"` becomes `1`; `"abc"` stays +/// `"abc"` and the server returns its own error, which is more useful +/// than a coercion that invented a number. Nothing is dropped, nothing +/// is added, and a value that already matches its declared type is +/// untouched. +public enum ToolArgumentCoercion { + // MARK: Public + + /// Coerce `arguments` to the types `schema` declares. + public static func coerce(_ arguments: JSONValue, to schema: JSONSchema) -> JSONValue { + switch schema { + case let .object(properties, _, _, _): + guard case let .object(values) = arguments else { + return arguments + } + var out: [String: JSONValue] = [:] + for (key, value) in values { + // A field the schema does not describe passes through + // untouched: guessing at its type would be inventing a + // contract that does not exist. + guard let propertySchema = properties[key] else { + out[key] = value + continue + } + out[key] = self.coerce(value, to: propertySchema) + } + return .object(out) + + case let .array(items, _): + guard case let .array(values) = arguments else { + return arguments + } + return .array(values.map { self.coerce($0, to: items) }) + + case .integer: + return self.asInteger(arguments) ?? arguments + + case .number: + return self.asNumber(arguments) ?? arguments + + case .boolean: + return self.asBoolean(arguments) ?? arguments + + case .string: + return self.asString(arguments) ?? arguments + + case let .oneOf(options), let .anyOf(options): + // Try each branch and keep the first that changes + // something. A union whose branches disagree is ambiguous, + // and the original is the safer answer. + for option in options { + let coerced = self.coerce(arguments, to: option) + if coerced != arguments { + return coerced + } + } + return arguments + + case .allOf, .null: + return arguments + } + } + + // MARK: Private + + private static func asInteger(_ value: JSONValue) -> JSONValue? { + switch value { + case .integer: + return nil + case let .string(text): + return Int64(text.trimmingCharacters(in: .whitespaces)).map { .integer($0) } + case let .number(double): + // Only when it is exactly an integer. Silently truncating + // 1.7 to 1 would change what the model asked for. + guard double.rounded() == double, double.magnitude < 9.2e18 else { + return nil + } + return .integer(Int64(double)) + default: + return nil + } + } + + private static func asNumber(_ value: JSONValue) -> JSONValue? { + switch value { + case .number: + nil + case let .string(text): + Double(text.trimmingCharacters(in: .whitespaces)).map { .number($0) } + case let .integer(int): + .number(Double(int)) + default: + nil + } + } + + /// Only the spellings JSON itself uses, plus the two a model + /// reliably means. "yes"/"no" are deliberately absent: they are a + /// guess about intent rather than a reading of the value. + private static func asBoolean(_ value: JSONValue) -> JSONValue? { + guard case let .string(text) = value else { + return nil + } + switch text.trimmingCharacters(in: .whitespaces).lowercased() { + case "true": return .bool(true) + case "false": return .bool(false) + default: return nil + } + } + + /// The mirror case: a model that sends a bare number where the + /// schema wants a string. Less common, equally cheap to fix. + private static func asString(_ value: JSONValue) -> JSONValue? { + switch value { + case let .integer(int): + .string(String(int)) + case let .number(double): + .string( + double.rounded() == double + ? String(Int64(double)) + : String(double) + ) + case let .bool(flag): + .string(flag ? "true" : "false") + default: + nil + } + } +} diff --git a/Sources/Aria/Context/ToolDefinitionCompactor.swift b/Sources/Aria/Context/ToolDefinitionCompactor.swift new file mode 100644 index 0000000..2c43ac9 --- /dev/null +++ b/Sources/Aria/Context/ToolDefinitionCompactor.swift @@ -0,0 +1,180 @@ +import Foundation + +// MARK: - ToolCompaction + +/// How much of a tool definition to give up when it will not fit. +/// +/// Ordered by what it costs the model. Each level keeps every *callable* +/// property of the schema — names, types, and the required list are +/// never touched — so a call the model could have made before it stays +/// legal after. What goes is guidance: the prose that helps it choose +/// well, then the optional surface it could have used. +/// +/// That ordering is the whole design. A tool the model calls with worse +/// arguments is recoverable; a tool whose arguments no longer validate +/// is not. +public enum ToolCompaction: Int, Sendable, Comparable, CaseIterable { + /// Exactly as the server published it. + case none = 0 + /// Drop per-property descriptions. The parameter names survive, and + /// on a well-named schema (`location`, `days`, `units`) they carry + /// most of what the description said anyway. + case propertyDescriptions = 1 + /// Also shorten the tool's own description to its first sentence. + /// + /// Tool descriptions are written for a catalogue page and are + /// front-loaded by convention: the opening sentence says what it + /// does, the rest is caveats and examples. Selection has already + /// happened by the time the model reads this, so the sentence that + /// distinguishes the tool matters more than the paragraph. + case shortDescription = 2 + /// Also drop optional properties, keeping the required ones. + /// + /// The model loses the ability to pass them — it does not lose the + /// ability to call the tool. Last because it is the only level that + /// removes capability rather than guidance. + case requiredOnly = 3 + + // MARK: Public + + public static func < (lhs: ToolCompaction, rhs: ToolCompaction) -> Bool { + lhs.rawValue < rhs.rawValue + } +} + +// MARK: - ToolDefinitionCompactor + +/// Shrinks tool definitions that are too large to send whole. +/// +/// A single MCP tool schema is routinely a thousand tokens — a full +/// paragraph of description plus a dozen documented parameters. On a +/// 4,096-token window that is a quarter of everything the model will +/// ever see, for one tool it may not even call. Six of them put 5,845 +/// tokens into that window and the request was refused outright. +/// +/// Budgeting cannot solve this. A share that fits six such tools does +/// not exist, so the choice was between refusing to send them and +/// sending one at a time; both are bad, and neither addresses the fact +/// that most of those tokens are prose the model does not need in order +/// to make the call. +/// +/// This is not truncation. Truncating a schema produces one the model +/// cannot satisfy; this removes only what is safe to remove, and stops +/// at the first level that fits. +public enum ToolDefinitionCompactor { + // MARK: Public + + /// Compact `definition` to the lightest level that fits `limit`. + /// + /// Returns the original when it already fits, and the most compact + /// form when nothing does — an oversized tool still goes, because + /// a tool the model cannot call is better than a turn that cannot + /// act. + public static func compact( + _ definition: ToolDefinition, + toFit limit: Int, + counter: any TokenCounter + ) -> (definition: ToolDefinition, applied: ToolCompaction) { + for level in ToolCompaction.allCases { + let candidate = self.apply(level, to: definition) + if counter.count(tool: candidate) <= limit { + return (candidate, level) + } + } + return (self.apply(.requiredOnly, to: definition), .requiredOnly) + } + + /// Apply one level, unconditionally. + public static func apply( + _ level: ToolCompaction, + to definition: ToolDefinition + ) -> ToolDefinition { + guard level != .none else { + return definition + } + return ToolDefinition( + name: definition.name, + description: level >= .shortDescription + ? self.firstSentence(of: definition.description) + : definition.description, + inputSchema: self.compact(schema: definition.inputSchema, level: level) + ) + } + + // MARK: Private + + /// Long enough for a real opening sentence, short enough that a + /// description written as one long run-on still shrinks. + private static let sentenceCap = 200 + + private static func compact(schema: JSONSchema, level: ToolCompaction) -> JSONSchema { + switch schema { + case let .object(properties, required, description, additionalProperties): + let requiredSet = Set(required) + var kept: [String: JSONSchema] = [:] + for (name, value) in properties { + // Required properties are never dropped. Removing one + // makes every call invalid, which is the one outcome + // this must not produce. + if level >= .requiredOnly, !requiredSet.contains(name) { + continue + } + kept[name] = self.compact(schema: value, level: level) + } + return .object( + properties: kept, + required: required, + description: level >= .propertyDescriptions ? nil : description, + additionalProperties: additionalProperties + ) + case let .string(description, enumValues): + // Enum values survive at every level: they are not + // documentation, they are the set of legal inputs, and a + // model guessing outside it produces a call that fails. + return .string( + description: level >= .propertyDescriptions ? nil : description, + enumValues: enumValues + ) + case let .number(description): + return .number(description: level >= .propertyDescriptions ? nil : description) + case let .integer(description): + return .integer(description: level >= .propertyDescriptions ? nil : description) + case let .boolean(description): + return .boolean(description: level >= .propertyDescriptions ? nil : description) + case let .array(items, description): + return .array( + items: self.compact(schema: items, level: level), + description: level >= .propertyDescriptions ? nil : description + ) + case let .oneOf(options): + return .oneOf(options.map { self.compact(schema: $0, level: level) }) + case let .anyOf(options): + return .anyOf(options.map { self.compact(schema: $0, level: level) }) + case let .allOf(options): + return .allOf(options.map { self.compact(schema: $0, level: level) }) + case .null: + return schema + } + } + + /// First sentence, or a hard cap when the text has no sentence + /// break — some servers write a single unpunctuated paragraph. + private static func firstSentence(of text: String) -> String { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + return trimmed + } + if let end = trimmed.firstIndex(where: { $0 == "." || $0 == "\n" }) { + let sentence = String(trimmed[trimmed.startIndex...end]) + .trimmingCharacters(in: .whitespacesAndNewlines) + // A three-word opener is a heading, not a description. + if sentence.count >= 24 { + return sentence + } + } + guard trimmed.count > Self.sentenceCap else { + return trimmed + } + return String(trimmed.prefix(Self.sentenceCap)) + "…" + } +} diff --git a/Sources/Aria/Providers/Tool.swift b/Sources/Aria/Providers/Tool.swift index 31b355f..0bd9dd3 100644 --- a/Sources/Aria/Providers/Tool.swift +++ b/Sources/Aria/Providers/Tool.swift @@ -138,6 +138,17 @@ public struct AnyTool: Sendable { public let definition: ToolDefinition public let invoke: @Sendable (JSONValue, ToolContext) async throws -> JSONValue + + /// The same tool, advertised differently. + /// + /// Used to send a compacted schema while keeping the original + /// invocation. What the model is *told* about a tool and what + /// happens when it calls one are separate concerns, and this is the + /// seam between them — the closure is untouched, so a compacted + /// definition cannot change behaviour, only description. + public func replacingDefinition(_ definition: ToolDefinition) -> AnyTool { + AnyTool(definition: definition, invoke: self.invoke) + } } extension AnyTool { diff --git a/Sources/AriaApple/Providers/FoundationModelsProvider.swift b/Sources/AriaApple/Providers/FoundationModelsProvider.swift index ec3fa73..8b2592c 100644 --- a/Sources/AriaApple/Providers/FoundationModelsProvider.swift +++ b/Sources/AriaApple/Providers/FoundationModelsProvider.swift @@ -276,6 +276,16 @@ if executableTools.isEmpty { fmTools = honourSelection ? [] : allTools } else { + // Name-matched only. + // + // Note what this *cannot* do: adopt a description the + // assembler rewrote. These instances come from factory + // closures supplied at construction, which captured + // their description then — so any later change to a + // definition reaches the `AnyTool` and not the tool + // registered here. That is why the assembler does not + // compact definitions; see the note in + // `ContextAssembler.fitting`. let selected = Set(executableTools.map(\.name)) fmTools = allTools.filter { selected.contains($0.name) } } diff --git a/Sources/AriaTesting/ContextAudit.swift b/Sources/AriaTesting/ContextAudit.swift new file mode 100644 index 0000000..d46085a --- /dev/null +++ b/Sources/AriaTesting/ContextAudit.swift @@ -0,0 +1,174 @@ +import Aria +import Foundation + +// MARK: - RequestRendering + +/// Serialises what a provider will actually put in front of the model. +/// +/// The point is to measure the *artifact that ships*, not the objects +/// the budget was computed from. Every budget failure worth debugging +/// so far has been the difference between those two: +/// +/// * a schema counted once as a schema and sent again inside the tool's +/// description, priced at a prose ratio both times; +/// * a definition compacted after the provider had already captured the +/// original, so the estimate shrank and the request did not; +/// * a tool costed and then dropped, or dropped and then sent. +/// +/// A counter that reads `ToolDefinition` cannot see any of those. A +/// counter that reads the rendered text sees all of them, which is why +/// the audit renders first and counts second. +public enum RequestRendering { + /// Approximate what a chat template emits for a request. + /// + /// Deliberately provider-agnostic and deliberately not exact: no + /// two templates agree on separators, and the audit is looking for + /// *large* discrepancies — a 4x under-count, not a 4% one. Being + /// roughly right about the whole is worth more than being exactly + /// right about one provider's punctuation. + public static func render(tools: [AnyTool], messages: [Message]) -> String { + var parts: [String] = [] + for tool in tools { + parts.append(self.render(tool: tool.definition)) + } + for message in messages { + parts.append(self.render(message: message)) + } + return parts.joined(separator: "\n") + } + + public static func render(tool: ToolDefinition) -> String { + var parts = ["{\"name\":\"\(tool.name)\",\"description\":\"\(tool.description)\""] + if let data = try? JSONEncoder().encode(tool.inputSchema), + let schema = String(data: data, encoding: .utf8) { + parts.append(",\"parameters\":\(schema)") + } + if let guidance = tool.promptGuidance { + parts.append(",\"guidance\":\"\(guidance)\"") + } + parts.append("}") + return parts.joined() + } + + public static func render(message: Message) -> String { + var parts = ["<|\(message.role)|>", message.textContent] + for call in message.toolCalls { + parts.append("{\"name\":\"\(call.name)\",\"arguments\":") + if let data = try? call.arguments.canonicalData(), + let json = String(data: data, encoding: .utf8) { + parts.append(json) + } + parts.append("}") + } + return parts.joined() + } +} + +// MARK: - ContextAudit + +/// Checks a token estimate against the request it is supposed to +/// describe. +/// +/// The estimate is the most load-bearing number in the context layer — +/// every trim, every cap and every drop is arithmetic on top of it — and +/// until now nothing verified it. When it was wrong the symptom appeared +/// somewhere else entirely: a provider refusing a turn that the budget +/// had declared comfortable, with the budget still reporting itself +/// satisfied afterwards. +/// +/// Direction matters. Over-estimating wastes budget and sends fewer +/// tools; under-estimating overflows the window and loses the whole +/// turn. So the assertions are one-sided: the estimate may be +/// conservative, and may not be optimistic. +public struct ContextAudit: Sendable { + // MARK: Lifecycle + + /// - Parameter reference: Ground truth over the rendered request. + /// Supply a real tokenizer where one is available; the default is + /// a character heuristic tuned for JSON-dense text, which is + /// enough to catch order-of-magnitude errors and not enough to + /// catch small ones. + public init( + counter: any TokenCounter = HeuristicTokenCounter(), + reference: @escaping @Sendable (String) -> Int = ContextAudit.defaultReference + ) { + self.counter = counter + self.reference = reference + } + + // MARK: Public + + public struct Report: Sendable { + public let estimatedTokens: Int + public let referenceTokens: Int + public let renderedCharacters: Int + + /// How far the estimate falls short. Above 1 means the request + /// is bigger than the budget believed. + public var underestimateFactor: Double { + guard self.estimatedTokens > 0 else { + return self.referenceTokens > 0 ? .infinity : 1 + } + return Double(self.referenceTokens) / Double(self.estimatedTokens) + } + + /// Characters the estimate implicitly priced per token. A value + /// far above the counter's configured ratio means it is reading + /// something other than what ships. + public var impliedCharactersPerToken: Double { + guard self.estimatedTokens > 0 else { + return .infinity + } + return Double(self.renderedCharacters) / Double(self.estimatedTokens) + } + + public func summary() -> String { + String( + format: "estimated %d · reference %d · %.2fx · %.1f chars/token implied", + self.estimatedTokens, + self.referenceTokens, + self.underestimateFactor, + self.impliedCharactersPerToken + ) + } + } + + /// JSON tokenises far denser than prose — quoted keys, braces and + /// nesting each become their own token. Measured against real + /// tokenizers this lands near three characters per token; the audit + /// uses it only to spot gross errors. + public static let defaultReference: @Sendable (String) -> Int = { text in + max(1, Int((Double(text.count) / 3.0).rounded(.up))) + } + + /// Audit a whole assembled request. + public func audit(_ assembled: AssembledContext) -> Report { + let rendered = RequestRendering.render( + tools: assembled.tools, + messages: assembled.messages + ) + let estimated = assembled.tools.reduce(0) { $0 + self.counter.count(tool: $1.definition) } + + assembled.messages.reduce(0) { $0 + self.counter.count(message: $1) } + return Report( + estimatedTokens: estimated, + referenceTokens: self.reference(rendered), + renderedCharacters: rendered.count + ) + } + + /// Audit one tool in isolation — the granularity that localises a + /// bad estimate to the definition that caused it. + public func audit(tool: ToolDefinition) -> Report { + let rendered = RequestRendering.render(tool: tool) + return Report( + estimatedTokens: self.counter.count(tool: tool), + referenceTokens: self.reference(rendered), + renderedCharacters: rendered.count + ) + } + + // MARK: Private + + private let counter: any TokenCounter + private let reference: @Sendable (String) -> Int +} diff --git a/Sources/AriaTools/Charts/ChartIntent.swift b/Sources/AriaTools/Charts/ChartIntent.swift new file mode 100644 index 0000000..04024a2 --- /dev/null +++ b/Sources/AriaTools/Charts/ChartIntent.swift @@ -0,0 +1,247 @@ +import Aria +import Foundation + +// MARK: - ChartKind + +/// The visualisations a model may ask for. +/// +/// An enum rather than a string, so an invalid chart type is +/// unrepresentable rather than a runtime failure. This is the whole +/// reason the model is not asked to author a Vega-Lite spec: given free +/// text it will eventually emit `"linechart"`, `"line_chart"`, or a +/// spec with a subtly wrong encoding block, and each of those is a +/// broken chart in front of a user. +public enum ChartKind: String, Codable, Sendable, CaseIterable { + case line + case bar + case area + case scatter + case pie + + // MARK: Internal + + /// Vega-Lite mark name. + var mark: String { + switch self { + case .line: "line" + case .bar: "bar" + case .area: "area" + case .scatter: "point" + case .pie: "arc" + } + } +} + +// MARK: - ChartIntent + +/// What the model emits when it decides something should be a chart. +/// +/// Deliberately tiny — roughly thirty tokens of output against the +/// several hundred a Vega-Lite spec would cost, which matters against a +/// 4,096-token window that also has to hold the data being charted. +/// +/// Note what is *absent*: the data. The model names a tool result it +/// has already seen and the fields within it; the host supplies the +/// rows. A model that retypes numbers eventually invents them, and an +/// invented number wearing the authority of a chart is worse than an +/// invented sentence. Keeping values out of the model's output makes +/// charts structurally grounded rather than grounded by inspection. +public struct ChartIntent: Codable, Hashable, Sendable { + // MARK: Lifecycle + + public init( + kind: ChartKind, + dataRef: String, + x: String, + y: String, + series: String? = nil, + title: String? = nil + ) { + self.kind = kind + self.dataRef = dataRef + self.x = x + self.y = y + self.series = series + self.title = title + } + + // MARK: Public + + public let kind: ChartKind + /// Identifier of a tool result the host is holding. + public let dataRef: String + /// Field name for the horizontal axis (or category, for pie). + public let x: String + /// Field name for the vertical axis (or magnitude, for pie). + public let y: String + /// Optional field to split the data into multiple series. + public let series: String? + public let title: String? +} + +// MARK: - ChartCompilerError + +public enum ChartCompilerError: Error, Equatable, Sendable { + /// The referenced tool result is not held by the host. + case unresolvedDataRef(String) + /// The rows contained no usable records. + case emptyData + /// A named field is absent from the data. + /// + /// Carries what *was* available, because the model naming a field + /// that does not exist is the most likely failure here and the + /// recovery is to tell it the real names. + case unknownField(name: String, available: [String]) +} + +// MARK: - ChartDataType + +/// Vega-Lite field type, inferred from the data rather than declared by +/// the model — the host has the values, so it can tell. +enum ChartDataType: String { + case quantitative + case temporal + case nominal +} + +// MARK: - VegaLiteCompiler + +/// Turns a `ChartIntent` plus real rows into a Vega-Lite spec. +/// +/// Improving charts — nicer axes, formatting, colour, legends — happens +/// here, and needs no prompt change and no model change. That +/// separation is the payoff for keeping the model's output small. +public enum VegaLiteCompiler { + // MARK: Public + + /// Vega-Lite schema this compiler emits against. + public static let schema = "https://vega.github.io/schema/vega-lite/v5.json" + + /// - Parameters: + /// - intent: What the model asked for. + /// - rows: The tool result's records, supplied by the host. + /// - Returns: A Vega-Lite spec as `JSONValue`, ready to serialise. + public static func compile( + _ intent: ChartIntent, + rows: [[String: JSONValue]] + ) throws -> JSONValue { + guard !rows.isEmpty else { + throw ChartCompilerError.emptyData + } + let available = Self.fieldNames(in: rows) + for field in [intent.x, intent.y] + (intent.series.map { [$0] } ?? []) { + guard available.contains(field) else { + throw ChartCompilerError.unknownField(name: field, available: available.sorted()) + } + } + + var encoding: [String: JSONValue] = [ + "x": Self.field(intent.x, type: Self.inferType(of: intent.x, in: rows)), + "y": Self.field(intent.y, type: Self.inferType(of: intent.y, in: rows)), + ] + if let series = intent.series { + encoding["color"] = Self.field(series, type: .nominal) + } + + // Pie has no axes: the fields mean angle and category instead. + if intent.kind == .pie { + encoding = [ + "theta": Self.field(intent.y, type: .quantitative), + "color": Self.field(intent.x, type: .nominal), + ] + } + + // Tooltips are on by default. An interactive chart the user + // cannot interrogate is a picture, and `tooltip: true` is the + // whole cost of not shipping one. + encoding["tooltip"] = .array( + ([intent.x, intent.y] + (intent.series.map { [$0] } ?? [])).map { name in + .object([ + "field": .string(name), + "type": .string(Self.inferType(of: name, in: rows).rawValue), + ]) + } + ) + + var spec: [String: JSONValue] = [ + "$schema": .string(Self.schema), + "data": .object(["values": .array(rows.map { .object($0) })]), + "mark": .object([ + "type": .string(intent.kind.mark), + "tooltip": .bool(true), + "point": .bool(intent.kind == .line), + ]), + "encoding": .object(encoding), + "width": .string("container"), + "autosize": .object(["type": .string("fit"), "contains": .string("padding")]), + ] + if let title = intent.title, !title.isEmpty { + spec["title"] = .string(title) + } + return .object(spec) + } + + // MARK: Private + + private static func field(_ name: String, type: ChartDataType) -> JSONValue { + .object(["field": .string(name), "type": .string(type.rawValue)]) + } + + private static func fieldNames(in rows: [[String: JSONValue]]) -> Set { + rows.reduce(into: Set()) { names, row in + names.formUnion(row.keys) + } + } + + /// Infer a field's type from its values. + /// + /// Order matters: numbers first, then anything date-shaped, then + /// nominal. A field is only quantitative if *every* present value + /// is numeric — one stray `"n/a"` makes the axis meaningless, and + /// nominal at least renders something honest. + private static func inferType( + of field: String, + in rows: [[String: JSONValue]] + ) -> ChartDataType { + let values = rows.compactMap { $0[field] }.filter { value in + if case .null = value { + return false + } + return true + } + guard !values.isEmpty else { + return .nominal + } + let allNumeric = values.allSatisfy { value in + switch value { + case .integer, .number: true + default: false + } + } + if allNumeric { + return .quantitative + } + let allDateLike = values.allSatisfy { value in + guard case let .string(text) = value else { + return false + } + return Self.looksLikeDate(text) + } + return allDateLike ? .temporal : .nominal + } + + /// ISO-8601-ish detection, deliberately narrow. A false positive + /// here turns a category axis into a broken time axis, which is a + /// worse outcome than a date rendered as a category. + private static func looksLikeDate(_ text: String) -> Bool { + guard text.count >= 8, let first = text.first, first.isNumber else { + return false + } + let separators = text.filter { $0 == "-" || $0 == "/" }.count + guard separators >= 2 else { + return false + } + let prefix = text.prefix(4) + return prefix.allSatisfy(\.isNumber) + } +} diff --git a/Sources/AriaTools/Charts/TabularExtraction.swift b/Sources/AriaTools/Charts/TabularExtraction.swift new file mode 100644 index 0000000..5ad113b --- /dev/null +++ b/Sources/AriaTools/Charts/TabularExtraction.swift @@ -0,0 +1,149 @@ +import Aria +import Foundation + +// MARK: - TabularExtraction + +/// Finds the row set inside an arbitrary tool result. +/// +/// Tools return whatever shape their author chose: a bare array, an +/// envelope like `{"items": [...]}`, something nested two levels down, +/// or a JSON string that has to be decoded first. A chart needs +/// `[[String: JSONValue]]`, and the model must not be the thing that +/// reshapes it — asking it to transcribe rows is asking it to invent +/// them. +/// +/// So the host looks. The rules are deliberately dull and predictable, +/// because a surprising rule here produces a chart of the wrong data, +/// which is worse than no chart. +public enum TabularExtraction { + // MARK: Public + + /// Where the rows came from, and what was left out. + public struct Extraction: Equatable, Sendable { + public let rows: [[String: JSONValue]] + /// Key path the rows were found at; empty when the result was + /// itself an array. Useful for telling a user *what* was + /// charted when a result had several candidate arrays. + public let path: [String] + /// Original row count when rows were dropped, else `nil`. + public let truncatedFrom: Int? + } + + public static let defaultMaximumRows = 2000 + + /// Extract the most plausible row set, or `nil` if there isn't one. + /// + /// - Parameter maximumRows: Charts of ten thousand points are + /// unreadable and slow to draw. Excess rows are dropped from the + /// end, and `extractDetailed` reports that it happened. + public static func rows( + from value: JSONValue, + maximumRows: Int = TabularExtraction.defaultMaximumRows + ) -> [[String: JSONValue]]? { + self.extractDetailed(from: value, maximumRows: maximumRows)?.rows + } + + /// As `rows(from:)`, but reports where the rows were found and + /// whether any were dropped — the caller usually wants to say so. + public static func extractDetailed( + from value: JSONValue, + maximumRows: Int = TabularExtraction.defaultMaximumRows + ) -> Extraction? { + guard let found = self.search(value, depth: 0, path: []) else { + return nil + } + let truncated = found.rows.count > maximumRows + return Extraction( + rows: truncated ? Array(found.rows.prefix(maximumRows)) : found.rows, + path: found.path, + truncatedFrom: truncated ? found.rows.count : nil + ) + } + + // MARK: Private + + private struct Found { + let rows: [[String: JSONValue]] + let path: [String] + } + + /// Keys checked first when several arrays are present. A result + /// with both `items` and `warnings` should chart the items. + private static let preferredKeys = [ + "rows", "items", "data", "results", "records", "series", "values", "entries", "points", + ] + + /// How deep to look. Two levels covers `{"data": {"items": [...]}}` + /// without wandering into unrelated nested structures. + private static let maximumDepth = 3 + + private static func search(_ value: JSONValue, depth: Int, path: [String]) -> Found? { + guard depth <= self.maximumDepth else { + return nil + } + switch value { + case let .array(items): + return self.rowsFromArray(items).map { Found(rows: $0, path: path) } + + case let .string(text): + // A tool that returns JSON as a string is common enough to + // handle, but only at the top: decoding strings found deep + // inside a payload would be guessing. + guard depth == 0, + let data = text.data(using: .utf8), + let decoded = try? JSONDecoder().decode(JSONValue.self, from: data) else { + return nil + } + return self.search(decoded, depth: depth + 1, path: path) + + case let .object(fields): + // Preferred keys first, in order, so the choice is stable + // rather than dependent on dictionary ordering. + for key in Self.preferredKeys { + guard let child = fields[key] else { + continue + } + if case let .array(items) = child, let rows = self.rowsFromArray(items) { + return Found(rows: rows, path: path + [key]) + } + } + // Then any array-valued key, sorted for determinism. + for key in fields.keys.sorted() { + if case let .array(items) = fields[key], let rows = self.rowsFromArray(items) { + return Found(rows: rows, path: path + [key]) + } + } + // Then recurse into objects, also sorted. + for key in fields.keys.sorted() { + guard let child = fields[key], case .object = child else { + continue + } + if let found = self.search(child, depth: depth + 1, path: path + [key]) { + return found + } + } + return nil + + default: + return nil + } + } + + /// An array is a row set only if it is non-empty and every element + /// is an object. A mixed array is not a table, and charting the + /// object-shaped subset of one would silently drop data. + private static func rowsFromArray(_ items: [JSONValue]) -> [[String: JSONValue]]? { + guard !items.isEmpty else { + return nil + } + var rows: [[String: JSONValue]] = [] + rows.reserveCapacity(items.count) + for item in items { + guard case let .object(fields) = item else { + return nil + } + rows.append(fields) + } + return rows + } +} diff --git a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift index f48e501..6786c6c 100644 --- a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift +++ b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift @@ -8,10 +8,16 @@ import MCP // MARK: - MCPClient /// MCP (Model Context Protocol) client bound to a single server -/// endpoint. One instance is constructed per call site; the workflow -/// compiler instantiates it per-step rather than pooling, which keeps -/// the engine stateless and matches MCP's session-per-call model for -/// read-mostly tool invocations. +/// endpoint. Instances stay cheap to construct — the value is a handle, +/// and the underlying connection lives in `MCPSessionPool` keyed by +/// endpoint + credential, so constructing one per call site costs +/// nothing and reuses the live session. +/// +/// It previously connected and ran a full `initialize` handshake per +/// call, on the stated grounds that this "matches MCP's session-per-call +/// model". MCP has no such model — it is a stateful session protocol, +/// and Streamable HTTP carries an `Mcp-Session-Id` across the session. +/// See `MCPSessionPool` for what that cost. /// /// Backed by the official `modelcontextprotocol/swift-sdk` /// (`HTTPClientTransport` over Streamable HTTP). We hand-rolled this @@ -33,20 +39,67 @@ import MCP public struct MCPClient: Sendable { // MARK: Lifecycle + /// - Parameter clientName: How this app introduces itself in the + /// MCP handshake. Defaults to the host bundle's name rather than a + /// literal, which is how every Niora request came to identify + /// itself as "Avyra" — the default was hardcoded in a shared + /// package and no call site overrode it. public init( serverURL: URL, credential: MCPCredential? = nil, - clientName: String = "Avyra", - clientVersion: String = "1.0" + clientName: String = MCPClient.defaultClientName, + clientVersion: String = "1.0", + streaming: Bool = false, + onToolsChanged: (@Sendable () -> Void)? = nil + ) { + self.serverURL = serverURL + self.credential = credential + self.clientName = clientName + self.clientVersion = clientVersion + self.transportFactory = nil + self.streaming = streaming + self.onToolsChanged = onToolsChanged + } + + /// Test seam: supply the transport instead of building an + /// `HTTPClientTransport` from the endpoint. + /// + /// Without this the mapping code below — pagination draining, + /// content-block conversion, prompt argument flattening — can only + /// be exercised against a live third-party server, which is to say + /// never in CI. + init( + serverURL: URL, + credential: MCPCredential?, + clientName: String, + clientVersion: String, + transportFactory: @escaping @Sendable () -> any Transport, + streaming: Bool = false, + onToolsChanged: (@Sendable () -> Void)? = nil ) { self.serverURL = serverURL self.credential = credential self.clientName = clientName self.clientVersion = clientVersion + self.transportFactory = transportFactory + self.streaming = streaming + self.onToolsChanged = onToolsChanged } // MARK: Public + /// The host app's name, or a neutral fallback off-app (tests, CLI). + public static let defaultClientName: String = { + let keys = ["CFBundleDisplayName", "CFBundleName"] + for key in keys { + if let name = Bundle.main.object(forInfoDictionaryKey: key) as? String, + !name.isEmpty { + return name + } + } + return "aria-mcp-client" + }() + /// Invoke `name` with the supplied arguments and return the /// concatenated text-content blocks — the canonical textual tool /// output, and the back-compatible shape every existing caller @@ -115,21 +168,109 @@ public struct MCPClient: Sendable { } } + /// Enumerate the server's resources, draining pagination. + /// + /// Distinct from tools: a tool is something the model *calls*, a + /// resource is content the host can read or render. MCP Apps + /// arrive here — a `ui://` resource with a + /// `text/html;profile=mcp-app` MIME type is an interactive surface + /// rather than data. + public func listResources() async throws -> [MCPResourceDescriptor] { + try await self.withConnectedClient { client in + var collected: [MCP.Resource] = [] + var cursor: String? + repeat { + let (resources, next) = try await client.listResources(cursor: cursor) + collected.append(contentsOf: resources) + cursor = next + } while cursor != nil + return collected.map { + MCPResourceDescriptor( + uri: $0.uri, + name: $0.name, + description: $0.description, + mimeType: $0.mimeType, + size: $0.size + ) + } + } + } + + /// Read one resource by URI. + /// + /// Returns every content block the server sends. A resource may be + /// text, base64 `blob`, or several parts — callers wanting the + /// markup of a UI resource want `firstHTMLResource`. + public func readResource(uri: String) async throws -> MCPCallResult { + try await self.withConnectedClient { client in + let contents = try await client.readResource(uri: uri) + return MCPCallResult( + content: contents.map { content in + if let text = content.text { + return .resource(MCPResourceContent( + uri: content.uri, + mimeType: content.mimeType, + text: text, + blob: nil + )) + } + return .resource(MCPResourceContent( + uri: content.uri, + mimeType: content.mimeType, + text: nil, + blob: content.blob + )) + }, + isError: false + ) + } + } + + /// Enumerate the server's prompt templates, draining pagination. + public func listPrompts() async throws -> [MCPPromptDescriptor] { + try await self.withConnectedClient { client in + var collected: [MCP.Prompt] = [] + var cursor: String? + repeat { + let (prompts, next) = try await client.listPrompts(cursor: cursor) + collected.append(contentsOf: prompts) + cursor = next + } while cursor != nil + return collected.map { prompt in + MCPPromptDescriptor( + name: prompt.name, + description: prompt.description, + arguments: (prompt.arguments ?? []).map { + MCPPromptDescriptor.Argument( + name: $0.name, + description: $0.description, + required: $0.required ?? false + ) + } + ) + } + } + } + // MARK: Private private let serverURL: URL private let credential: MCPCredential? private let clientName: String private let clientVersion: String + private let transportFactory: (@Sendable () -> any Transport)? + private let streaming: Bool + private let onToolsChanged: (@Sendable () -> Void)? private static func makeTransport( endpoint: URL, - credential: MCPCredential? + credential: MCPCredential?, + streaming: Bool ) -> MCP.HTTPClientTransport { let authorization = Self.authorizationHeader(for: credential) return MCP.HTTPClientTransport( endpoint: endpoint, - streaming: false, + streaming: streaming, requestModifier: { request in guard let authorization else { return request @@ -262,26 +403,73 @@ public struct MCPClient: Sendable { return MCPError.networkFailure(error.localizedDescription) } - /// Stand up a fresh SDK client + transport, run the - /// initialize handshake, hand the connected client to `body`, then - /// tear everything down — on both the success and failure paths so - /// the transport's resources don't leak. SDK errors are normalised - /// into our `MCPError` so callers see one error taxonomy. + /// Distinguishes credentials without putting secrets in a + /// dictionary key. Collisions only cost a needless reconnect. + private static func fingerprint(_ credential: MCPCredential?) -> Int { + var hasher = Hasher() + switch credential { + case .none: + hasher.combine(0) + case let .bearer(token): + hasher.combine(1) + hasher.combine(token) + case let .basic(username, password): + hasher.combine(2) + hasher.combine(username) + hasher.combine(password) + } + return hasher.finalize() + } + + /// Run `body` against a pooled, connected client. SDK errors are + /// normalised into our `MCPError` so callers see one taxonomy + /// regardless of whether the failure came from our code or the SDK. + /// + /// The session is *not* torn down afterwards — that is the point. + /// `MCPSessionPool` owns its lifetime, reuses it for the next call, + /// and evicts it once idle. private func withConnectedClient( - _ body: (MCP.Client) async throws -> T + _ body: @Sendable @escaping (MCP.Client) async throws -> T ) async throws -> T { - let client = MCP.Client(name: self.clientName, version: self.clientVersion) - let transport = Self.makeTransport( - endpoint: self.serverURL, - credential: self.credential - ) + let endpoint = self.serverURL + let credential = self.credential + let override = self.transportFactory + let streaming = self.streaming + // Built up front with an explicit type: inlining this at the + // call site defeated the type checker outright. + let onConnect: (@Sendable (MCP.Client) async -> Void)? = + if let notify = self.onToolsChanged { + { client in + // The server tells us its tools changed; the app + // re-lists rather than serving a stale cache until + // someone opens Settings and taps refresh. + await client.onNotification(ToolListChangedNotification.self) { _ in + notify() + } + } + } else { + nil + } do { - _ = try await client.connect(transport: transport) - let result = try await body(client) - await client.disconnect() - return result + return try await MCPSessionPool.shared.withClient( + key: MCPSessionKey( + endpoint: endpoint, + credentialFingerprint: Self.fingerprint(credential), + clientName: self.clientName, + clientVersion: self.clientVersion, + streaming: streaming + ), + makeTransport: { + override?() ?? Self.makeTransport( + endpoint: endpoint, + credential: credential, + streaming: streaming + ) + }, + onConnect: onConnect, + body: body + ) } catch { - await client.disconnect() throw Self.mapError(error) } } diff --git a/Sources/WorkflowKit/Engine/MCP/MCPResource.swift b/Sources/WorkflowKit/Engine/MCP/MCPResource.swift new file mode 100644 index 0000000..c319971 --- /dev/null +++ b/Sources/WorkflowKit/Engine/MCP/MCPResource.swift @@ -0,0 +1,102 @@ +import Foundation + +// MARK: - MCPResourceDescriptor + +/// A resource a server advertises, as returned by `resources/list`. +/// +/// Mirrors `MCPToolDescriptor`: a plain value with no SDK types in it, +/// so persistence and UI never import `MCP`. +public struct MCPResourceDescriptor: Codable, Hashable, Sendable, Identifiable { + // MARK: Lifecycle + + public init( + uri: String, + name: String, + description: String? = nil, + mimeType: String? = nil, + size: Int? = nil + ) { + self.uri = uri + self.name = name + self.description = description + self.mimeType = mimeType + self.size = size + } + + // MARK: Public + + public let uri: String + public let name: String + public let description: String? + public let mimeType: String? + public let size: Int? + + public var id: String { + self.uri + } + + /// Whether this resource is an interactive UI surface rather than + /// data to read. + /// + /// [MCP Apps](https://blog.modelcontextprotocol.io/posts/2025-11-21-mcp-apps/) + /// marks these with a `ui://` URI and the + /// `text/html;profile=mcp-app` MIME type. The profile parameter is + /// what distinguishes an interactive app from a server that merely + /// serves HTML, so both halves are checked — a bare `text/html` + /// resource is a document, and rendering it as an app would grant + /// it a message channel it never asked for. + public var isInteractiveUI: Bool { + guard self.uri.hasPrefix("ui://") else { + return false + } + guard let mimeType else { + return false + } + return mimeType.replacingOccurrences(of: " ", with: "") + .lowercased() + .contains("profile=mcp-app") + } +} + +// MARK: - MCPPromptDescriptor + +/// A prompt template a server advertises. +/// +/// Prompts are the surface Discover recipes map onto: a named, described +/// entry point with typed arguments, which is exactly the shape a +/// recipe needs. +public struct MCPPromptDescriptor: Codable, Hashable, Sendable, Identifiable { + // MARK: Lifecycle + + public init(name: String, description: String? = nil, arguments: [Argument] = []) { + self.name = name + self.description = description + self.arguments = arguments + } + + // MARK: Public + + public struct Argument: Codable, Hashable, Sendable { + // MARK: Lifecycle + + public init(name: String, description: String? = nil, required: Bool = false) { + self.name = name + self.description = description + self.required = required + } + + // MARK: Public + + public let name: String + public let description: String? + public let required: Bool + } + + public let name: String + public let description: String? + public let arguments: [Argument] + + public var id: String { + self.name + } +} diff --git a/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift b/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift new file mode 100644 index 0000000..91cb385 --- /dev/null +++ b/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift @@ -0,0 +1,198 @@ +import Foundation +import MCP + +// MARK: - MCPSessionKey + +/// Identity of a poolable session. +/// +/// The credential is part of the key because the `Authorization` header +/// is baked into the transport when it is built — a refreshed OAuth +/// token cannot reuse the old transport, it needs a new one. Keying on +/// a fingerprint rather than the credential itself keeps secrets out of +/// dictionary keys and out of any future debug description. +struct MCPSessionKey: Hashable { + let endpoint: URL + let credentialFingerprint: Int + let clientName: String + let clientVersion: String + /// Streaming and non-streaming transports are different + /// connections, not a setting on one. Sharing a session between + /// them would hand a caller expecting server-initiated + /// notifications a transport that cannot deliver them. + var streaming = false +} + +// MARK: - MCPSessionPool + +/// Keeps one connected MCP client per server, instead of standing up a +/// fresh one for every call. +/// +/// The client this replaced ran a full `initialize` handshake per +/// request and tore the transport down afterwards, with a comment +/// claiming this "matches MCP's session-per-call model". MCP has no +/// such model: it is a stateful session protocol — `initialize`, then +/// operate, then shut down — and Streamable HTTP carries an +/// `Mcp-Session-Id` across the session. Reconnecting per call meant: +/// +/// * a handshake on every tool call, so a turn calling three tools +/// paid three; +/// * any server holding per-session state (a cache, a cursor, an auth +/// context) seeing each call as a brand-new session, and +/// misbehaving quietly rather than erroring. +/// +/// Sessions are evicted after `idleTimeout` so a server that is used +/// once in a conversation does not hold a connection open for the life +/// of the process. +actor MCPSessionPool { + // MARK: Internal + + static let shared = MCPSessionPool() + + /// Run `body` against a connected client for `key`, reusing an + /// existing session when one is live. + /// + /// Retries **once** on a connection-shaped failure. A pooled + /// connection can be closed by the server, a proxy, or the OS + /// between calls, and the caller cannot distinguish "the session + /// went stale" from "this request is bad" — so the pool absorbs it. + /// The retry is deliberately not a general one: a `serverError` is + /// the server answering, and answering twice would be wrong. + /// `makeTransport` is a factory rather than a value because a + /// retry needs a *fresh* transport — a disconnected one cannot be + /// reconnected. Taking `any Transport` also lets tests drive the + /// pool without a live server. + func withClient( + key: MCPSessionKey, + makeTransport: @Sendable () -> any Transport, + onConnect: (@Sendable (Client) async -> Void)? = nil, + body: @Sendable (Client) async throws -> T + ) async throws -> T { + self.evictExpired() + do { + let client = try await self.connectedClient( + for: key, + makeTransport: makeTransport, + onConnect: onConnect + ) + let value = try await body(client) + self.sessions[key]?.touch() + return value + } catch { + guard Self.isConnectionFailure(error) else { + throw error + } + await self.drop(key) + let client = try await self.connectedClient( + for: key, + makeTransport: makeTransport, + onConnect: onConnect + ) + let value = try await body(client) + self.sessions[key]?.touch() + return value + } + } + + /// Close every session. Call when the host app backgrounds — an + /// open connection across a suspension is a connection that will be + /// found dead on resume, and paying the reconnect at that point is + /// cheaper than discovering it mid-turn. + func closeAll() async { + let live = self.sessions.values + self.sessions.removeAll() + for session in live { + await session.client.disconnect() + } + } + + /// Drop one session — used after an auth change or a failure. + func drop(_ key: MCPSessionKey) async { + guard let session = self.sessions.removeValue(forKey: key) else { + return + } + await session.client.disconnect() + } + + // MARK: Private + + private final class Session { + // MARK: Lifecycle + + init(client: Client) { + self.client = client + self.lastUsed = Date() + } + + // MARK: Internal + + let client: Client + private(set) var lastUsed: Date + + func touch() { + self.lastUsed = Date() + } + } + + /// Long enough to cover a conversational turn and the follow-up + /// that usually accompanies it; short enough that an app left open + /// on a screen is not holding sockets to every configured server. + private static let idleTimeout: TimeInterval = 180 + + private var sessions: [MCPSessionKey: Session] = [:] + + /// Connection-shaped failures, worth one retry on a fresh session. + /// A `serverError` is excluded on purpose: the server responded, + /// and repeating a request it already answered risks doing a + /// non-idempotent thing twice. + private static func isConnectionFailure(_ error: Error) -> Bool { + if let sdk = error as? MCP.MCPError { + switch sdk { + case .connectionClosed, .transportError: + return true + default: + return false + } + } + if let ours = error as? MCPError, case .networkFailure = ours { + return true + } + return false + } + + private func connectedClient( + for key: MCPSessionKey, + makeTransport: @Sendable () -> any Transport, + onConnect: (@Sendable (Client) async -> Void)? = nil + ) async throws -> Client { + if let existing = self.sessions[key] { + return existing.client + } + let client = Client(name: key.clientName, version: key.clientVersion) + // Handlers are registered *before* connecting: a server may + // emit `tools/list_changed` immediately after initialize, and a + // handler attached afterwards would miss it. + await onConnect?(client) + _ = try await client.connect(transport: makeTransport()) + self.sessions[key] = Session(client: client) + return client + } + + private func evictExpired() { + let cutoff = Date().addingTimeInterval(-Self.idleTimeout) + let stale = self.sessions.filter { $0.value.lastUsed < cutoff } + guard !stale.isEmpty else { + return + } + for key in stale.keys { + self.sessions.removeValue(forKey: key) + } + // Disconnect off the actor: teardown is I/O and nothing else + // needs to wait for it. + let clients = stale.values.map(\.client) + Task { + for client in clients { + await client.disconnect() + } + } + } +} diff --git a/Tests/AriaTests/ContextAuditTests.swift b/Tests/AriaTests/ContextAuditTests.swift new file mode 100644 index 0000000..719e190 --- /dev/null +++ b/Tests/AriaTests/ContextAuditTests.swift @@ -0,0 +1,149 @@ +@testable import Aria +import AriaTesting +import XCTest + +/// Does the token estimate describe the request that actually ships? +/// +/// Nothing checked this until now, and every budget failure worth +/// debugging turned out to be the gap between the two: a turn priced at +/// 1,366 tokens refused at 5,362, with the budget reporting itself +/// satisfied afterwards. +/// +/// The assertions are one-sided on purpose. Over-estimating wastes +/// budget and sends fewer tools; under-estimating overflows the window +/// and loses the whole turn. +final class ContextAuditTests: XCTestCase { + // MARK: - The systemic risk + + /// A realistic MCP schema, costed and rendered. + func testMCPSizedToolIsNotWildlyUnderestimated() { + let report = ContextAudit().audit(tool: Self.mcpTool) + XCTAssertLessThan( + report.underestimateFactor, + 1.5, + "estimate is optimistic: \(report.summary())" + ) + } + + /// The specific error that priced a JSON blob as prose. A tool + /// whose description *contains* its schema — how the MCP bridge + /// presents them to FoundationModels — is the case that broke. + func testSchemaInsideADescriptionIsNotPricedAsProse() { + let inlined = ToolDefinition( + name: "weather_mcp__get_weather_summary", + description: """ + Get a weather summary. + + Input schema: + \(Self.schemaJSON) + """, + inputSchema: Self.mcpTool.inputSchema + ) + let report = ContextAudit().audit(tool: inlined) + XCTAssertLessThan( + report.underestimateFactor, + 1.5, + "description-embedded schema underpriced: \(report.summary())" + ) + } + + /// Six of them, the surface that produced the field failure. + func testARealisticToolSurfaceIsNotUnderestimated() async { + let tools = (0 ..< 6).map { index in + AnyTool( + definition: ToolDefinition( + name: "weather_mcp__tool_\(index)", + description: Self.mcpTool.description, + inputSchema: Self.mcpTool.inputSchema + ), + invoke: { _, _ in .object([:]) } + ) + } + var state = AgentState() + state.messages = [.user("what is the weather in dublin")] + let assembled = await DefaultContextAssembler(unrankedFillLimit: 0).assemble( + systemPrompt: "You are a concise, helpful assistant.", + tools: tools, + state: state, + budget: ContextBudget(total: 8192, reservedForOutput: 768, maxTools: 6) + ) + let report = ContextAudit().audit(assembled) + XCTAssertLessThan( + report.underestimateFactor, + 1.5, + "assembled request underpriced: \(report.summary())" + ) + } + + // MARK: - Sanity + + /// Prose should not be *over*-estimated into uselessness either — + /// a counter that doubles everything would pass the checks above + /// while sending a third of the tools it could. + func testProseIsNotWildlyOverestimated() { + let chatty = ToolDefinition( + name: "note", + description: String(repeating: "Write a short note for later. ", count: 20), + inputSchema: .object( + properties: ["text": .string(description: "The note body.")], + required: ["text"] + ) + ) + let report = ContextAudit().audit(tool: chatty) + XCTAssertGreaterThan( + report.underestimateFactor, + 0.4, + "estimate is wasteful: \(report.summary())" + ) + } + + /// A tool-calling turn costs its call, not its (empty) text. This + /// is the message most likely to overflow and the easiest to price + /// at nearly zero. + func testToolCallingTurnIsNotPricedAtZero() { + let counter = HeuristicTokenCounter() + let message = Message.assistant("", toolCalls: [ + ToolCall( + id: "1", + name: "weather_mcp__get_weather_summary", + arguments: .object([ + "city_name": .string("Dublin"), + "days": .integer(7), + "detail": .string("summary"), + ]) + ), + ]) + XCTAssertGreaterThan(counter.count(message: message), 10) + } + + // MARK: Private + + private static let schemaJSON = """ + {"type":"object","properties":{"city_name":{"type":"string","description":"City name, \ + postal code, or coordinates."},"days":{"type":"integer","description":"How many days of \ + forecast, 1 to 16."},"units":{"type":"string","enum":["metric","imperial"],"description":\ + "Measurement system."},"detail":{"type":"string","description":"summary, standard or full."}},\ + "required":["city_name"]} + """ + + private static let mcpTool = ToolDefinition( + name: "weather_mcp__get_weather_summary", + description: """ + Retrieve a plain-language weather summary for a location, including current \ + conditions, today's forecast, precipitation probability, wind, humidity, UV index \ + and any active severe weather alerts. Accepts a city name or coordinates. + """, + inputSchema: .object( + properties: [ + "city_name": .string(description: "City name, postal code, or coordinates."), + "days": .integer(description: "How many days of forecast, 1 to 16."), + "units": .string( + description: "Measurement system.", + enumValues: ["metric", "imperial"] + ), + "detail": .string(description: "summary, standard or full."), + ], + required: ["city_name"] + ) + ) +} diff --git a/Tests/AriaTests/ImageContextTests.swift b/Tests/AriaTests/ImageContextTests.swift new file mode 100644 index 0000000..61361e5 --- /dev/null +++ b/Tests/AriaTests/ImageContextTests.swift @@ -0,0 +1,114 @@ +@testable import Aria +import XCTest + +/// Images were invisible to the context layer in two different ways, +/// and together they meant an attached photo either overflowed the +/// window or arrived with no tools at all. +final class ImageContextTests: XCTestCase { + // MARK: - Cost + + /// An image counted at zero is invisible to every budget decision + /// downstream: history windowing never drops it because it looks + /// free, and the request overflows with the budget satisfied. + func testAnImageCostsSomething() { + let counter = HeuristicTokenCounter() + let withImage = Message(role: .user, content: [ + .text("what is this?"), + .image(ImageContent(source: .data(Data(repeating: 0, count: 64), mimeType: "image/jpeg"))), + ]) + let withoutImage = Message(role: .user, content: [.text("what is this?")]) + + XCTAssertGreaterThan( + counter.count(message: withImage), + counter.count(message: withoutImage) + 500, + "an image priced at nearly nothing will never be windowed out" + ) + } + + /// Two images cost more than one. Sounds trivial; a per-message + /// constant would pass the test above and fail this one. + func testImagesAreCountedIndividually() { + let counter = HeuristicTokenCounter() + let image = ContentPart.image( + ImageContent(source: .data(Data(repeating: 0, count: 64), mimeType: "image/jpeg")) + ) + let one = Message(role: .user, content: [image]) + let two = Message(role: .user, content: [image, image]) + XCTAssertGreaterThan(counter.count(message: two), counter.count(message: one)) + } + + // MARK: - Selection + + /// A photo with no caption gives the ranker nothing to rank on. + /// That is not "nothing matched" — the ranker was never given a + /// question — and treating it as a decision sent zero tools. + func testImageOnlyTurnStillGetsTools() async { + // More tools than `maxTools`, or the assembler short-circuits + // and returns everything unranked — which passes this test + // without ever reaching the code it is about. + let tools = (0 ..< 12).map { index in + AnyTool( + definition: ToolDefinition( + name: "tool_\(index)", + description: "Does a thing.", + inputSchema: .object(properties: [:], required: []) + ), + invoke: { _, _ in .object([:]) } + ) + } + var state = AgentState() + state.messages = [Message(role: .user, content: [ + .image(ImageContent(source: .data(Data(repeating: 0, count: 64), mimeType: "image/jpeg"))), + ])] + + let assembled = await DefaultContextAssembler(unrankedFillLimit: 0).assemble( + systemPrompt: "You are helpful.", + tools: tools, + state: state, + budget: ContextBudget(total: 8192, reservedForOutput: 768, maxTools: 6) + ) + XCTAssertFalse( + assembled.tools.isEmpty, + "an uncaptioned image arrived with no way to act on it" + ) + } + + /// A captioned image still ranks on its caption — the fallback must + /// not swallow the normal path. + func testCaptionedImageStillRanksOnItsText() async { + let weather = AnyTool( + definition: ToolDefinition( + name: "get_weather", + description: "Get the weather forecast for a location.", + inputSchema: .object(properties: [:], required: []) + ), + invoke: { _, _ in .object([:]) } + ) + // Enough tools that ranking actually runs: under `maxTools` + // with room to spare, the assembler returns everything + // unranked, and this would be testing the early exit. + let filler = (0 ..< 10).map { index in + AnyTool( + definition: ToolDefinition( + name: "unrelated_tool_\(index)", + description: "Encode or decode base64 payloads.", + inputSchema: .object(properties: [:], required: []) + ), + invoke: { _, _ in .object([:]) } + ) + } + var state = AgentState() + state.messages = [Message(role: .user, content: [ + .text("what is the weather here"), + .image(ImageContent(source: .data(Data(repeating: 0, count: 64), mimeType: "image/jpeg"))), + ])] + + let assembled = await DefaultContextAssembler(unrankedFillLimit: 0).assemble( + systemPrompt: "You are helpful.", + tools: filler + [weather], + state: state, + budget: ContextBudget(total: 8192, reservedForOutput: 768, maxTools: 6) + ) + XCTAssertEqual(assembled.tools.first?.name, "get_weather") + } +} diff --git a/Tests/AriaTests/ToolArgumentCoercionTests.swift b/Tests/AriaTests/ToolArgumentCoercionTests.swift new file mode 100644 index 0000000..4ff5dc7 --- /dev/null +++ b/Tests/AriaTests/ToolArgumentCoercionTests.swift @@ -0,0 +1,125 @@ +@testable import Aria +import XCTest + +/// The failure these exist for: a 0.8B model called a weather tool with +/// `{"days": "1", "latitude": "56.35"}` and the server refused it — +/// *"Invalid days: must be a finite number, received string"*. Right +/// tool, right intent, right values, defeated by quotation marks. +final class ToolArgumentCoercionTests: XCTestCase { + // MARK: - The regression + + func testQuotedNumbersAreCoercedToTheirDeclaredTypes() { + let coerced = ToolArgumentCoercion.coerce( + .object([ + "days": .string("1"), + "latitude": .string("56.35"), + "city_name": .string("Dublin"), + "include_alerts": .string("true"), + ]), + to: Self.weatherSchema + ) + XCTAssertEqual(coerced, .object([ + "days": .integer(1), + "latitude": .number(56.35), + "city_name": .string("Dublin"), + "include_alerts": .bool(true), + ])) + } + + // MARK: - Restraint + + /// A value that is not losslessly convertible is left alone. The + /// server's own error beats a number this invented. + func testUnconvertibleValuesAreLeftForTheServerToReject() { + let coerced = ToolArgumentCoercion.coerce( + .object(["days": .string("next week")]), + to: Self.weatherSchema + ) + XCTAssertEqual(coerced, .object(["days": .string("next week")])) + } + + /// Truncating 1.7 to 1 would change what the model asked for. + func testFractionalValuesAreNotTruncatedIntoIntegers() { + let coerced = ToolArgumentCoercion.coerce( + .object(["days": .number(1.7)]), + to: Self.weatherSchema + ) + XCTAssertEqual(coerced, .object(["days": .number(1.7)])) + } + + /// "yes" is a guess about intent, not a reading of the value. + func testOnlyJSONBooleanSpellingsAreAccepted() { + for text in ["yes", "1", "on", "Y"] { + let coerced = ToolArgumentCoercion.coerce( + .object(["include_alerts": .string(text)]), + to: Self.weatherSchema + ) + XCTAssertEqual( + coerced, + .object(["include_alerts": .string(text)]), + "\"\(text)\" should not have been read as a boolean" + ) + } + } + + /// A field the schema does not describe has no contract to enforce. + func testUndeclaredFieldsPassThroughUntouched() { + let coerced = ToolArgumentCoercion.coerce( + .object(["mystery": .string("42")]), + to: Self.weatherSchema + ) + XCTAssertEqual(coerced, .object(["mystery": .string("42")])) + } + + func testCorrectlyTypedArgumentsAreUnchanged() { + let already = JSONValue.object([ + "days": .integer(3), + "latitude": .number(37.7), + "city_name": .string("Dublin"), + ]) + XCTAssertEqual(ToolArgumentCoercion.coerce(already, to: Self.weatherSchema), already) + } + + // MARK: - Shape + + func testNestedObjectsAndArraysAreCoercedRecursively() { + let schema = JSONSchema.object( + properties: [ + "points": .array(items: .object( + properties: ["value": .number(description: nil)], + required: [] + )) + ], + required: [] + ) + let coerced = ToolArgumentCoercion.coerce( + .object(["points": .array([.object(["value": .string("2.5")])])]), + to: schema + ) + XCTAssertEqual( + coerced, + .object(["points": .array([.object(["value": .number(2.5)])])]) + ) + } + + /// The mirror case: a bare number where a string is declared. + func testNumbersAreStringifiedWhenTheSchemaWantsText() { + let coerced = ToolArgumentCoercion.coerce( + .object(["city_name": .integer(94568)]), + to: Self.weatherSchema + ) + XCTAssertEqual(coerced, .object(["city_name": .string("94568")])) + } + + // MARK: Private + + private static let weatherSchema = JSONSchema.object( + properties: [ + "city_name": .string(description: nil, enumValues: nil), + "days": .integer(description: nil), + "latitude": .number(description: nil), + "include_alerts": .boolean(description: nil), + ], + required: ["city_name"] + ) +} diff --git a/Tests/AriaTests/ToolBudgetCeilingTests.swift b/Tests/AriaTests/ToolBudgetCeilingTests.swift new file mode 100644 index 0000000..0e510c3 --- /dev/null +++ b/Tests/AriaTests/ToolBudgetCeilingTests.swift @@ -0,0 +1,196 @@ +@testable import Aria +import XCTest + +/// `maxTools` bounds how *many* tools are sent; `toolTokenLimit` bounds +/// what they *cost*. Only the count was enforced once ranking returned +/// something, so a small number of large schemas could still overflow +/// the window — which is exactly what a connected MCP server did: +/// 5,845 tokens into a 4,096-token window, with the budget reporting +/// itself satisfied. +final class ToolBudgetCeilingTests: XCTestCase { + /// The regression. Six large tools sit under `maxTools: 6` and far + /// over the token ceiling. + func testRankedToolsAreTrimmedToTheTokenCeiling() async throws { + let tools = (0 ..< 6).map { Self.tool(named: "weather_\($0)", descriptionLength: 1200) } + let budget = ContextBudget(total: 4096, reservedForOutput: 768, maxTools: 6) + let assembler = DefaultContextAssembler(unrankedFillLimit: 0) + + let request = await assembler.assemble( + systemPrompt: "You are helpful.", + tools: tools, + state: Self.state(query: "what is the weather"), + budget: budget + ) + + let counter = HeuristicTokenCounter() + let spent = request.tools.reduce(0) { $0 + counter.count(tool: $1.definition) } + XCTAssertLessThanOrEqual( + spent, + budget.toolTokenLimit, + "selected \(request.tools.count) tools costing \(spent) against a ceiling of \(budget.toolTokenLimit)" + ) + XCTAssertFalse( + request.tools.isEmpty, + "trimming to nothing would make the turn unable to act" + ) + } + + /// The regression from the field: ranking chose the right tool and + /// the share admitted none of them, leaving a turn that could only + /// answer from imagination. + /// + /// A single MCP schema on a real server runs to roughly a thousand + /// tokens, which is over the share once a pinned tool has taken its + /// cut. Respecting the share exactly is right when it costs a + /// fourth tool and wrong when it costs the only one. + func testTopRankedToolSurvivesEvenWhenItBustsTheShare() async throws { + let huge = Self.tool(named: "weather_summary", descriptionLength: 6000) + let pinned = Self.tool(named: "load_skill", descriptionLength: 40) + let assembler = DefaultContextAssembler( + pinnedToolNames: ["load_skill"], + unrankedFillLimit: 0 + ) + var state = AgentState() + state.messages = [.user("weather summary please")] + + let assembled = await assembler.assemble( + systemPrompt: "You are helpful.", + tools: [pinned, huge], + state: state, + budget: ContextBudget(total: 4096, reservedForOutput: 768, maxTools: 6) + ) + XCTAssertTrue( + assembled.tools.contains { $0.name == "weather_summary" }, + "the only tool that could answer was trimmed away: \(assembled.tools.map(\.name))" + ) + } + + /// Survives every compaction level: many required properties with + /// long names, which no level is permitted to remove. + private static func incompressibleTool(named name: String) -> AnyTool { + var properties: [String: JSONSchema] = [:] + var required: [String] = [] + for index in 0 ..< 220 { + let key = "detailed_configuration_parameter_number_\(index)_for_the_request" + properties[key] = .string(description: nil, enumValues: nil) + required.append(key) + } + return AnyTool( + definition: ToolDefinition( + name: name, + description: "Weather.", + inputSchema: .object(properties: properties, required: required) + ), + invoke: { _, _ in .object([:]) } + ) + } + + /// Sized and shaped like a published MCP schema: a paragraph of + /// description and several documented parameters. + private static func mcpSizedTool(named name: String) -> AnyTool { + AnyTool( + definition: ToolDefinition( + name: name, + description: String( + repeating: "Retrieve weather forecast conditions for a location, " + + "including temperature, precipitation, wind, humidity and alerts. ", + count: 20 + ), + inputSchema: .object( + properties: [ + "location": .string( + description: String(repeating: "City name or coordinates. ", count: 24), + enumValues: nil + ), + "units": .string( + description: String(repeating: "Measurement system to use. ", count: 24), + enumValues: ["metric", "imperial"] + ), + "detail": .string( + description: String(repeating: "How much detail to return. ", count: 24), + enumValues: nil + ), + ], + required: ["location"] + ) + ), + invoke: { _, _ in .object([:]) } + ) + } + + /// The field regression: three MCP tools worth 4,533 tokens went + /// out against a 2,996-token budget, history collapsed to nothing + /// to make room, and the provider refused the request anyway. + /// + /// The share is breakable by design — sending no usable tool is + /// worse than overspending it. The *budget* is not. + func testToolsNeverExceedTheHardBudget() async throws { + // Incompressible on purpose: compaction may drop descriptions + // and optional properties, never required ones. A tool whose + // required surface alone busts the budget is the case the hard + // bound exists for. + let tools = (0 ..< 6).map { Self.incompressibleTool(named: "weather_mcp__tool_\($0)") } + let budget = ContextBudget(total: 4096, reservedForOutput: 768, maxTools: 6) + var state = AgentState() + state.messages = [.user("weather forecast for dublin please")] + + let assembled = await DefaultContextAssembler(unrankedFillLimit: 0).assemble( + systemPrompt: "You are helpful.", + tools: tools, + state: state, + budget: budget + ) + + let counter = HeuristicTokenCounter() + let toolTokens = assembled.tools.reduce(0) { $0 + counter.count(tool: $1.definition) } + let messageTokens = assembled.messages.reduce(0) { $0 + counter.count(message: $1) } + XCTAssertLessThanOrEqual( + toolTokens + messageTokens, + budget.available, + "assembled \(toolTokens) tool + \(messageTokens) message tokens against \(budget.available)" + ) + // And the user's turn survived, which is the point of the + // reservation — a request that cannot carry the question is not + // a smaller request, it is a broken one. + XCTAssertTrue( + assembled.messages.contains { $0.role == .user }, + "the user message was dropped to make room for tools" + ) + } + + /// Small tools are unaffected — the ceiling must not become a cap + /// that starves ordinary surfaces. + func testSmallToolSurfacesAreNotTrimmed() async throws { + let tools = (0 ..< 5).map { Self.tool(named: "t\($0)", descriptionLength: 40) } + let assembler = DefaultContextAssembler(unrankedFillLimit: 0) + let request = await assembler.assemble( + systemPrompt: "You are helpful.", + tools: tools, + state: Self.state(query: "t1 t2 t3"), + budget: ContextBudget(total: 8192, reservedForOutput: 768, maxTools: 12) + ) + XCTAssertEqual(request.tools.count, 5) + } + + // MARK: Private + + private static func state(query: String) -> AgentState { + var state = AgentState() + state.messages = [.user(query)] + return state + } + + private static func tool(named name: String, descriptionLength: Int) -> AnyTool { + AnyTool( + definition: ToolDefinition( + name: name, + description: String(repeating: "weather forecast data ", count: descriptionLength / 22), + inputSchema: .object( + properties: ["location": .string(description: "Place name")], + required: ["location"] + ) + ), + invoke: { _, _ in .object([:]) } + ) + } +} diff --git a/Tests/AriaTests/ToolDefinitionCompactorTests.swift b/Tests/AriaTests/ToolDefinitionCompactorTests.swift new file mode 100644 index 0000000..6753654 --- /dev/null +++ b/Tests/AriaTests/ToolDefinitionCompactorTests.swift @@ -0,0 +1,131 @@ +@testable import Aria +import XCTest + +/// A published MCP tool schema is routinely a thousand tokens — a +/// paragraph of description plus a dozen documented parameters. On a +/// 4,096-token window that is a quarter of everything, for one tool +/// the model may not call. +/// +/// The constraint these pin: compaction removes *guidance*, never +/// *capability*. A call the model could make before must still +/// validate after, because a tool it calls with worse arguments is +/// recoverable and a tool whose arguments no longer validate is not. +final class ToolDefinitionCompactorTests: XCTestCase { + // MARK: - Safety + + /// The one outcome this must never produce. + func testRequiredPropertiesSurviveEveryLevel() { + for level in ToolCompaction.allCases { + let compacted = ToolDefinitionCompactor.apply(level, to: Self.weatherTool) + guard case let .object(properties, required, _, _) = compacted.inputSchema else { + return XCTFail("schema shape changed at \(level)") + } + XCTAssertEqual(required, ["location"], "required list changed at \(level)") + XCTAssertNotNil(properties["location"], "required property dropped at \(level)") + } + } + + /// Enum values are the set of legal inputs, not documentation. A + /// model guessing outside them produces a call that fails. + func testEnumValuesSurviveEveryLevel() { + for level in ToolCompaction.allCases { + let compacted = ToolDefinitionCompactor.apply(level, to: Self.weatherTool) + guard case let .object(properties, _, _, _) = compacted.inputSchema, + case let .string(_, enumValues) = properties["units"] ?? .null else { + // `units` is optional and legitimately gone at the last + // level; its absence is not a failure. + XCTAssertEqual(level, .requiredOnly, "units vanished early at \(level)") + continue + } + XCTAssertEqual(enumValues, ["metric", "imperial"], "enum lost at \(level)") + } + } + + func testNameIsNeverChanged() { + for level in ToolCompaction.allCases { + XCTAssertEqual( + ToolDefinitionCompactor.apply(level, to: Self.weatherTool).name, + "weather_mcp__get_weather_summary" + ) + } + } + + // MARK: - Effect + + /// Each level must actually be cheaper, or it is not worth the + /// capability it costs. + func testEachLevelIsStrictlyCheaper() { + let counter = HeuristicTokenCounter() + let costs = ToolCompaction.allCases.map { + counter.count(tool: ToolDefinitionCompactor.apply($0, to: Self.weatherTool)) + } + for (index, cost) in costs.enumerated().dropFirst() { + XCTAssertLessThan(cost, costs[index - 1], "level \(index) did not shrink anything") + } + // The headline: the thing that made this necessary. + XCTAssertLessThan( + Double(costs.last ?? 0), + Double(costs[0]) * 0.5, + "full compaction saved less than half: \(costs)" + ) + } + + /// Stops at the lightest level that fits rather than flattening + /// everything — capability is only given up when it buys something. + func testStopsAtTheLightestLevelThatFits() { + let counter = HeuristicTokenCounter() + let full = counter.count(tool: Self.weatherTool) + let (_, applied) = ToolDefinitionCompactor.compact( + Self.weatherTool, + toFit: full, + counter: counter + ) + XCTAssertEqual(applied, .none, "compacted a tool that already fit") + } + + /// A description with no sentence break still shrinks. + func testUnpunctuatedDescriptionIsCapped() { + let rambling = ToolDefinition( + name: "t", + description: String(repeating: "words and more words ", count: 60), + inputSchema: .object(properties: [:], required: []) + ) + let compacted = ToolDefinitionCompactor.apply(.shortDescription, to: rambling) + XCTAssertLessThan(compacted.description.count, rambling.description.count / 2) + } + + // MARK: Private + + /// Shaped like the real thing that caused the problem. + private static let weatherTool = ToolDefinition( + name: "weather_mcp__get_weather_summary", + description: """ + Retrieve a plain-language weather summary for a location. Includes current \ + conditions, today's forecast, precipitation probability, wind, humidity, UV \ + index and any active severe weather alerts. Accepts a city name or \ + coordinates. Data comes from NOAA and Open-Meteo and is cached for fifteen \ + minutes. See also get_forecast for multi-day output. + """, + inputSchema: .object( + properties: [ + "location": .string( + description: "City name, postal code, or 'lat,lon' coordinates. Required.", + enumValues: nil + ), + "units": .string( + description: "Measurement system for temperature, wind and precipitation.", + enumValues: ["metric", "imperial"] + ), + "detail": .string( + description: "How much to return: summary is one paragraph, full includes hourly breakdowns.", + enumValues: ["summary", "standard", "full"] + ), + "include_alerts": .boolean( + description: "Whether to include active severe weather alerts in the response." + ), + ], + required: ["location"], + description: "Arguments for the weather summary request." + ) + ) +} diff --git a/Tests/AriaTests/ToolRelevanceCorpusTests.swift b/Tests/AriaTests/ToolRelevanceCorpusTests.swift new file mode 100644 index 0000000..4a31cc6 --- /dev/null +++ b/Tests/AriaTests/ToolRelevanceCorpusTests.swift @@ -0,0 +1,114 @@ +@testable import Aria +import XCTest + +/// Queries that should send *no* tool, and queries that should. +/// +/// The eval corpus measures whether the right tool is offered. It has +/// almost nothing on the opposite failure — a tool offered when none is +/// relevant — and that is the one users notice. Asked "What's the +/// tracking number?", the ranker offered `calculator`, the model called +/// it with `(4.45 - 4.45) / 16` scavenged from a fasting summary, and +/// the reply was `0`. +/// +/// The two known failures are asserted as *expected* failures rather +/// than deleted or quietly passed. They are real, they are not fixed, +/// and a corpus that hides them is worth less than one that states +/// them. +final class ToolRelevanceCorpusTests: XCTestCase { + // MARK: - Nothing should rank + + /// Conversation, gratitude, and questions about the conversation + /// itself. These already work. + func testConversationalTurnsRankNothing() async { + for query in [ + "How are you?", + "Thanks!", + "What did we talk about earlier?", + "Explain quantum computing", + ] { + let ranked = await LexicalToolSelector().select( + from: Self.surface, + query: query, + limit: 6 + ) + XCTAssertTrue(ranked.isEmpty, "\"\(query)\" offered \(ranked.map(\.name))") + } + } + + /// A single common word matching a single tool. + /// + /// "number" appears in `calculator`, `uuid` and `unit_converter`; + /// "fact" appears in `remember_fact`. Neither query is *about* + /// those tools, and one term of overlap out of two is not + /// responsiveness — but nothing in the ranker says so yet. + /// + /// The obvious fix — require two matched terms — breaks "remind me + /// to buy milk", where one matched term is the whole intent. That + /// is why this is recorded rather than patched: the rule needs + /// measuring against more than the four cases that motivated it. + func testSingleCommonWordShouldNotRankATool() async { + XCTExpectFailure("known gap: no relevance floor on a single weak term") + for query in ["What's the tracking number?", "Tell me a fun fact"] { + let ranked = await LexicalToolSelector().select( + from: Self.surface, + query: query, + limit: 6 + ) + XCTAssertTrue(ranked.isEmpty, "\"\(query)\" offered \(ranked.map(\.name))") + } + } + + // MARK: - Something should rank + + /// The other side of the trade. Any relevance floor added later has + /// to keep every one of these. + func testRealRequestsStillRankTheRightTool() async { + let expectations = [ + "Get me weather summary": "weather_mcp__get_weather_summary", + "How am I doing with fasting today?": "niora__get_fasting_status", + "log 500ml of water": "niora__log_water", + "what time is it": "current_time", + ] + for (query, expected) in expectations { + let ranked = await LexicalToolSelector().select( + from: Self.surface, + query: query, + limit: 6 + ) + XCTAssertEqual(ranked.first?.name, expected, "\"\(query)\" ranked \(ranked.map(\.name))") + } + } + + // MARK: Private + + /// The surface from a real device: built-ins, a weather MCP server, + /// Niora, and two workflow remixes. + private static let surface: [ToolDefinition] = [ + ("current_time", "Get the current date and time."), + ("remember_fact", "Save a durable fact about the user."), + ("http_request", "Fetch a URL over HTTP."), + ("calculator", "Evaluate an arithmetic expression and return the resulting number."), + ("json_path", "Extract values from JSON with a path expression."), + ("regex", "Match or replace text with a regular expression."), + ("uuid", "Generate a unique identifier number."), + ("unit_converter", "Convert a number between units of measurement."), + ("random_picker", "Pick a random item from a list."), + ("date_math", "Add or subtract intervals from a date."), + ("weather_mcp__get_weather_summary", "Get a plain-language weather summary for a location."), + ("weather_mcp__get_forecast", "Get a multi-day weather forecast for a location."), + ("weather_mcp__get_current_conditions", "Get current weather conditions for a location."), + ("weather_mcp__get_alerts", "Get active severe weather alerts for a location."), + ("niora__get_fasting_status", "Whether the user is fasting now, and progress toward target."), + ("niora__start_fast", "Begin a new fast for the user."), + ("niora__end_fast", "End the user's current fast."), + ("niora__log_water", "Log a water intake entry in millilitres."), + ("niora__get_hydration_today", "How much water the user has drunk today."), + ("niora__get_profile", "The user's profile: display name, goals, dietary preferences."), + ].map { + ToolDefinition( + name: $0.0, + description: $0.1, + inputSchema: .object(properties: [:], required: []) + ) + } +} diff --git a/Tests/AriaToolsTests/TabularExtractionTests.swift b/Tests/AriaToolsTests/TabularExtractionTests.swift new file mode 100644 index 0000000..6fb0b39 --- /dev/null +++ b/Tests/AriaToolsTests/TabularExtractionTests.swift @@ -0,0 +1,97 @@ +import Aria +@testable import AriaTools +import XCTest + +final class TabularExtractionTests: XCTestCase { + /// A bare array is the simplest tool result shape. + func testBareArray() { + let value = JSONValue.array([Self.row(1), Self.row(2)]) + XCTAssertEqual(TabularExtraction.rows(from: value)?.count, 2) + XCTAssertEqual(TabularExtraction.extractDetailed(from: value)?.path, []) + } + + /// Envelope keys are checked in a fixed order, so a result with + /// several arrays charts the same thing every time rather than + /// depending on dictionary ordering. + func testPrefersConventionalEnvelopeKeys() { + let value = JSONValue.object([ + "warnings": .array([Self.row(9)]), + "items": .array([Self.row(1), Self.row(2)]), + ]) + let found = TabularExtraction.extractDetailed(from: value) + XCTAssertEqual(found?.path, ["items"]) + XCTAssertEqual(found?.rows.count, 2) + } + + func testFindsNestedRows() { + let value = JSONValue.object([ + "data": .object(["records": .array([Self.row(1)])]), + ]) + XCTAssertEqual(TabularExtraction.extractDetailed(from: value)?.path, ["data", "records"]) + } + + /// Tools that hand back JSON as a string are common enough to + /// support — but only at the top level, since decoding strings + /// buried inside a payload would be guessing. + func testDecodesTopLevelJSONString() { + let value = JSONValue.string(#"{"items":[{"x":1,"y":2}]}"#) + XCTAssertEqual(TabularExtraction.rows(from: value)?.count, 1) + } + + /// A mixed array is not a table. Charting only its object elements + /// would silently drop data, which is worse than declining. + func testMixedArrayIsNotATable() { + let value = JSONValue.array([Self.row(1), .string("nope")]) + XCTAssertNil(TabularExtraction.rows(from: value)) + } + + func testEmptyArrayIsNotATable() { + XCTAssertNil(TabularExtraction.rows(from: .array([]))) + } + + func testScalarsAndProseAreNotTables() { + XCTAssertNil(TabularExtraction.rows(from: .integer(4))) + XCTAssertNil(TabularExtraction.rows(from: .string("no data today"))) + XCTAssertNil(TabularExtraction.rows(from: .object(["note": .string("hi")]))) + } + + /// Ten thousand points is an unreadable chart and a slow one. The + /// excess is dropped, and the caller is told so it can say so. + func testTruncatesAndReportsOriginalCount() { + let rows = (0 ..< 50).map { Self.row($0) } + let found = TabularExtraction.extractDetailed( + from: .array(rows), + maximumRows: 10 + ) + XCTAssertEqual(found?.rows.count, 10) + XCTAssertEqual(found?.truncatedFrom, 50) + } + + func testNoTruncationReportsNil() { + let found = TabularExtraction.extractDetailed(from: .array([Self.row(1)])) + XCTAssertNil(found?.truncatedFrom) + } + + /// End to end: a realistic MCP tool result should chart. + func testRealisticToolResultCompilesToAChart() throws { + let payload = JSONValue.object([ + "series": .array([ + .object(["date": .string("2026-01-01"), "weight_kg": .number(80.2)]), + .object(["date": .string("2026-01-08"), "weight_kg": .number(79.6)]), + ]), + ]) + let rows = try XCTUnwrap(TabularExtraction.rows(from: payload)) + XCTAssertNoThrow( + try VegaLiteCompiler.compile( + ChartIntent(kind: .line, dataRef: "r1", x: "date", y: "weight_kg"), + rows: rows + ) + ) + } + + // MARK: Private + + private static func row(_ n: Int) -> JSONValue { + .object(["x": .integer(Int64(n)), "y": .integer(Int64(n * 2))]) + } +} diff --git a/Tests/AriaToolsTests/VegaLiteCompilerTests.swift b/Tests/AriaToolsTests/VegaLiteCompilerTests.swift new file mode 100644 index 0000000..41fe933 --- /dev/null +++ b/Tests/AriaToolsTests/VegaLiteCompilerTests.swift @@ -0,0 +1,152 @@ +import Aria +@testable import AriaTools +import XCTest + +final class VegaLiteCompilerTests: XCTestCase { + // MARK: - Grounding + + /// The model names fields; it never supplies values. If it names a + /// field that is not there, the chart must fail loudly rather than + /// render an empty axis — and the error carries the real field + /// names so the failure is recoverable. + func testUnknownFieldFailsWithTheAvailableNames() { + XCTAssertThrowsError( + try VegaLiteCompiler.compile( + ChartIntent(kind: .line, dataRef: "r1", x: "day", y: "value"), + rows: Self.rows + ) + ) { error in + XCTAssertEqual( + error as? ChartCompilerError, + // `x` is validated first, so "day" is the name reported. + .unknownField(name: "day", available: ["date", "region", "value"]) + ) + } + } + + func testEmptyDataFails() { + XCTAssertThrowsError( + try VegaLiteCompiler.compile( + ChartIntent(kind: .bar, dataRef: "r1", x: "date", y: "value"), + rows: [] + ) + ) { XCTAssertEqual($0 as? ChartCompilerError, .emptyData) } + } + + /// Values in the spec must come from the rows the host supplied. + func testDataIsInlinedFromTheHostRows() throws { + let spec = try VegaLiteCompiler.compile( + ChartIntent(kind: .line, dataRef: "r1", x: "date", y: "value"), + rows: Self.rows + ) + guard case let .object(top) = spec, + case let .object(data) = top["data"], + case let .array(values) = data["values"] else { + return XCTFail("expected inlined data values") + } + XCTAssertEqual(values.count, 3) + } + + // MARK: - Type inference + + /// The host has the values, so it infers types rather than asking + /// the model to declare them. + func testInfersQuantitativeTemporalAndNominal() throws { + let spec = try VegaLiteCompiler.compile( + ChartIntent(kind: .line, dataRef: "r1", x: "date", y: "value", series: "region"), + rows: Self.rows + ) + XCTAssertEqual(Self.type(of: "x", in: spec), "temporal") + XCTAssertEqual(Self.type(of: "y", in: spec), "quantitative") + XCTAssertEqual(Self.type(of: "color", in: spec), "nominal") + } + + /// One non-numeric value makes the whole axis untrustworthy, so the + /// field degrades to nominal rather than rendering a broken scale. + func testOneBadValueDemotesQuantitativeToNominal() throws { + var rows = Self.rows + rows[1]["value"] = .string("n/a") + let spec = try VegaLiteCompiler.compile( + ChartIntent(kind: .bar, dataRef: "r1", x: "region", y: "value"), + rows: rows + ) + XCTAssertEqual(Self.type(of: "y", in: spec), "nominal") + } + + /// A category that merely starts with a digit must not become a + /// time axis — a false temporal is worse than a plain category. + func testDigitLeadingCategoriesAreNotTemporal() throws { + let rows: [[String: JSONValue]] = [ + ["label": .string("3 items"), "n": .integer(1)], + ["label": .string("4 items"), "n": .integer(2)], + ] + let spec = try VegaLiteCompiler.compile( + ChartIntent(kind: .bar, dataRef: "r1", x: "label", y: "n"), + rows: rows + ) + XCTAssertEqual(Self.type(of: "x", in: spec), "nominal") + } + + // MARK: - Shape + + /// Pie has no axes — the same two fields mean angle and category. + func testPieUsesThetaAndColorRatherThanAxes() throws { + let spec = try VegaLiteCompiler.compile( + ChartIntent(kind: .pie, dataRef: "r1", x: "region", y: "value"), + rows: Self.rows + ) + guard case let .object(top) = spec, case let .object(encoding) = top["encoding"] else { + return XCTFail("no encoding") + } + XCTAssertNotNil(encoding["theta"]) + XCTAssertNil(encoding["x"]) + } + + /// An interactive chart the user cannot interrogate is a picture. + func testTooltipsAreAlwaysEnabled() throws { + let spec = try VegaLiteCompiler.compile( + ChartIntent(kind: .scatter, dataRef: "r1", x: "date", y: "value"), + rows: Self.rows + ) + guard case let .object(top) = spec, + case let .object(encoding) = top["encoding"], + case let .array(tooltip) = encoding["tooltip"], + case let .object(mark) = top["mark"] else { + return XCTFail("no tooltip") + } + XCTAssertEqual(tooltip.count, 2) + XCTAssertEqual(mark["tooltip"], .bool(true)) + } + + /// Every kind must compile — the enum exists so none of them can be + /// a runtime surprise. + func testEveryChartKindCompiles() throws { + for kind in ChartKind.allCases { + XCTAssertNoThrow( + try VegaLiteCompiler.compile( + ChartIntent(kind: kind, dataRef: "r1", x: "region", y: "value"), + rows: Self.rows + ), + "\(kind) failed to compile" + ) + } + } + + // MARK: Private + + private static let rows: [[String: JSONValue]] = [ + ["date": .string("2026-01-01"), "value": .integer(3), "region": .string("west")], + ["date": .string("2026-01-02"), "value": .integer(5), "region": .string("west")], + ["date": .string("2026-01-03"), "value": .integer(4), "region": .string("east")], + ] + + private static func type(of channel: String, in spec: JSONValue) -> String? { + guard case let .object(top) = spec, + case let .object(encoding) = top["encoding"], + case let .object(field) = encoding[channel], + case let .string(type) = field["type"] else { + return nil + } + return type + } +} diff --git a/Tests/WorkflowKitTests/MCPResourceTests.swift b/Tests/WorkflowKitTests/MCPResourceTests.swift new file mode 100644 index 0000000..48b3cfa --- /dev/null +++ b/Tests/WorkflowKitTests/MCPResourceTests.swift @@ -0,0 +1,71 @@ +import Foundation +import MCP +@testable import WorkflowKit +import XCTest + +// MARK: - MCPResourceTests + +final class MCPResourceTests: XCTestCase { + // MARK: - isInteractiveUI + + /// MCP Apps are identified by *both* a `ui://` URI and the + /// `profile=mcp-app` MIME parameter. Treating plain `text/html` as + /// an app would hand a message channel to a server that only meant + /// to serve a document. + func testInteractiveUIRequiresBothURISchemeAndProfile() { + XCTAssertTrue(Self.resource("ui://x", "text/html;profile=mcp-app").isInteractiveUI) + XCTAssertTrue(Self.resource("ui://x", "text/html; profile=mcp-app").isInteractiveUI) + XCTAssertFalse(Self.resource("ui://x", "text/html").isInteractiveUI) + XCTAssertFalse(Self.resource("https://x", "text/html;profile=mcp-app").isInteractiveUI) + XCTAssertFalse(Self.resource("ui://x", nil).isInteractiveUI) + } + + // MARK: - Against the fake server + + func testListResourcesDrainsPagination() async throws { + let resources = try await Self.withFakeServer { client in + try await client.listResources() + } + XCTAssertEqual(resources.count, 2, "second page was dropped") + XCTAssertTrue(resources.contains { $0.isInteractiveUI }) + XCTAssertTrue(resources.contains { $0.uri == "file://data.csv" }) + } + + func testReadResourceReturnsMarkup() async throws { + let result = try await Self.withFakeServer { client in + try await client.readResource(uri: "ui://dashboard") + } + XCTAssertEqual(result.firstHTMLResource?.text, "

hi

") + } + + func testListPromptsCarriesArguments() async throws { + let prompts = try await Self.withFakeServer { client in + try await client.listPrompts() + } + XCTAssertEqual(prompts.first?.name, "summarise") + XCTAssertEqual(prompts.first?.arguments.first?.name, "uri") + XCTAssertEqual(prompts.first?.arguments.first?.required, true) + } + + // MARK: Private + + private static func resource(_ uri: String, _ mime: String?) -> MCPResourceDescriptor { + MCPResourceDescriptor(uri: uri, name: "n", mimeType: mime) + } + + /// A real `MCPClient` wired to the in-process fake, so the + /// mapping under test is the shipping code path. + private static func withFakeServer( + _ body: @Sendable (MCPClient) async throws -> T + ) async throws -> T { + let counter = FakeMCPTransport.ConnectCounter() + let client = MCPClient( + serverURL: URL(string: "https://example.test/mcp")!, + credential: nil, + clientName: "test-\(UUID().uuidString)", + clientVersion: "1.0", + transportFactory: { FakeMCPTransport(counter: counter) } + ) + return try await body(client) + } +} diff --git a/Tests/WorkflowKitTests/MCPSessionPoolTests.swift b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift new file mode 100644 index 0000000..7454e87 --- /dev/null +++ b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift @@ -0,0 +1,272 @@ +import Foundation +import Logging +import MCP +@testable import WorkflowKit +import XCTest + +// MARK: - FakeMCPTransport + +/// An in-process MCP server, just complete enough to complete a +/// handshake and answer `tools/list`. +/// +/// Worth the ~80 lines: the pool's whole job is deciding *when* to +/// connect, and that is invisible against a live server and untestable +/// against none. Every connect is counted here, which is the assertion +/// the pool actually needs. +actor FakeMCPTransport: Transport { + // MARK: Lifecycle + + init(counter: ConnectCounter, failSend: Bool = false) { + self.counter = counter + self.failSend = failSend + } + + // MARK: Internal + + /// Shared across the transports one factory produces, so a retry's + /// fresh transport still reports into the same tally. + final class ConnectCounter: @unchecked Sendable { + private let lock = NSLock() + private var value = 0 + + var count: Int { + self.lock.lock() + defer { self.lock.unlock() } + return self.value + } + + func increment() { + self.lock.lock() + self.value += 1 + self.lock.unlock() + } + } + + nonisolated var logger: Logger { Logger(label: "fake") } + + func connect() async throws { + self.counter.increment() + } + + func disconnect() async { + self.continuation?.finish() + self.continuation = nil + } + + func send(_ data: Data) async throws { + if self.failSend { + throw MCP.MCPError.connectionClosed + } + guard + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let method = object["method"] as? String + else { + return + } + // Notifications carry no id and expect no reply. + guard let id = object["id"] else { + return + } + let result: [String: Any] + switch method { + case "initialize": + result = [ + "protocolVersion": "2025-06-18", + "capabilities": [ + "tools": [:] as [String: Any], + "resources": [:] as [String: Any], + "prompts": [:] as [String: Any], + ], + "serverInfo": ["name": "fake", "version": "1.0"], + ] + case "tools/list": + result = ["tools": [[ + "name": "echo", + "description": "Echo input.", + "inputSchema": ["type": "object"] as [String: Any], + ]]] + case "resources/list": + // Two pages, so pagination is actually exercised rather + // than assumed: a client that ignores `nextCursor` sees + // only half a server's resources. + let params = object["params"] as? [String: Any] + if params?["cursor"] as? String == "page2" { + result = ["resources": [[ + "uri": "file://data.csv", "name": "Data", "mimeType": "text/csv", + ]]] + } else { + result = [ + "resources": [[ + "uri": "ui://dashboard", "name": "Dashboard", + "mimeType": "text/html;profile=mcp-app", + ]], + "nextCursor": "page2", + ] + } + case "resources/read": + result = ["contents": [[ + "uri": "ui://dashboard", + "mimeType": "text/html;profile=mcp-app", + "text": "

hi

", + ]]] + case "prompts/list": + result = ["prompts": [[ + "name": "summarise", + "description": "Summarise a document.", + "arguments": [["name": "uri", "description": "What to read", "required": true]], + ]]] + default: + result = [:] + } + let response: [String: Any] = ["jsonrpc": "2.0", "id": id, "result": result] + guard let encoded = try? JSONSerialization.data(withJSONObject: response) else { + return + } + self.continuation?.yield(encoded) + } + + func receive() -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + self.continuation = continuation + } + } + + // MARK: Private + + private let counter: ConnectCounter + private let failSend: Bool + private var continuation: AsyncThrowingStream.Continuation? +} + +// MARK: - MCPSessionPoolTests + +final class MCPSessionPoolTests: XCTestCase { + /// The point of the whole change: two calls, one handshake. + /// + /// Before pooling, every call connected and disconnected — so a + /// turn invoking three tools ran three `initialize` handshakes and + /// presented itself to the server as three separate sessions. + func testReusesOneSessionAcrossCalls() async throws { + let counter = FakeMCPTransport.ConnectCounter() + let pool = MCPSessionPool() + let key = Self.key() + + for _ in 0 ..< 3 { + _ = try await pool.withClient( + key: key, + makeTransport: { FakeMCPTransport(counter: counter) }, + body: { client in try await client.listTools(cursor: nil).0.count } + ) + } + XCTAssertEqual(counter.count, 1, "expected one handshake for three calls") + await pool.closeAll() + } + + /// A refreshed OAuth token bakes a new `Authorization` header into + /// a new transport, so it must not ride the old session. + func testDifferentCredentialGetsItsOwnSession() async throws { + let counter = FakeMCPTransport.ConnectCounter() + let pool = MCPSessionPool() + + for fingerprint in [1, 2] { + _ = try await pool.withClient( + key: Self.key(fingerprint: fingerprint), + makeTransport: { FakeMCPTransport(counter: counter) }, + body: { client in try await client.listTools(cursor: nil).0.count } + ) + } + XCTAssertEqual(counter.count, 2) + await pool.closeAll() + } + + /// A pooled connection can die between calls. The caller cannot + /// tell that from a bad request, so the pool absorbs it once. + func testRetriesOnceOnConnectionFailure() async throws { + let counter = FakeMCPTransport.ConnectCounter() + let pool = MCPSessionPool() + let attempts = FakeMCPTransport.ConnectCounter() + + let value: Int = try await pool.withClient( + key: Self.key(), + makeTransport: { FakeMCPTransport(counter: counter) }, + body: { client in + attempts.increment() + if attempts.count == 1 { + throw MCP.MCPError.connectionClosed + } + return try await client.listTools(cursor: nil).0.count + } + ) + XCTAssertEqual(value, 1) + XCTAssertEqual(attempts.count, 2, "should have retried exactly once") + XCTAssertEqual(counter.count, 2, "retry needs a fresh transport") + await pool.closeAll() + } + + /// A server error is the server *answering*. Retrying it would + /// risk performing a non-idempotent action twice. + func testDoesNotRetryServerErrors() async throws { + let counter = FakeMCPTransport.ConnectCounter() + let pool = MCPSessionPool() + let attempts = FakeMCPTransport.ConnectCounter() + + do { + _ = try await pool.withClient( + key: Self.key(), + makeTransport: { FakeMCPTransport(counter: counter) }, + body: { _ -> Int in + attempts.increment() + throw MCP.MCPError.serverError(code: -32000, message: "nope") + } + ) + XCTFail("expected the server error to propagate") + } catch { + XCTAssertEqual(attempts.count, 1, "server errors must not be retried") + } + await pool.closeAll() + } + + // MARK: Private + + private static func key(fingerprint: Int = 0) -> MCPSessionKey { + MCPSessionKey( + endpoint: URL(string: "https://example.test/mcp")!, + credentialFingerprint: fingerprint, + clientName: "test", + clientVersion: "1.0" + ) + } +} + +// MARK: - MCPStreamingSessionTests + +final class MCPStreamingSessionTests: XCTestCase { + /// Streaming and non-streaming are different transports, so they + /// must not share a session. A caller that turned streaming on to + /// receive `tools/list_changed` would otherwise be handed a pooled + /// connection that cannot deliver notifications. + func testStreamingModeGetsItsOwnSession() async throws { + let counter = FakeMCPTransport.ConnectCounter() + let pool = MCPSessionPool() + let base = MCPSessionKey( + endpoint: URL(string: "https://example.test/mcp")!, + credentialFingerprint: 0, + clientName: "test", + clientVersion: "1.0", + streaming: false + ) + var streamingKey = base + streamingKey.streaming = true + + for key in [base, streamingKey, base] { + _ = try await pool.withClient( + key: key, + makeTransport: { FakeMCPTransport(counter: counter) }, + body: { client in try await client.listTools(cursor: nil).0.count } + ) + } + // Two sessions for three calls: the third reuses the first. + XCTAssertEqual(counter.count, 2) + await pool.closeAll() + } +}