Harden model-facing decoding against common LLM tool-call quirks (0.2.2)

Sources/CodeMode/API/BridgeModels.swift

@@ -151,6 +151,73 @@ public struct JavaScriptExecutionRequest: Sendable, Codable, Equatable {
         self.timeoutMs = timeoutMs
         self.context = context
     }
+
+    private enum CodingKeys: String, CodingKey {
+        case code
+        case allowedCapabilities
+        case allowedCapabilityKeys
+        case timeoutMs
+        case context
+    }
+
+    // Custom decoding tolerant of common LLM tool-call quirks, and aligned with the
+    // advertised schema where only `code` and `allowedCapabilities` are required:
+    //  - `allowedCapabilityKeys`, `timeoutMs`, and `context` may be omitted (the
+    //    synthesized decoder wrongly required all three, so an otherwise-valid call
+    //    that left them out was rejected).
+    //  - `timeoutMs` accepts a JSON number or a numeric string ("10000").
+    //  - an unknown capability ID in `allowedCapabilities` produces a clear,
+    //    actionable error naming the offending value(s) instead of an opaque
+    //    Codable failure.
+    public init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+
+        self.code = try container.decode(String.self, forKey: .code)
+
+        let rawCapabilities = try container.decode([String].self, forKey: .allowedCapabilities)
+        var capabilities: [CapabilityID] = []
+        var unknown: [String] = []
+        for raw in rawCapabilities {
+            if let id = CapabilityID(rawValue: raw) {
+                capabilities.append(id)
+            } else {
+                unknown.append(raw)
+            }
+        }
+        guard unknown.isEmpty else {
+            throw DecodingError.dataCorruptedError(
+                forKey: .allowedCapabilities,
+                in: container,
+                debugDescription: "Unknown capability ID(s): \(unknown.joined(separator: ", ")). Use searchJavaScriptAPI to discover valid capability IDs."
+            )
+        }
+        self.allowedCapabilities = capabilities
+
+        self.allowedCapabilityKeys = try container.decodeIfPresent([CodeModeCapabilityKey].self, forKey: .allowedCapabilityKeys) ?? []
+        self.timeoutMs = try Self.decodeTimeoutMs(from: container) ?? 10_000
+        self.context = try container.decodeIfPresent(ExecutionContext.self, forKey: .context) ?? .init()
+    }
+
+    private static func decodeTimeoutMs(from container: KeyedDecodingContainer<CodingKeys>) throws -> Int? {
+        guard container.contains(.timeoutMs), try !container.decodeNil(forKey: .timeoutMs) else {
+            return nil
+        }
+        if let value = try? container.decode(Int.self, forKey: .timeoutMs) {
+            return value
+        }
+        if let value = try? container.decode(Double.self, forKey: .timeoutMs) {
+            return Int(value)
+        }
+        if let string = try? container.decode(String.self, forKey: .timeoutMs),
+           let value = Double(string.trimmingCharacters(in: .whitespaces)) {
+            return Int(value)
+        }
+        throw DecodingError.dataCorruptedError(
+            forKey: .timeoutMs,
+            in: container,
+            debugDescription: "timeoutMs must be an integer or a numeric string."
+        )
+    }
 }
 
 public enum JavaScriptExecutionEvent: Sendable, Equatable {

Sources/CodeMode/Registry/CapabilityRegistry.swift

@@ -559,14 +559,70 @@ public final class CapabilityRegistry: @unchecked Sendable {
     }
 
     private func invoke(_ function: RegisteredCodeModeFunction, arguments: [String: JSONValue], context: BridgeInvocationContext) throws -> JSONValue {
-        try validateArguments(arguments, for: function)
+        // Coerce common LLM output quirks (a number sent as a string, a bool sent
+        // as "true"/"false") into the declared type before validation, so both the
+        // type check and the bridge see canonical values instead of rejecting the
+        // call. Only declared `.number`/`.bool` arguments are touched.
+        let normalized = normalizedArguments(arguments, for: function)
+        try validateArguments(normalized, for: function)
         try validatePermissions(function.requiredPermissions, context: context)
         try context.checkCancellation()
         return try mapCodeModeFunctionError {
-            try function.handler(arguments, context)
+            try function.handler(normalized, context)
         }
     }
 
+    private func normalizedArguments(_ arguments: [String: JSONValue], for function: RegisteredCodeModeFunction) -> [String: JSONValue] {
+        var result = arguments
+        for (path, expectedType) in function.argumentTypes where expectedType == .number || expectedType == .bool {
+            guard let value = value(atPath: path, in: result),
+                  let coerced = Self.coerceScalarString(value, to: expectedType)
+            else {
+                continue
+            }
+            result = Self.setValue(coerced, atPath: path, in: result)
+        }
+        return result
+    }
+
+    /// Coerces a JSON string into the declared scalar type. JSON numbers/bools are
+    /// already the right type and are left untouched (returns nil). Booleans accept
+    /// only case-insensitive "true"/"false" — not 1/0 or "yes"/"no" — so safety
+    /// gates like `confirmed` still require an explicit affirmative.
+    private static func coerceScalarString(_ value: JSONValue, to type: CapabilityArgumentType) -> JSONValue? {
+        guard case let .string(raw) = value else { return nil }
+        let trimmed = raw.trimmingCharacters(in: .whitespaces)
+        switch type {
+        case .number:
+            guard let number = Double(trimmed) else { return nil }
+            return .number(number)
+        case .bool:
+            switch trimmed.lowercased() {
+            case "true": return .bool(true)
+            case "false": return .bool(false)
+            default: return nil
+            }
+        default:
+            return nil
+        }
+    }
+
+    /// Replaces the value at a dotted path (e.g. `options.timeoutMs`), rebuilding
+    /// the intermediate objects. Missing intermediates are treated as empty objects.
+    private static func setValue(_ newValue: JSONValue, atPath path: String, in root: [String: JSONValue]) -> [String: JSONValue] {
+        let segments = path.split(separator: ".").map(String.init)
+        guard let first = segments.first else { return root }
+        var result = root
+        if segments.count == 1 {
+            result[first] = newValue
+        } else {
+            let rest = segments.dropFirst().joined(separator: ".")
+            let child = result[first]?.objectValue ?? [:]
+            result[first] = .object(setValue(newValue, atPath: rest, in: child))
+        }
+        return result
+    }
+
     private func mapCodeModeFunctionError(_ body: () throws -> JSONValue) throws -> JSONValue {
         do {
             return try body()

Tests/CodeModeTests/CapabilityRegistryTests.swift

@@ -431,9 +431,12 @@ private func jsNames(for capability: CapabilityID) -> [String] {
     defer { cleanup(sandbox) }
 
     do {
+        // A bool where a number is expected: a genuine type mismatch that is not
+        // one of the coerced quirks (numeric string / "true"/"false"), so it must
+        // still be rejected.
         _ = try registry.invoke(
             "weather.read",
-            arguments: ["latitude": .string("37.0"), "longitude": .number(-122.0)],
+            arguments: ["latitude": .bool(true), "longitude": .number(-122.0)],
             context: context
         )
         Issue.record("Expected type mismatch validation to throw")

Tests/CodeModeTests/DefensiveDecodingTests.swift

@@ -0,0 +1,150 @@
+import Foundation
+import Testing
+@testable import CodeMode
+
+// MARK: - JavaScriptExecutionRequest decoding (Fix 2/3)
+
+private func decodeRequest(_ json: String) throws -> JavaScriptExecutionRequest {
+    try JSONDecoder().decode(JavaScriptExecutionRequest.self, from: Data(json.utf8))
+}
+
+@Test func executionRequestDecodesTimeoutFromNumericString() throws {
+    // The exact shape that was failing in the field: timeoutMs quoted, and
+    // allowedCapabilityKeys / context omitted.
+    let request = try decodeRequest(#"""
+    {"allowedCapabilities":["calendar.write"],"code":"return 1;","timeoutMs":"10000"}
+    """#)
+    #expect(request.timeoutMs == 10_000)
+    #expect(request.allowedCapabilities == [.calendarWrite])
+    #expect(request.allowedCapabilityKeys.isEmpty)
+}
+
+@Test func executionRequestDecodesTimeoutFromNumber() throws {
+    let request = try decodeRequest(#"{"code":"return 1;","allowedCapabilities":[],"timeoutMs":5000}"#)
+    #expect(request.timeoutMs == 5_000)
+}
+
+@Test func executionRequestTreatsOmittedOptionalsAsDefaults() throws {
+    // Only the two schema-required fields present.
+    let request = try decodeRequest(#"{"code":"return 1;","allowedCapabilities":["fs.read"]}"#)
+    #expect(request.timeoutMs == 10_000)
+    #expect(request.allowedCapabilityKeys.isEmpty)
+    #expect(request.context == ExecutionContext())
+    #expect(request.allowedCapabilities == [.fsRead])
+}
+
+@Test func executionRequestTreatsNullTimeoutAsDefault() throws {
+    let request = try decodeRequest(#"{"code":"return 1;","allowedCapabilities":[],"timeoutMs":null}"#)
+    #expect(request.timeoutMs == 10_000)
+}
+
+@Test func executionRequestReportsUnknownCapabilityClearly() {
+    do {
+        _ = try decodeRequest(#"{"code":"return 1;","allowedCapabilities":["calendar.write","calendarWrite","not.a.cap"]}"#)
+        Issue.record("Expected decoding to throw on unknown capability IDs")
+    } catch {
+        let message = "\(error)"
+        // The good spelling must not appear; the two bad ones must, by name.
+        #expect(message.contains("calendarWrite"))
+        #expect(message.contains("not.a.cap"))
+    }
+}
+
+@Test func executionRequestRejectsNonNumericTimeout() {
+    #expect(throws: (any Error).self) {
+        _ = try decodeRequest(#"{"code":"return 1;","allowedCapabilities":[],"timeoutMs":"soon"}"#)
+    }
+}
+
+@Test func executionRequestRoundTripsThroughCodable() throws {
+    let original = JavaScriptExecutionRequest(
+        code: "return 1;",
+        allowedCapabilities: [.fsRead, .networkFetch],
+        allowedCapabilityKeys: ["custom.tool"],
+        timeoutMs: 7_500,
+        context: ExecutionContext(userID: "u", sessionID: "s")
+    )
+    let data = try JSONEncoder().encode(original)
+    let decoded = try JSONDecoder().decode(JavaScriptExecutionRequest.self, from: data)
+    #expect(decoded == original)
+}
+
+// MARK: - Registry argument coercion (Fix 1)
+
+/// A synthetic capability whose handler echoes the (already-normalized) arguments
+/// it receives, so tests can observe coercion end-to-end.
+private func makeEchoRegistry(argumentTypes: [String: CapabilityArgumentType]) -> CapabilityRegistry {
+    let descriptor = CapabilityDescriptor(
+        id: .fsRead,
+        title: "Echo",
+        summary: "Echoes received arguments",
+        tags: ["test"],
+        example: "noop",
+        optionalArguments: Array(argumentTypes.keys.filter { $0.contains(".") == false }),
+        argumentTypes: argumentTypes
+    )
+    return CapabilityRegistry(
+        registrations: [
+            CapabilityRegistration(descriptor: descriptor) { args, _ in .object(args) }
+        ]
+    )
+}
+
+@Test func registryCoercesNumericStringToNumber() throws {
+    let registry = makeEchoRegistry(argumentTypes: ["count": .number])
+    let (context, sandbox) = try makeInvocationContext(allowedCapabilities: [.fsRead])
+    defer { cleanup(sandbox) }
+
+    let result = try registry.invoke(
+        CapabilityID.fsRead.rawValue,
+        arguments: ["count": .string(" 20 ")],
+        context: context
+    )
+    #expect(result.objectValue?["count"] == .number(20))
+}
+
+@Test func registryCoercesBoolStringsCaseInsensitively() throws {
+    let registry = makeEchoRegistry(argumentTypes: ["flag": .bool])
+    let (context, sandbox) = try makeInvocationContext(allowedCapabilities: [.fsRead])
+    defer { cleanup(sandbox) }
+
+    for (input, expected) in [("true", true), ("TRUE", true), ("False", false)] {
+        let result = try registry.invoke(
+            CapabilityID.fsRead.rawValue,
+            arguments: ["flag": .string(input)],
+            context: context
+        )
+        #expect(result.objectValue?["flag"] == .bool(expected))
+    }
+}
+
+@Test func registryDoesNotCoerceNonBooleanStringsOrBadNumbers() throws {
+    let registry = makeEchoRegistry(argumentTypes: ["count": .number, "flag": .bool])
+    let (context, sandbox) = try makeInvocationContext(allowedCapabilities: [.fsRead])
+    defer { cleanup(sandbox) }
+
+    // "yes"/"1"/"0" are intentionally NOT booleans, so validation still rejects them
+    // (safety gates like `confirmed` keep requiring an explicit true/false).
+    for bad in ["yes", "1", "0"] {
+        #expect(throws: (any Error).self) {
+            _ = try registry.invoke(CapabilityID.fsRead.rawValue, arguments: ["flag": .string(bad)], context: context)
+        }
+    }
+    // A non-numeric string stays rejected too.
+    #expect(throws: (any Error).self) {
+        _ = try registry.invoke(CapabilityID.fsRead.rawValue, arguments: ["count": .string("abc")], context: context)
+    }
+}
+
+@Test func registryCoercesNestedDottedNumericString() throws {
+    let registry = makeEchoRegistry(argumentTypes: ["options": .object, "options.timeoutMs": .number])
+    let (context, sandbox) = try makeInvocationContext(allowedCapabilities: [.fsRead])
+    defer { cleanup(sandbox) }
+
+    let result = try registry.invoke(
+        CapabilityID.fsRead.rawValue,
+        arguments: ["options": .object(["timeoutMs": .string("5000")])],
+        context: context
+    )
+    #expect(result.objectValue?["options"]?.objectValue?["timeoutMs"] == .number(5000))
+}