From 009ca29dfe85c79da473b682f8ee3d0fafeb90cb Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 11:09:06 -0700
Subject: [PATCH 01/17] fix(mcp): keep one session per server instead of one
 per call
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`MCPClient` ran a full `initialize` handshake on every request and tore
the transport down afterwards, on the stated grounds that this "matches
MCP's session-per-call model".

MCP has no such model. It is a stateful session protocol — initialize,
operate, shut down — and Streamable HTTP carries an `Mcp-Session-Id`
across the session. What the per-call approach actually bought:

- A handshake on every tool call. 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. That fails
  quietly rather than erroring, which is the worst way to fail.

`MCPSessionPool` keeps one connected client per (endpoint, credential,
client identity) and evicts after three minutes idle. The credential is
in the key because the `Authorization` header is baked into the
transport at construction, so a refreshed OAuth token needs a new
transport and must not ride the old session. It is keyed by fingerprint
rather than by value so secrets stay out of dictionary keys.

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 retries connection
failures once, on a fresh transport. It deliberately does not retry
`serverError`: that is the server *answering*, and repeating a request
it already processed risks doing a non-idempotent thing twice.

`makeTransport` is a factory taking `any Transport` rather than a
concrete `HTTPClientTransport` value, because a retry needs a fresh
transport (a disconnected one cannot be reconnected) and because it
makes the pool testable. `FakeMCPTransport` is an in-process server
just complete enough to answer a handshake and `tools/list`, counting
every connect — the pool's whole job is deciding *when* to connect,
which is invisible against a live server and untestable against none.

Also fixes client identity. `clientName` defaulted to the literal
"Avyra" in a shared package and no call site overrode it, so every
Niora request introduced itself as Avyra to the server. It now reads
the host bundle's name, which fixes both apps without touching a
single call site.

387 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 .../WorkflowKit/Engine/MCP/MCPClient.swift    |  87 ++++++--
 .../Engine/MCP/MCPSessionPool.swift           | 179 +++++++++++++++
 .../MCPSessionPoolTests.swift                 | 205 ++++++++++++++++++
 3 files changed, 450 insertions(+), 21 deletions(-)
 create mode 100644 Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift
 create mode 100644 Tests/WorkflowKitTests/MCPSessionPoolTests.swift

diff --git a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
index f48e501..dac3f32 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,10 +39,15 @@ 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",
+        clientName: String = MCPClient.defaultClientName,
         clientVersion: String = "1.0"
     ) {
         self.serverURL = serverURL
@@ -47,6 +58,18 @@ public struct MCPClient: Sendable {
 
     // 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
@@ -262,26 +285,48 @@ 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<T: Sendable>(
-        _ 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
         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
+                ),
+                makeTransport: { Self.makeTransport(endpoint: endpoint, credential: credential) },
+                body: body
+            )
         } catch {
-            await client.disconnect()
             throw Self.mapError(error)
         }
     }
diff --git a/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift b/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift
new file mode 100644
index 0000000..adf1999
--- /dev/null
+++ b/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift
@@ -0,0 +1,179 @@
+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
+}
+
+// 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<T: Sendable>(
+        key: MCPSessionKey,
+        makeTransport: @Sendable () -> any Transport,
+        body: @Sendable (Client) async throws -> T
+    ) async throws -> T {
+        self.evictExpired()
+        do {
+            let client = try await self.connectedClient(for: key, makeTransport: makeTransport)
+            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)
+            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
+    ) async throws -> Client {
+        if let existing = self.sessions[key] {
+            return existing.client
+        }
+        let client = Client(name: key.clientName, version: key.clientVersion)
+        _ = 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/WorkflowKitTests/MCPSessionPoolTests.swift b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
new file mode 100644
index 0000000..3a748dc
--- /dev/null
+++ b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
@@ -0,0 +1,205 @@
+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]],
+                "serverInfo": ["name": "fake", "version": "1.0"],
+            ]
+        case "tools/list":
+            result = ["tools": [[
+                "name": "echo",
+                "description": "Echo input.",
+                "inputSchema": ["type": "object"] as [String: Any],
+            ]]]
+        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<Data, Swift.Error> {
+        AsyncThrowingStream { continuation in
+            self.continuation = continuation
+        }
+    }
+
+    // MARK: Private
+
+    private let counter: ConnectCounter
+    private let failSend: Bool
+    private var continuation: AsyncThrowingStream<Data, Swift.Error>.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"
+        )
+    }
+}

From 3446c68299acaddf35379e1e0e7cbe706bb040c7 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 11:12:34 -0700
Subject: [PATCH 02/17] feat(mcp): read resources and prompts, not just tools
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The client implemented `tools/*` and nothing else. Tool *results* could
carry an embedded resource, but resources could not be listed or read on
their own, and prompts were absent entirely — so a server's content was
invisible unless a tool happened to hand it over.

Adds `listResources`, `readResource`, and `listPrompts`, all draining
pagination the way `listTools` already did. Descriptors are plain values
with no SDK types in them, so persistence and UI never import `MCP`.

`MCPResourceDescriptor.isInteractiveUI` identifies an MCP App, and
requires *both* halves of the signal: a `ui://` URI and a
`profile=mcp-app` MIME parameter. Checking only the MIME type would
treat any server serving `text/html` as an interactive surface and hand
it a message channel it never asked for; checking only the scheme would
miss that the profile is what makes it an app.

Prompts land now rather than later because they are the shape Discover
recipes already want: a named, described entry point with typed
arguments.

Testing needed a seam. `MCPClient` built its own `HTTPClientTransport`
from the endpoint, so pagination draining, content-block conversion and
prompt-argument flattening could only be exercised against a live
third-party server — which is to say never in CI. An internal
initializer now accepts a transport factory, and the fake server
answers `resources/list` across *two* pages specifically so a client
that ignored `nextCursor` would fail the test rather than silently
returning half a server's resources.

391 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 .../WorkflowKit/Engine/MCP/MCPClient.swift    | 112 +++++++++++++++++-
 .../WorkflowKit/Engine/MCP/MCPResource.swift  | 102 ++++++++++++++++
 Tests/WorkflowKitTests/MCPResourceTests.swift |  71 +++++++++++
 .../MCPSessionPoolTests.swift                 |  36 +++++-
 4 files changed, 319 insertions(+), 2 deletions(-)
 create mode 100644 Sources/WorkflowKit/Engine/MCP/MCPResource.swift
 create mode 100644 Tests/WorkflowKitTests/MCPResourceTests.swift

