From 6000f02a6750089f282d78761f32f20b12ed95cd Mon Sep 17 00:00:00 2001 From: Zac White Date: Sun, 12 Jul 2026 23:20:25 -0700 Subject: [PATCH] Harden model-facing decoding against common LLM tool-call quirks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defensive fixes so ordinary model output stops nuking whole calls: 1. Registry coerces declared .number/.bool arguments from strings before validation: a numeric string ("20") becomes a number and "true"/"false" (case-insensitive only — not 1/0 or "yes") becomes a bool, at every declared argument including nested dotted paths. Both the type check and the bridge then see canonical values. Safety gates like "confirmed" still require an explicit true/false. 2. JavaScriptExecutionRequest gets a tolerant init(from:): timeoutMs accepts a number or a numeric string, and — matching the advertised schema where only code + allowedCapabilities are required — allowedCapabilityKeys, timeoutMs, and context may be omitted. The synthesized decoder wrongly required all three, so a valid call that left them out (as the field report did) was rejected before any JS ran. 3. An unknown/misspelled capability ID in allowedCapabilities now yields a clear error naming the offending value(s) and pointing at searchJavaScriptAPI, instead of an opaque Codable failure. Audited the other model-facing Codable input (JavaScriptAPISearchRequest): just a code string, nothing to harden. Adds DefensiveDecodingTests (11 cases) and retargets registryValidationRejectsWrongArgumentType at a bool (a genuine type mismatch, not a coerced quirk). 241 tests pass; deterministic evals 49/49. --- Sources/CodeMode/API/BridgeModels.swift | 67 ++++++++ .../Registry/CapabilityRegistry.swift | 60 ++++++- .../CapabilityRegistryTests.swift | 5 +- .../DefensiveDecodingTests.swift | 150 ++++++++++++++++++ Tools/CodeModeEval/Package.resolved | 4 +- 5 files changed, 281 insertions(+), 5 deletions(-) create mode 100644 Tests/CodeModeTests/DefensiveDecodingTests.swift diff --git a/Sources/CodeMode/API/BridgeModels.swift b/Sources/CodeMode/API/BridgeModels.swift index 9127d1f..e922dd9 100644 --- a/Sources/CodeMode/API/BridgeModels.swift +++ b/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) 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 { diff --git a/Sources/CodeMode/Registry/CapabilityRegistry.swift b/Sources/CodeMode/Registry/CapabilityRegistry.swift index 7e7d489..b50f76d 100644 --- a/Sources/CodeMode/Registry/CapabilityRegistry.swift +++ b/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() diff --git a/Tests/CodeModeTests/CapabilityRegistryTests.swift b/Tests/CodeModeTests/CapabilityRegistryTests.swift index 089e4d4..d191fe1 100644 --- a/Tests/CodeModeTests/CapabilityRegistryTests.swift +++ b/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") diff --git a/Tests/CodeModeTests/DefensiveDecodingTests.swift b/Tests/CodeModeTests/DefensiveDecodingTests.swift new file mode 100644 index 0000000..8a4485c --- /dev/null +++ b/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)) +} diff --git a/Tools/CodeModeEval/Package.resolved b/Tools/CodeModeEval/Package.resolved index 524d693..49387ad 100644 --- a/Tools/CodeModeEval/Package.resolved +++ b/Tools/CodeModeEval/Package.resolved @@ -15,8 +15,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/swiftlang/swift-syntax.git", "state" : { - "revision" : "f99ae8aa18f0cf0d53481901f88a0991dc3bd4a2", - "version" : "601.0.1" + "revision" : "79e4b74a295b6eb74a8b585e3a39d29e70c1dbd1", + "version" : "603.0.2" } } ],