ios: bound assistant image sets and keep output tools quiet

2026-10-20 14:31:15 -07:00
parent ee845c9fa3
commit b743523fd5
4 changed files with 59 additions and 3 deletions
@@ -105,9 +105,10 @@ extension Message {
              let info = metadata?.objectValue,
              let entries = info["agentAttachments"]?.arrayValue else { return [] }
        if role == .assistant {
-            guard info["kind"]?.stringValue == "message", info["status"]?.stringValue == "complete" else { return [] }
+            guard info["kind"]?.stringValue == "message", info["status"]?.stringValue == "complete",
+                  entries.count <= 8 else { return [] }
        }
-        return entries.compactMap { entry in
+        let attachments: [AgentAttachment] = entries.compactMap { entry in
            // Inspect the original JSON before Codable can discard unknown fields. Output
            // descriptors are deliberately path/URL-free; legacy user uploads stay compatible.
            if role == .assistant {
@@ -120,6 +121,14 @@ extension Message {
                  (try? attachment.validate()) != nil else { return nil }
            return attachment
        }
+        if role == .assistant {
+            // Reject the entire output set, including partially invalid arrays. The
+            // count and per-image bounds above also bound this aggregate calculation.
+            guard attachments.count == entries.count,
+                  Set(attachments.map(\.id)).count == attachments.count,
+                  attachments.reduce(0, { $0 + $1.sizeBytes }) <= 32 * 1024 * 1024 else { return [] }
+        }
+        return attachments
    }
 }

@@ -239,7 +239,7 @@ struct AgentPresentation {
    private(set) var suppressed = false
    private(set) var busySince: Date?
    private var seenTools: Set<String> = []
-    static let quietTools: Set<String> = ["user_message", "notify_parent", "clock_wait", "event_decision"]
+    static let quietTools: Set<String> = ["user_message", "output_image", "notify_parent", "clock_wait", "event_decision"]

    mutating func receive(_ event: AgentEvent, before: AgentRoomState, after: AgentRoomState, replay: Bool, now: Date) {
        switch event {
@@ -178,6 +178,26 @@ private func requestBytes(_ request: URLRequest) -> Data {
        for descriptor in descriptors {
            row.metadata = .object(["kind": .string("message"), "status": .string("complete"), "agentAttachments": .array([descriptor])])
            #expect(row.agentAttachments.isEmpty)
+            row.metadata = .object(["kind": .string("message"), "status": .string("complete"), "agentAttachments": .array([image.metadataValue, descriptor])])
+            #expect(row.agentAttachments.isEmpty) // No valid prefix from an invalid output set.
+        }
+        let maximum = (0..<8).map { index in
+            var item = image; item.id = "image-\(index)"; item.sizeBytes = 4 * 1024 * 1024
+            return item
+        }
+        row.metadata = .object(["kind": .string("message"), "status": .string("complete"), "agentAttachments": .array(maximum.map(\.metadataValue))])
+        #expect(row.agentAttachments == maximum) // Exactly eight IDs and 32 MiB.
+        let tooMany = (0..<9).map { index in
+            var item = image; item.id = "image-\(index)"
+            return item
+        }
+        var tooLarge = maximum; tooLarge[0].sizeBytes += 1
+        for invalid in [tooMany, [image, image], tooLarge] {
+            row.role = .assistant
+            row.metadata = .object(["kind": .string("message"), "status": .string("complete"), "agentAttachments": .array(invalid.map(\.metadataValue))])
+            #expect(row.agentAttachments.isEmpty)
+            row.role = .user
+            #expect(row.agentAttachments == invalid) // These new output limits do not change legacy input parsing.
        }
        row.role = .user
        row.metadata = .object(["agentAttachments": .array([sampleAttachment().metadataValue])])
@@ -347,6 +347,33 @@ struct AgentTests {
        #expect(presentation.isTyping(now: now.addingTimeInterval(7.3)))
    }

+    @Test func outputImageActivityKeepsImageReplyTypingSuppressed() throws {
+        var presentation = AgentPresentation(), state = AgentRoomState()
+        let now = Date(timeIntervalSince1970: 100)
+        state.activity.isThinking = true
+        presentation.sent(now: now, activity: state.activity)
+        #expect(presentation.isTyping(now: now.addingTimeInterval(0.3)))
+        let image = AgentAttachment(id: "image", kind: .image, filename: "image.png", mimeType: "image/png", sizeBytes: 12, sha256: String(repeating: "a", count: 64))
+        var reply = agentMessage("image-reply", entry: "1", info: ["status": .string("complete"), "agentAttachments": .array([image.metadataValue])])
+        reply.content = ""
+        let page = agentPage([reply], cursor: "1")
+        var after = state; try after.apply(page)
+        presentation.receive(.messages(page), before: state, after: after, replay: false, now: now.addingTimeInterval(1))
+        state = after
+        #expect(presentation.suppressed)
+
+        after.activity.tools = [.init(id: "output-activity", name: "output_image", status: "initiated")]
+        presentation.receive(.activity("root", after.activity), before: state, after: after, replay: false, now: now.addingTimeInterval(2))
+        state = after
+        #expect(!presentation.isTyping(now: now.addingTimeInterval(3)))
+
+        let tool = agentPage([agentMessage("output-tool", entry: "2", role: .tool, info: ["toolName": .string("output_image"), "status": .string("initiated")])], cursor: "2")
+        try after.apply(tool)
+        presentation.receive(.messages(tool), before: state, after: after, replay: false, now: now.addingTimeInterval(4))
+        #expect(presentation.suppressed)
+        #expect(!presentation.isTyping(now: now.addingTimeInterval(5)))
+    }
+
    @Test func socketURLAndAuthentication() throws {
        let config = APIConfiguration(baseURL: URL(string: "https://sybil.bajor.cloud/api/?accidental=query")!, authToken: " token ")
        #expect(try config.agentWebSocketURL().absoluteString == "wss://sybil.bajor.cloud/api/v1/agent/ws")