diff --git a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
index dac3f32..6174155 100644
--- a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
+++ b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
@@ -54,6 +54,28 @@ public struct MCPClient: Sendable {
         self.credential = credential
         self.clientName = clientName
         self.clientVersion = clientVersion
+        self.transportFactory = nil
+    }
+
+    /// 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
+    ) {
+        self.serverURL = serverURL
+        self.credential = credential
+        self.clientName = clientName
+        self.clientVersion = clientVersion
+        self.transportFactory = transportFactory
     }
 
     // MARK: Public
@@ -138,12 +160,97 @@ 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 static func makeTransport(
         endpoint: URL,
@@ -315,6 +422,7 @@ public struct MCPClient: Sendable {
     ) async throws -> T {
         let endpoint = self.serverURL
         let credential = self.credential
+        let override = self.transportFactory
         do {
             return try await MCPSessionPool.shared.withClient(
                 key: MCPSessionKey(
@@ -323,7 +431,9 @@ public struct MCPClient: Sendable {
                     clientName: self.clientName,
                     clientVersion: self.clientVersion
                 ),
-                makeTransport: { Self.makeTransport(endpoint: endpoint, credential: credential) },
+                makeTransport: {
+                    override?() ?? Self.makeTransport(endpoint: endpoint, credential: credential)
+                },
                 body: body
             )
         } catch {
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/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, "<h1>hi</h1>")
+    }
+
+    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<T: Sendable>(
+        _ 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
index 3a748dc..0f83cd9 100644
--- a/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
+++ b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
@@ -72,7 +72,11 @@ actor FakeMCPTransport: Transport {
         case "initialize":
             result = [
                 "protocolVersion": "2025-06-18",
-                "capabilities": ["tools": [:] as [String: Any]],
+                "capabilities": [
+                    "tools": [:] as [String: Any],
+                    "resources": [:] as [String: Any],
+                    "prompts": [:] as [String: Any],
+                ],
                 "serverInfo": ["name": "fake", "version": "1.0"],
             ]
         case "tools/list":
@@ -81,6 +85,36 @@ actor FakeMCPTransport: Transport {
                 "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": "<h1>hi</h1>",
+            ]]]
+        case "prompts/list":
+            result = ["prompts": [[
+                "name": "summarise",
+                "description": "Summarise a document.",
+                "arguments": [["name": "uri", "description": "What to read", "required": true]],
+            ]]]
         default:
             result = [:]
         }

From 95a1195d9c9fdda28216b3daf99a4303318322b7 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 11:15:10 -0700
Subject: [PATCH 03/17] feat(charts): compile a constrained ChartIntent into
 Vega-Lite
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Agents should be able to say "this wants a chart" and have one appear.
Two constraints make that safe, and both are about what the model is
*not* allowed to do.

**It does not author the spec.** `ChartIntent` is an enum plus four
field names — roughly thirty tokens 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. An invalid chart type becomes
unrepresentable rather than a runtime surprise; given free text a small
model eventually emits `"linechart"` or a subtly wrong encoding block.
The codebase already reached this conclusion once — `BriefActionSuggestion`
is "kept simple so the model never has to author an `ActionTarget`".

**It does not supply the data.** The intent names a tool result and the
fields within it; the host injects 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. This makes
charts *structurally* grounded — no gate has to check them, because the
values never passed through the model.

Consequences that follow from the host owning the data:

- Field types are inferred, not declared. A field is quantitative only
  if every present value is numeric — one stray "n/a" degrades it to
  nominal rather than rendering a broken scale.
- Date detection is deliberately narrow. A false temporal turns a
  category axis into a broken time axis, which is worse than a date
  drawn as a category, so "3 items" stays nominal.
- A named field that does not exist throws, carrying the field names
  that *do*, since that is the likely failure and the recovery is
  telling the model the real names.

Tooltips are always on: an interactive chart the user cannot interrogate
is a picture, and that is the entire argument against Swift Charts here,
where every interaction would be hand-built.

Improving charts now happens in the compiler — no prompt change, no
model change. That separation is the payoff for keeping the output
small.

400 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/AriaTools/Charts/ChartIntent.swift    | 247 ++++++++++++++++++
 .../VegaLiteCompilerTests.swift               | 152 +++++++++++
 2 files changed, 399 insertions(+)
 create mode 100644 Sources/AriaTools/Charts/ChartIntent.swift
 create mode 100644 Tests/AriaToolsTests/VegaLiteCompilerTests.swift

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<String> {
+        rows.reduce(into: Set<String>()) { 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/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
+    }
+}

From 282d924b51762e4fc4a40bec31efc6dedd183fca Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 11:21:00 -0700
Subject: [PATCH 04/17] feat(charts): find the row set inside an arbitrary tool
 result
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A chart needs `[[String: JSONValue]]`. Tools return whatever shape their
author chose — a bare array, `{"items": [...]}`, something two levels
down, or JSON handed back as a string.

The model must not be what reshapes it. Asking a model to transcribe
rows into a chart call is asking it to invent them, and an invented
number wearing the authority of a chart is worse than an invented
sentence. So the host looks instead, and the rules are deliberately dull:

- Conventional envelope keys are tried in a fixed order, so a result
  carrying both `items` and `warnings` charts the items — and charts the
  same thing on every run rather than depending on dictionary ordering.
  Remaining keys are searched sorted, for the same reason.
- A JSON-in-a-string result is decoded, but only at the top level.
  Decoding strings found deep inside a payload would be guessing.
- A mixed array is *not* a table. Charting only its object-shaped
  elements would silently drop data, which is worse than declining.
- Rows are capped, and the original count is reported rather than
  swallowed — ten thousand points is an unreadable chart and a slow one,
  but the user should be told what they are looking at.

The one thing this refuses to do is guess creatively. A surprising rule
here draws a chart of the wrong data, and a chart of the wrong data is
believed.

410 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 .../AriaTools/Charts/TabularExtraction.swift  | 149 ++++++++++++++++++
 .../TabularExtractionTests.swift              |  97 ++++++++++++
 2 files changed, 246 insertions(+)
 create mode 100644 Sources/AriaTools/Charts/TabularExtraction.swift
 create mode 100644 Tests/AriaToolsTests/TabularExtractionTests.swift

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/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))])
+    }
+}

From 2536f81e6fde2055652cde647751c6503e9b9d9b Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 11:30:56 -0700
Subject: [PATCH 05/17] feat(mcp): optional streaming, so tools stop going
 stale
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

With `streaming: false` there is no long-lived listener, so
`notifications/tools/list_changed` can never arrive. That is why a
server changing its tools stays stale in the app until someone opens
Settings and taps refresh — the client had no way to be told.

Streaming is now per-server and off by default, because the existing
note is right that many servers do not support the GET SSE listener.
This has to degrade, not fail: a server without it keeps working exactly
as before.

The session key includes the streaming mode. Streaming and
non-streaming are different transports, not a setting on one, and
sharing a session between them would hand a caller that turned
streaming on specifically to receive notifications a pooled connection
that cannot deliver them.

Notification handlers register *before* connect. A server may emit
`tools/list_changed` immediately after initialize, and a handler
attached afterwards would miss exactly the case it exists for.

One implementation note worth keeping: building the `onConnect` closure
inline at the call site made the type checker fail outright — "failed to
produce diagnostic for expression". It is constructed up front with an
explicit type instead.

411 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 .../WorkflowKit/Engine/MCP/MCPClient.swift    | 45 ++++++++++++++++---
 .../Engine/MCP/MCPSessionPool.swift           | 25 +++++++++--
 .../MCPSessionPoolTests.swift                 | 33 ++++++++++++++
 3 files changed, 94 insertions(+), 9 deletions(-)

diff --git a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
index 6174155..6786c6c 100644
--- a/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
+++ b/Sources/WorkflowKit/Engine/MCP/MCPClient.swift
@@ -48,13 +48,17 @@ public struct MCPClient: Sendable {
         serverURL: URL,
         credential: MCPCredential? = nil,
         clientName: String = MCPClient.defaultClientName,
-        clientVersion: String = "1.0"
+        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
@@ -69,13 +73,17 @@ public struct MCPClient: Sendable {
         credential: MCPCredential?,
         clientName: String,
         clientVersion: String,
-        transportFactory: @escaping @Sendable () -> any Transport
+        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
@@ -251,15 +259,18 @@ public struct MCPClient: Sendable {
     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
@@ -423,17 +434,39 @@ public struct MCPClient: Sendable {
         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 {
             return try await MCPSessionPool.shared.withClient(
                 key: MCPSessionKey(
                     endpoint: endpoint,
                     credentialFingerprint: Self.fingerprint(credential),
                     clientName: self.clientName,
-                    clientVersion: self.clientVersion
+                    clientVersion: self.clientVersion,
+                    streaming: streaming
                 ),
                 makeTransport: {
-                    override?() ?? Self.makeTransport(endpoint: endpoint, credential: credential)
+                    override?() ?? Self.makeTransport(
+                        endpoint: endpoint,
+                        credential: credential,
+                        streaming: streaming
+                    )
                 },
+                onConnect: onConnect,
                 body: body
             )
         } catch {
diff --git a/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift b/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift
index adf1999..91cb385 100644
--- a/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift
+++ b/Sources/WorkflowKit/Engine/MCP/MCPSessionPool.swift
@@ -15,6 +15,11 @@ struct MCPSessionKey: Hashable {
     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
@@ -59,11 +64,16 @@ actor MCPSessionPool {
     func withClient<T: Sendable>(
         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)
+            let client = try await self.connectedClient(
+                for: key,
+                makeTransport: makeTransport,
+                onConnect: onConnect
+            )
             let value = try await body(client)
             self.sessions[key]?.touch()
             return value
@@ -72,7 +82,11 @@ actor MCPSessionPool {
                 throw error
             }
             await self.drop(key)
-            let client = try await self.connectedClient(for: key, makeTransport: makeTransport)
+            let client = try await self.connectedClient(
+                for: key,
+                makeTransport: makeTransport,
+                onConnect: onConnect
+            )
             let value = try await body(client)
             self.sessions[key]?.touch()
             return value
@@ -147,12 +161,17 @@ actor MCPSessionPool {
 
     private func connectedClient(
         for key: MCPSessionKey,
-        makeTransport: @Sendable () -> any Transport
+        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
diff --git a/Tests/WorkflowKitTests/MCPSessionPoolTests.swift b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
index 0f83cd9..7454e87 100644
--- a/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
+++ b/Tests/WorkflowKitTests/MCPSessionPoolTests.swift
@@ -237,3 +237,36 @@ final class MCPSessionPoolTests: XCTestCase {
         )
     }
 }
+
+// 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()
+    }
+}

From 525825c377cd123030397d4c1250f17fa7c0122f Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 14:48:36 -0700
Subject: [PATCH 06/17] fix(context): enforce the token ceiling on ranked
 tools, not just filler
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`maxTools` bounds how *many* tools are sent. `toolTokenLimit` bounds
what they cost. Only the first was enforced once ranking returned
anything:

    guard selected.count == requiredTools.count else {
        return selected        // no cost check, ever
    }

The ceiling was consulted for the "everything already fits" early exit
and for filler tools, then skipped entirely on the path almost every
real turn takes. Six tools is a fine cap until the six are MCP schemas.
A connected weather server (17 tools) put **5,845 tokens into a
4,096-token window** and the turn was refused outright — with the budget
reporting itself satisfied, because six is less than six.

That is the failure mode this whole layer exists to prevent, so it is
worth being precise about what was wrong: not the budget, not the
selector, not the consumer's configuration. The assembler counted tools
and never weighed them.

Ranked tools are now taken in order until the next one does not fit, and
the loop *stops* rather than skipping ahead to a smaller candidate.
Skipping would quietly prefer cheap tools over relevant ones, which
inverts the purpose of ranking.

Required tools are exempt. They are pinned or already invoked this turn,
and dropping one breaks the conversation outright, whereas overshooting
risks a refusal that trimming everything else may still avoid.

The regression test was checked against the unfixed code rather than
assumed: it selects 6 tools costing 2,130 against a ceiling of 1,198 and
fails, then passes once trimmed. A test for a budget bug that has never
been seen to fail is not evidence of anything.

413 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift  | 47 ++++++++++++-
 Tests/AriaTests/ToolBudgetCeilingTests.swift | 73 ++++++++++++++++++++
 2 files changed, 119 insertions(+), 1 deletion(-)
 create mode 100644 Tests/AriaTests/ToolBudgetCeilingTests.swift

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index bdd9160..0e6f45d 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -327,6 +327,46 @@ 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) }
+        for tool in selected where !requiredNames.contains(tool.name) {
+            let next = cost(tool)
+            guard spent + next <= ceiling else {
+                break
+            }
+            kept.append(tool)
+            spent += next
+        }
+        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 {
@@ -406,7 +446,12 @@ 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) }
+            )
         }
 
         // Nothing scored. How much to read into that depends on the
diff --git a/Tests/AriaTests/ToolBudgetCeilingTests.swift b/Tests/AriaTests/ToolBudgetCeilingTests.swift
new file mode 100644
index 0000000..9654354
--- /dev/null
+++ b/Tests/AriaTests/ToolBudgetCeilingTests.swift
@@ -0,0 +1,73 @@
+@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"
+        )
+    }
+
+    /// 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([:]) }
+        )
+    }
+}

From 4071d140e43796c4bbc0e5ca4803ffa7a5df5b11 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 15:17:52 -0700
Subject: [PATCH 07/17] feat(context): record which tools were offered, not
 just how many
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`toolsOffered` was a count. A count is enough to notice a surprise and
never enough to explain one.

A run asked for a weather summary with a weather server connected and
sent exactly one tool — `load_skill`, which was pinned, meaning the
ranker had scored *nothing* across twenty-seven candidates. From the
diagnostic alone there was no way to tell whether the ranker had buried
the weather tools or whether they had never been candidates at all.
Those have nothing in common except the symptom, and opposite fixes:
one is a ranking bug, the other is a tool that was never registered.

`ContextAllocation.offeredToolNames` carries the candidate list, so the
question is answerable by reading the export instead of by reasoning
backwards from a number.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift |  1 +
 Sources/Aria/Context/ContextBudget.swift    | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index 0e6f45d..e6e09a4 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -161,6 +161,7 @@ public struct DefaultContextAssembler: ContextAssembler {
             toolsOffered: tools.count,
             toolsSelected: selectedTools.count,
             selectedToolNames: selectedTools.map(\.name),
+            offeredToolNames: tools.map(\.name),
             messagesDropped: history.dropped,
             memoriesDropped: 0,
             toolResultsTruncated: truncated
diff --git a/Sources/Aria/Context/ContextBudget.swift b/Sources/Aria/Context/ContextBudget.swift
index 127f5a9..ff78bb5 100644
--- a/Sources/Aria/Context/ContextBudget.swift
+++ b/Sources/Aria/Context/ContextBudget.swift
@@ -152,6 +152,7 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
         toolsOffered: Int = 0,
         toolsSelected: Int = 0,
         selectedToolNames: [String] = [],
+        offeredToolNames: [String] = [],
         messagesDropped: Int = 0,
         memoriesDropped: Int = 0,
         toolResultsTruncated: Int = 0
@@ -164,6 +165,7 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
         self.toolsOffered = toolsOffered
         self.toolsSelected = toolsSelected
         self.selectedToolNames = selectedToolNames
+        self.offeredToolNames = offeredToolNames
         self.messagesDropped = messagesDropped
         self.memoriesDropped = memoriesDropped
         self.toolResultsTruncated = toolResultsTruncated
@@ -192,6 +194,16 @@ 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]
 
     /// History messages dropped to fit.
     public let messagesDropped: Int

From 233df3b5e44eb0c371e156f1a76cad16ddd5b435 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 15:50:51 -0700
Subject: [PATCH 08/17] feat(context): separate "the ranker found nothing" from
 "the budget cut it"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A weather question with a weather server connected sent one tool. The
export said 35 offered, 1 selected — and that is where the trail ended,
because two completely different failures produce that same pair of
numbers:

- ranking found nothing, so only the pinned tool survived; or
- ranking found the right tools and the token ceiling trimmed them away.

One is a retrieval problem, the other a budget problem, and the fix for
either is useless against the other. Three turns and a reproduction were
spent distinguishing them by hand, and the answer was still not
conclusive.

`ContextAllocation` now carries `rankedToolNames` — what ranking chose
*before* trimming — alongside `maxTools` and `toolTokenLimit`, the two
caps in force. Empty ranked names means retrieval; full ranked names
with a short selection means budget. No inference required.

`selectTools` returns the ranking beside the selection rather than
stashing it, since `DefaultContextAssembler` is a struct and scratch
state would not compile — and a tuple keeps the two facts adjacent,
which is how they are read.

Worth recording for the corpus: the assembler was reproduced against
the exact 35-tool surface from the field export, pinned `load_skill`,
`maxTools: 6`, `unrankedFillLimit: 0`, and it selected `load_skill`
plus five weather tools correctly. Whatever produced the field result
is not in this code path.

414 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift | 37 ++++++++++++++-------
 Sources/Aria/Context/ContextBudget.swift    | 18 ++++++++++
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index e6e09a4..1af16e8 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -115,7 +115,7 @@ public struct DefaultContextAssembler: ContextAssembler {
         state: AgentState,
         budget: ContextBudget
     ) async -> AssembledContext {
-        let selectedTools = await self.selectTools(
+        let (selectedTools, rankedNames) = await self.selectTools(
             from: tools,
             state: state,
             budget: budget
@@ -162,6 +162,9 @@ public struct DefaultContextAssembler: ContextAssembler {
             toolsSelected: selectedTools.count,
             selectedToolNames: selectedTools.map(\.name),
             offeredToolNames: tools.map(\.name),
+            rankedToolNames: rankedNames,
+            maxTools: budget.maxTools,
+            toolTokenLimit: budget.toolTokenLimit,
             messagesDropped: history.dropped,
             memoriesDropped: 0,
             toolResultsTruncated: truncated
@@ -378,13 +381,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
@@ -400,7 +409,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
@@ -414,7 +423,7 @@ 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 ranked = await self.selector.select(
@@ -424,6 +433,7 @@ public struct DefaultContextAssembler: ContextAssembler {
         )
         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.
         //
@@ -447,11 +457,14 @@ 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 Self.fitting(
-                selected,
-                required: requiredTools,
-                ceiling: ceiling,
-                cost: { self.tokenCounter.count(tool: $0.definition) }
+            return (
+                Self.fitting(
+                    selected,
+                    required: requiredTools,
+                    ceiling: ceiling,
+                    cost: { self.tokenCounter.count(tool: $0.definition) }
+                ),
+                rankedNames
             )
         }
 
@@ -459,7 +472,7 @@ public struct DefaultContextAssembler: ContextAssembler {
         // 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) }
@@ -479,7 +492,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 ff78bb5..f1deb65 100644
--- a/Sources/Aria/Context/ContextBudget.swift
+++ b/Sources/Aria/Context/ContextBudget.swift
@@ -153,6 +153,9 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
         toolsSelected: Int = 0,
         selectedToolNames: [String] = [],
         offeredToolNames: [String] = [],
+        rankedToolNames: [String] = [],
+        maxTools: Int? = nil,
+        toolTokenLimit: Int = 0,
         messagesDropped: Int = 0,
         memoriesDropped: Int = 0,
         toolResultsTruncated: Int = 0
@@ -166,6 +169,9 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
         self.toolsSelected = toolsSelected
         self.selectedToolNames = selectedToolNames
         self.offeredToolNames = offeredToolNames
+        self.rankedToolNames = rankedToolNames
+        self.maxTools = maxTools
+        self.toolTokenLimit = toolTokenLimit
         self.messagesDropped = messagesDropped
         self.memoriesDropped = memoriesDropped
         self.toolResultsTruncated = toolResultsTruncated
@@ -204,6 +210,18 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
     /// 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
 
     /// History messages dropped to fit.
     public let messagesDropped: Int

From fd7968850ab193977cfbc5fb9b10d7bb5c77bbe3 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 16:04:59 -0700
Subject: [PATCH 09/17] fix(context): never trim away the only tool that could
 answer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This morning's ceiling fix traded an overflow for a turn that cannot
act, which is the worse failure.

The field diagnostic, once it carried enough to be conclusive:

    RANKED before trim: [load_skill, weather_mcp__get_weather_summary,
                         weather_mcp__get_forecast, ...]
    SENT:               [load_skill]
    tool ceiling: 1198   tools sent: 118 tokens

Ranking was right — it put `get_weather_summary` second for "Get
weather summary". Every schema on that MCP server runs to roughly a
thousand tokens, over the 1,080 left after the pinned tool, so `fitting`
broke on the first candidate and returned the pinned tool alone. The
model was handed no way to fetch weather and answered from imagination,
filling a template with `[Current Date]` and `[Temperature Range]`.

The trim was correct by its own rule, and the rule was wrong. The
ceiling is a *share* of the budget — 40% by default — not the context
window. Honouring it exactly is right when it costs a fourth tool and
wrong when it costs the only one: no amount of budget discipline
redeems a request that cannot do the thing it was asked to do.

So when the share admits nothing, the top-ranked candidate is admitted
anyway. Deliberately the top-ranked one rather than the largest that
fits — preferring a cheap tool to a relevant one inverts ranking, and
this is a last resort, not a second policy.

Checked against the unfixed code rather than assumed: it reproduces the
field result exactly, `["load_skill"]` and nothing else.

Two things this leaves open, stated rather than buried. A thousand-token
tool schema is enormous, and six of them are what put 5,845 tokens into
a 4,096-token window earlier — capping or summarising oversized schemas
is the real fix and is not this one. And admitting an over-share tool
can still overflow a small window; it is strictly better than sending
nothing, and strictly worse than schemas that fit.

414 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift  | 24 ++++++++++++++++
 Tests/AriaTests/ToolBudgetCeilingTests.swift | 30 ++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index 1af16e8..c4feffb 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -368,6 +368,30 @@ public struct DefaultContextAssembler: ContextAssembler {
             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
     }
 
diff --git a/Tests/AriaTests/ToolBudgetCeilingTests.swift b/Tests/AriaTests/ToolBudgetCeilingTests.swift
index 9654354..6a328e1 100644
--- a/Tests/AriaTests/ToolBudgetCeilingTests.swift
+++ b/Tests/AriaTests/ToolBudgetCeilingTests.swift
@@ -35,6 +35,36 @@ final class ToolBudgetCeilingTests: XCTestCase {
         )
     }
 
+    /// 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))"
+        )
+    }
+
     /// Small tools are unaffected — the ceiling must not become a cap
     /// that starves ordinary surfaces.
     func testSmallToolSurfacesAreNotTrimmed() async throws {

From 1fe5654b405cef8e3f59a73f19e5a43adeb103c3 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 16:15:16 -0700
Subject: [PATCH 10/17] feat(context): compact oversized tool schemas instead
 of dropping them
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A published MCP tool schema is routinely a thousand tokens: a paragraph
of description plus a dozen documented parameters. On the 4,096-token
window an edge device actually has, that is a quarter of everything for
one tool the model may not even call. Six of them produced a 5,845-token
request that was refused outright, and the budget fix that followed
turned that into a turn with one tool and no way to answer.

Budgeting cannot solve this. A share that fits six such tools does not
exist. But most of those tokens are prose the model does not need in
order to make the call — so shrink the definition rather than choosing
between overflow and starvation.

`ToolDefinitionCompactor` applies the lightest level that fits:

  1. drop per-property descriptions — parameter names carry most of it
  2. also shorten the tool description to its first sentence
  3. also drop optional properties

The ordering is the design, and the invariant is that **a call the
model could make before must still validate after**. Names, types, the
required list and enum values survive every level: enums are not
documentation, they are the set of legal inputs, and a model guessing
outside one produces a call that fails. What is given up is guidance
first and optional capability last, because a tool called with worse
arguments is recoverable and a tool whose arguments no longer validate
is not.

`AnyTool.replacingDefinition` keeps the invocation closure untouched, so
a compacted definition can change what the model is *told* and never
what happens when it calls.

Verified in both directions, after a first attempt that passed
vacuously: with schemas sized like the real ones, removing the
compaction hook yields one tool and 1,390 tokens against a 1,198
ceiling — the field failure exactly — and restoring it yields several
tools under budget.

Still open: this shrinks what a server published, it does not make
servers publish less. A tool whose *required* parameters alone exceed
the budget is beyond it.

421 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift   |  25 ++-
 .../Context/ToolDefinitionCompactor.swift     | 180 ++++++++++++++++++
 Sources/Aria/Providers/Tool.swift             |  11 ++
 Tests/AriaTests/ToolBudgetCeilingTests.swift  |  65 +++++++
 .../ToolDefinitionCompactorTests.swift        | 131 +++++++++++++
 5 files changed, 408 insertions(+), 4 deletions(-)
 create mode 100644 Sources/Aria/Context/ToolDefinitionCompactor.swift
 create mode 100644 Tests/AriaTests/ToolDefinitionCompactorTests.swift

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index c4feffb..569d4ad 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -355,17 +355,26 @@ public struct DefaultContextAssembler: ContextAssembler {
         _ selected: [AnyTool],
         required: [AnyTool],
         ceiling: Int,
-        cost: (AnyTool) -> Int
+        cost: (AnyTool) -> Int,
+        compact: (AnyTool, Int) -> (AnyTool, ToolCompaction)
     ) -> [AnyTool] {
         let requiredNames = Set(required.map(\.name))
         var kept = required
         var spent = required.reduce(0) { $0 + cost($1) }
         for tool in selected where !requiredNames.contains(tool.name) {
-            let next = cost(tool)
+            // Try to make it fit before deciding it does not.
+            //
+            // A published MCP schema is mostly prose — a paragraph of
+            // description and a dozen documented parameters — and the
+            // model needs almost none of it to make the call. Dropping
+            // that first turns "no room for this tool" into "room for
+            // three", without removing anything the call depends on.
+            let (candidate, level) = compact(tool, ceiling - spent)
+            let next = cost(candidate)
             guard spent + next <= ceiling else {
                 break
             }
-            kept.append(tool)
+            kept.append(level == .none ? tool : candidate)
             spent += next
         }
 
@@ -486,7 +495,15 @@ public struct DefaultContextAssembler: ContextAssembler {
                     selected,
                     required: requiredTools,
                     ceiling: ceiling,
-                    cost: { self.tokenCounter.count(tool: $0.definition) }
+                    cost: { self.tokenCounter.count(tool: $0.definition) },
+                    compact: { tool, room in
+                        let (definition, level) = ToolDefinitionCompactor.compact(
+                            tool.definition,
+                            toFit: room,
+                            counter: self.tokenCounter
+                        )
+                        return (tool.replacingDefinition(definition), level)
+                    }
                 ),
                 rankedNames
             )
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/Tests/AriaTests/ToolBudgetCeilingTests.swift b/Tests/AriaTests/ToolBudgetCeilingTests.swift
index 6a328e1..66b07e7 100644
--- a/Tests/AriaTests/ToolBudgetCeilingTests.swift
+++ b/Tests/AriaTests/ToolBudgetCeilingTests.swift
@@ -65,6 +65,71 @@ final class ToolBudgetCeilingTests: XCTestCase {
         )
     }
 
+    /// The end-to-end case this was all for: six MCP-sized schemas on a
+    /// 4,096-token window. Uncompacted, one of them alone busts the
+    /// share and the model gets a single tool. Compacted, several fit.
+    func testCompactionLetsSeveralMCPToolsThrough() async throws {
+        let pinned = Self.tool(named: "load_skill", descriptionLength: 40)
+        let weather = (0 ..< 6).map { Self.mcpSizedTool(named: "weather_mcp__tool_\($0)") }
+        let assembler = DefaultContextAssembler(
+            pinnedToolNames: ["load_skill"],
+            unrankedFillLimit: 0
+        )
+        var state = AgentState()
+        state.messages = [.user("weather forecast for dublin")]
+
+        let budget = ContextBudget(total: 4096, reservedForOutput: 768, maxTools: 6)
+        let assembled = await assembler.assemble(
+            systemPrompt: "You are helpful.",
+            tools: [pinned] + weather,
+            state: state,
+            budget: budget
+        )
+
+        let sent = assembled.tools.filter { $0.name.hasPrefix("weather_mcp__") }
+        XCTAssertGreaterThan(
+            sent.count,
+            1,
+            "compaction bought nothing: \(assembled.tools.map(\.name))"
+        )
+        let counter = HeuristicTokenCounter()
+        let cost = assembled.tools.reduce(0) { $0 + counter.count(tool: $1.definition) }
+        XCTAssertLessThanOrEqual(cost, budget.toolTokenLimit)
+    }
+
+    /// 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([:]) }
+        )
+    }
+
     /// Small tools are unaffected — the ceiling must not become a cap
     /// that starves ordinary surfaces.
     func testSmallToolSurfacesAreNotTrimmed() async throws {
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."
+        )
+    )
+}

From 1bde96f484885472b84edbf32de7a3363dcc35a6 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 16:26:34 -0700
Subject: [PATCH 11/17] fix(context): tools may break the share, never the
 budget
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two rules composed into the worst possible outcome.

The share (`toolTokenLimit`, 40% of the budget) is deliberately
breakable: sending no usable tool is worse than overspending it, and a
turn that cannot act is a turn that answers from imagination.

But nothing bounded the result. A run selected three MCP tools worth
4,533 tokens against a 1,198 share and a 2,996 budget; `remaining` went
to zero, history was discarded to make room, and the provider refused
the request anyway at 4,727 of 4,096. The conversation was thrown away
*and* the turn failed — strictly worse than either failure alone.

`withinHardLimit` drops tools, lowest-ranked first, until the request
physically fits, recomposing the prompt 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.

So the hierarchy is now explicit. The share is advice. The budget is
not. Compaction tries to make both satisfiable before either has to
give.

Verified in both directions, after two vacuous attempts — the first
fixture compacted small enough to fit, the second was one incompressible
tool that still fit. It took a tool whose *required* surface alone
exceeds the whole budget to reproduce it: 10,401 tokens against 2,996
without the bound, within budget with it.

That is three tests today that looked green while asserting nothing.
Reverting the fix before trusting the test is cheap; believing a green
test that never failed is not.

422 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift  | 48 +++++++++++++++-
 Tests/AriaTests/ToolBudgetCeilingTests.swift | 60 ++++++++++++++++++++
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index 569d4ad..10cde94 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, rankedNames) = 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
@@ -404,6 +421,35 @@ public struct DefaultContextAssembler: ContextAssembler {
         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 {
diff --git a/Tests/AriaTests/ToolBudgetCeilingTests.swift b/Tests/AriaTests/ToolBudgetCeilingTests.swift
index 66b07e7..90ecf19 100644
--- a/Tests/AriaTests/ToolBudgetCeilingTests.swift
+++ b/Tests/AriaTests/ToolBudgetCeilingTests.swift
@@ -97,6 +97,26 @@ final class ToolBudgetCeilingTests: XCTestCase {
         XCTAssertLessThanOrEqual(cost, budget.toolTokenLimit)
     }
 
+    /// 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 {
@@ -130,6 +150,46 @@ final class ToolBudgetCeilingTests: XCTestCase {
         )
     }
 
+    /// 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 {

From 2d0d3adae367b2004ed37995eedc569ef76269c7 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 17:11:10 -0700
Subject: [PATCH 12/17] fix(tools): coerce model arguments to the schema they
 were given
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A 0.8B model found the right tool, formed the right intent, and picked
the right values — then the call failed on quotation marks:

    {"days": "1", "latitude": "56.35", "longitude": "-4.11"}
    → "Invalid days: must be a finite number, received string"

Nothing upstream was wrong. The Qwen parser preserves JSON types
faithfully; the model simply quoted its numbers, which small models do.
The schema already says what each field is, so the host can reconcile
that instead of asking the model to try again — the same division the
rest of this layer runs on: the model decides *what*, the host owns the
shape.

Only lossless conversions. `"1"` becomes `1`; `"next week"` stays a
string and the server returns its own error, which is more useful than a
number this invented. Specifically not done:

- `1.7` is never truncated to an integer — that changes the request.
- `"yes"` is not read as `true`. It is a guess about intent rather than
  a reading of the value, and JSON has two spellings for a reason.
- A field the schema does not describe passes through untouched; there
  is no contract to enforce and inventing one is worse than nothing.

Also handles the mirror case — a bare number where the schema declares a
string — since it costs nothing and fails the same way.

422 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Agent/AgentToolExecution.swift   |  14 +-
 .../Aria/Context/ToolArgumentCoercion.swift   | 147 ++++++++++++++++++
 .../AriaTests/ToolArgumentCoercionTests.swift | 125 +++++++++++++++
 3 files changed, 285 insertions(+), 1 deletion(-)
 create mode 100644 Sources/Aria/Context/ToolArgumentCoercion.swift
 create mode 100644 Tests/AriaTests/ToolArgumentCoercionTests.swift

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/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/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"]
+    )
+}

From 4c36f140ac42622c374b2a293f50e5d9dd027561 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 17:32:49 -0700
Subject: [PATCH 13/17] fix(context): stop compacting definitions the provider
 will not send
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Compaction shrank what the assembler *counted* and not what the
provider *sent*, which is worse than not compacting at all.

`FoundationModelsProvider` builds its typed tools from factory closures
handed to it at construction. Those closures captured their description
then. Swapping a definition afterwards reaches the `AnyTool` the
assembler returns and cannot reach the tool the provider registers — so
the budget was costing compacted tools while the request carried full
ones.

Measured, from the field: a turn priced at 1,366 estimated tokens was
refused at 5,362 actual — 3.9x under, with six tools trimmed to exactly
the 1,198 ceiling and six full ones sent.

An accurate budget with fewer tools beats an inaccurate budget with
more. A refused turn sends none at all.

`ToolDefinitionCompactor` stays, with its tests. It is correct and it
will be useful to a provider that builds its tool list from
`AssembledContext.tools` rather than from closures fixed at
construction — which is the real fix, and is not a one-line change to
make safely.

Also documents the constraint at both ends: the provider's filter now
says why it can only match on name, and `fitting` says why it does not
compact. This cost a full debugging cycle to find; the next person
should not have to repeat it.

429 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift   | 45 ++++++++++---------
 .../Providers/FoundationModelsProvider.swift  | 10 +++++
 Tests/AriaTests/ToolBudgetCeilingTests.swift  | 32 -------------
 3 files changed, 34 insertions(+), 53 deletions(-)

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index 10cde94..18828bb 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -372,26 +372,37 @@ public struct DefaultContextAssembler: ContextAssembler {
         _ selected: [AnyTool],
         required: [AnyTool],
         ceiling: Int,
-        cost: (AnyTool) -> Int,
-        compact: (AnyTool, Int) -> (AnyTool, ToolCompaction)
+        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) {
-            // Try to make it fit before deciding it does not.
-            //
-            // A published MCP schema is mostly prose — a paragraph of
-            // description and a dozen documented parameters — and the
-            // model needs almost none of it to make the call. Dropping
-            // that first turns "no room for this tool" into "room for
-            // three", without removing anything the call depends on.
-            let (candidate, level) = compact(tool, ceiling - spent)
-            let next = cost(candidate)
+            let next = cost(tool)
             guard spent + next <= ceiling else {
                 break
             }
-            kept.append(level == .none ? tool : candidate)
+            kept.append(tool)
             spent += next
         }
 
@@ -541,15 +552,7 @@ public struct DefaultContextAssembler: ContextAssembler {
                     selected,
                     required: requiredTools,
                     ceiling: ceiling,
-                    cost: { self.tokenCounter.count(tool: $0.definition) },
-                    compact: { tool, room in
-                        let (definition, level) = ToolDefinitionCompactor.compact(
-                            tool.definition,
-                            toFit: room,
-                            counter: self.tokenCounter
-                        )
-                        return (tool.replacingDefinition(definition), level)
-                    }
+                    cost: { self.tokenCounter.count(tool: $0.definition) }
                 ),
                 rankedNames
             )
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/Tests/AriaTests/ToolBudgetCeilingTests.swift b/Tests/AriaTests/ToolBudgetCeilingTests.swift
index 90ecf19..0e510c3 100644
--- a/Tests/AriaTests/ToolBudgetCeilingTests.swift
+++ b/Tests/AriaTests/ToolBudgetCeilingTests.swift
@@ -65,38 +65,6 @@ final class ToolBudgetCeilingTests: XCTestCase {
         )
     }
 
-    /// The end-to-end case this was all for: six MCP-sized schemas on a
-    /// 4,096-token window. Uncompacted, one of them alone busts the
-    /// share and the model gets a single tool. Compacted, several fit.
-    func testCompactionLetsSeveralMCPToolsThrough() async throws {
-        let pinned = Self.tool(named: "load_skill", descriptionLength: 40)
-        let weather = (0 ..< 6).map { Self.mcpSizedTool(named: "weather_mcp__tool_\($0)") }
-        let assembler = DefaultContextAssembler(
-            pinnedToolNames: ["load_skill"],
-            unrankedFillLimit: 0
-        )
-        var state = AgentState()
-        state.messages = [.user("weather forecast for dublin")]
-
-        let budget = ContextBudget(total: 4096, reservedForOutput: 768, maxTools: 6)
-        let assembled = await assembler.assemble(
-            systemPrompt: "You are helpful.",
-            tools: [pinned] + weather,
-            state: state,
-            budget: budget
-        )
-
-        let sent = assembled.tools.filter { $0.name.hasPrefix("weather_mcp__") }
-        XCTAssertGreaterThan(
-            sent.count,
-            1,
-            "compaction bought nothing: \(assembled.tools.map(\.name))"
-        )
-        let counter = HeuristicTokenCounter()
-        let cost = assembled.tools.reduce(0) { $0 + counter.count(tool: $1.definition) }
-        XCTAssertLessThanOrEqual(cost, budget.toolTokenLimit)
-    }
-
     /// 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 {

From db826a5390b67ce646ed7f6f3304b159b7558fae Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 19:26:36 -0700
Subject: [PATCH 14/17] feat(testing): audit the token estimate against the
 request it describes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The estimate is the most load-bearing number in the context layer —
every trim, cap and drop is arithmetic on top of it — and nothing
verified it. When it was wrong the symptom appeared somewhere else
entirely: a provider refusing a turn the budget had called comfortable,
with the budget still reporting itself satisfied afterwards.

`ContextAudit` renders the request the way a chat template will and
compares that against what the counter charged. Rendering first is the
point: a counter that reads `ToolDefinition` objects cannot see a schema
counted once as a schema and sent again inside a description, and that
double-shipping priced as prose is a real error this catches.

Assertions are one-sided, because the directions are not symmetric.
Over-estimating wastes budget and sends fewer tools; under-estimating
overflows the window and loses the turn.

What it does **not** catch, stated plainly: it renders the tools the
assembler produced, so it cannot see a provider sending something else.
That was the 1,366-vs-5,362 failure — definitions compacted after the
provider had captured the originals — and both sides of this audit would
have agreed while the request was four times larger. Closing that needs
the provider's own count, which is the companion change in the app.

The reference counter is a character heuristic, so this catches
order-of-magnitude errors and not small ones. It takes a closure so a
real tokenizer can be supplied where one exists.

434 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/AriaTesting/ContextAudit.swift  | 174 ++++++++++++++++++++++++
 Tests/AriaTests/ContextAuditTests.swift | 149 ++++++++++++++++++++
 2 files changed, 323 insertions(+)
 create mode 100644 Sources/AriaTesting/ContextAudit.swift
 create mode 100644 Tests/AriaTests/ContextAuditTests.swift

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/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"]
+        )
+    )
+}

From 3df63e02216ace6945370f203603e8cf50a20414 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 20:51:43 -0700
Subject: [PATCH 15/17] test(context): corpus for the failure the evals never
 measured
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.

From the field: asked "What's the tracking number?", the ranker offered
`calculator`, the model called it with `(4.45 - 4.45) / 16` — numbers
scavenged out of a fasting summary earlier in the conversation — and
answered `0`.

Measured against a real 34-tool device surface, the pattern is exact.
Everything that fails does so on one common word matching one tool:
"number" appears in `calculator`, `uuid` and `unit_converter`; "fact"
appears in `remember_fact`. Everything that succeeds matches on
distinctive terms, usually more than one. Conversational turns —
"Thanks!", "How are you?", "Explain quantum computing" — already rank
nothing and now stay that way.

The two failures are recorded as `XCTExpectFailure` rather than fixed
here, and that is deliberate. The obvious rule — require two matched
terms — breaks "remind me to buy milk", where one matched term is the
whole intent. Picking a threshold against the four cases that happened
to motivate it is how the stopword list came to delete the subject of
every fasting query, twice.

So this states the gap instead of guessing at it. The other half of the
corpus is the constraint any future floor has to satisfy: four real
requests that must keep ranking their tool first.

435 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 .../AriaTests/ToolRelevanceCorpusTests.swift  | 114 ++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 Tests/AriaTests/ToolRelevanceCorpusTests.swift

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: [])
+        )
+    }
+}

From 3296face914a4eae3c2af44c4fb93e48f8dc3ccd Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 21:28:30 -0700
Subject: [PATCH 16/17] fix(context): make images visible to the context layer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Attached images were invisible in two independent ways, and each alone
is enough to break a turn.

**They cost nothing.** `count(message:)` measured `textContent` and the
tool calls, so an image contributed zero. Everything downstream is
arithmetic on that number: history windowing never dropped an image
message because it looked free, and the request overflowed with the
budget reporting itself satisfied. Images are now charged
`tokensPerImage`, a deliberately large constant — a 1024x1024 image runs
to roughly a thousand tokens on the Qwen-VL family, and nothing about
that is derivable from the bytes we hold. Over-charging costs a few
tools; under-charging costs the turn.

**A photo with no caption got no tools at all.** The ranking query is
the latest user text, which is empty for an image-only turn. The
selector returns nothing for an empty query, `unrankedFillLimit: 0`
reads that as a decision, and the model received an image with no way to
act on it.

Those are different situations. "Nothing matched" is a ranker verdict;
an empty query means the ranker was never given a question. The empty
case now falls through to what fits, which is the same reading the
definition-only provider path already takes.

A captioned image is unaffected — it still ranks on its caption, which
has its own test, because a fallback that swallowed the normal path
would be worse than the bug.

Both tests were vacuous on the first attempt and were rewritten until
they failed against the unfixed code. The trap each time: a surface
smaller than `maxTools` short-circuits the assembler and returns
everything unranked, so a selection test with four tools and a cap of
six proves nothing. That is the fourth time today.

441 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift |  22 +++-
 Sources/Aria/Context/TokenCounter.swift     |  22 ++++
 Tests/AriaTests/ImageContextTests.swift     | 114 ++++++++++++++++++++
 3 files changed, 157 insertions(+), 1 deletion(-)
 create mode 100644 Tests/AriaTests/ImageContextTests.swift

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index 18828bb..72144e2 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -516,9 +516,29 @@ public struct DefaultContextAssembler: ContextAssembler {
             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 })
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/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")
+    }
+}

From d306a6d09206d09ae9d06f367077a78af67f6620 Mon Sep 17 00:00:00 2001
From: Prasad Pamidi <prasad.pamidi01@gmail.com>
Date: Sun, 9 Aug 2026 23:14:52 -0700
Subject: [PATCH 17/17] feat(context): record how many images reached the
 request

Asked "What's the tracking number" over a receipt photo, a model
answered from the caption alone and invented coordinates in the Bay of
Biscay to call a weather tool with. The diagnostic could not say whether
the image had reached it: the allocation described tokens, tools and
caps, and was silent on the one thing that mattered.

`ContextAllocation.imageCount` closes that. Zero with an attachment on
screen localises the loss to assembly; non-zero moves it downstream to
the provider or the model. Without it the question needs a code read
and a guess, which is how the last four of these went.

441 tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013UQx5NRp1nBm1Y7sHRJu5L
---
 Sources/Aria/Context/ContextAssembler.swift |  8 ++++++++
 Sources/Aria/Context/ContextBudget.swift    | 11 +++++++++++
 2 files changed, 19 insertions(+)

diff --git a/Sources/Aria/Context/ContextAssembler.swift b/Sources/Aria/Context/ContextAssembler.swift
index 72144e2..d0ec9c2 100644
--- a/Sources/Aria/Context/ContextAssembler.swift
+++ b/Sources/Aria/Context/ContextAssembler.swift
@@ -182,6 +182,14 @@ public struct DefaultContextAssembler: ContextAssembler {
             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
diff --git a/Sources/Aria/Context/ContextBudget.swift b/Sources/Aria/Context/ContextBudget.swift
index f1deb65..8580b08 100644
--- a/Sources/Aria/Context/ContextBudget.swift
+++ b/Sources/Aria/Context/ContextBudget.swift
@@ -156,6 +156,7 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
         rankedToolNames: [String] = [],
         maxTools: Int? = nil,
         toolTokenLimit: Int = 0,
+        imageCount: Int = 0,
         messagesDropped: Int = 0,
         memoriesDropped: Int = 0,
         toolResultsTruncated: Int = 0
@@ -172,6 +173,7 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
         self.rankedToolNames = rankedToolNames
         self.maxTools = maxTools
         self.toolTokenLimit = toolTokenLimit
+        self.imageCount = imageCount
         self.messagesDropped = messagesDropped
         self.memoriesDropped = memoriesDropped
         self.toolResultsTruncated = toolResultsTruncated
@@ -222,6 +224,15 @@ public struct ContextAllocation: Sendable, Equatable, Codable {
     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