Claude implementation of A2A in Swift

a2a.swift

// A2A Protocol Implementation Example in Swift using Hummingbird
//
// This is a Swift implementation of Google's Agent-to-Agent (A2A) protocol using Hummingbird 2.
// The example demonstrates two components:
// 1. An A2A Server: An agent that exposes an HTTP endpoint implementing A2A protocol methods
// 2. An A2A Client: An application that consumes A2A services

import Foundation
import Hummingbird
import NIOCore
import NIOPosix
import NIOHTTP1
import AsyncHTTPClient
import NIOFoundationCompat

// MARK: - Models

// Task State Enumeration
enum TaskState: String, Codable {
    case submitted
    case working
    case inputRequired = "input-required"
    case completed
    case failed
    case canceled
}

// Agent Card Structure
struct AgentCard: Codable {
    let name: String
    let description: String
    let url: String
    let version: String
    let capabilities: Capabilities
    let defaultInputModes: [String]
    let defaultOutputModes: [String]
    let skills: [Skill]
    
    struct Capabilities: Codable {
        let streaming: Bool
        let pushNotifications: Bool
        let stateTransitionHistory: Bool
    }
    
    struct Skill: Codable {
        let id: String
        let name: String
        let description: String
        let tags: [String]
        let examples: [String]
    }
}

// Message Part Structure
struct MessagePart: Codable {
    let type: String
    let text: String?
}

// Message Structure
struct Message: Codable {
    let role: String
    let parts: [MessagePart]
}

// Task Structure
struct Task {
    let id: String
    var state: TaskState
    let message: Message
    let createdAt: Date
    var response: Message?
}

// Request Structures
struct TaskSendRequest: Codable {
    let message: Message
}

struct TaskGetRequest: Codable {
    let taskId: String
    
    enum CodingKeys: String, CodingKey {
        case taskId = "task_id"
    }
}

struct TaskCancelRequest: Codable {
    let taskId: String
    
    enum CodingKeys: String, CodingKey {
        case taskId = "task_id"
    }
}

// Response Structures
struct TaskSendResponse: Codable {
    let taskId: String
    
    enum CodingKeys: String, CodingKey {
        case taskId = "task_id"
    }
}

struct TaskGetResponse: Codable {
    let taskId: String
    let state: String
    let message: Message?
    
    enum CodingKeys: String, CodingKey {
        case taskId = "task_id"
        case state
        case message
    }
}

struct ErrorResponse: Codable {
    let error: String
    let code: Int
}

// MARK: - Server Implementation

// In-memory store for tasks
class TaskStore {
    private var tasks: [String: Task] = [:]
    private let lock = NSLock()
    
    func addTask(task: Task) {
        lock.lock()
        defer { lock.unlock() }
        tasks[task.id] = task
    }
    
    func getTask(id: String) -> Task? {
        lock.lock()
        defer { lock.unlock() }
        return tasks[id]
    }
    
    func updateTask(id: String, state: TaskState, response: Message? = nil) -> Bool {
        lock.lock()
        defer { lock.unlock() }
        
        guard var task = tasks[id] else {
            return false
        }
        
        task.state = state
        if let response = response {
            task.response = response
        }
        
        tasks[id] = task
        return true
    }
}

// A2A Server Application
actor A2AServer {
    private let app: HBApplication
    private let taskStore = TaskStore()
    
    // Agent Card Definition
    private let agentCard = AgentCard(
        name: "Echo Agent",
        description: "A simple agent that echoes back messages with a friendly greeting",
        url: "http://localhost:8080/",
        version: "1.0.0",
        capabilities: AgentCard.Capabilities(
            streaming: false,
            pushNotifications: false,
            stateTransitionHistory: false
        ),
        defaultInputModes: [
            "text",
            "text/plain"
        ],
        defaultOutputModes: [
            "text",
            "text/plain"
        ],
        skills: [
            AgentCard.Skill(
                id: "echo",
                name: "Echo Service",
                description: "Echoes back the input with a friendly greeting",
                tags: ["echo", "test"],
                examples: [
                    "Hello, how are you?",
                    "What's your name?"
                ]
            )
        ]
    )
    
    init() {
        // Create Hummingbird application
        app = HBApplication(
            configuration: .init(
                address: .hostname("localhost", port: 8080),
                serverName: "A2A Server"
            )
        )
        
        // Configure JSON encoder/decoder
        let encoder = JSONEncoder()
        let decoder = JSONDecoder()
        
        app.encoder = encoder
        app.decoder = decoder
        
        // Register routes
        setupRoutes()
    }
    
    private func setupRoutes() {
        // Expose the Agent Card at the standard /.well-known/agent.json path
        app.router.get("/.well-known/agent.json") { request -> AgentCard in
            return self.agentCard
        }
        
        // A2A protocol endpoint to send a task
        app.router.post("/tasks/send") { request -> HBResponse in
            do {
                let taskRequest = try request.decode(as: TaskSendRequest.self)
                
                // Extract first message from payload
                let messageContent = taskRequest.message.parts.first?.text ?? ""
                
                // Create a new task ID
                let taskId = UUID().uuidString
                
                // Create a task in "submitted" state
                let task = Task(
                    id: taskId,
                    state: .submitted,
                    message: taskRequest.message,
                    createdAt: Date(),
                    response: nil
                )
                
                self.taskStore.addTask(task: task)
                
                // Process the task (normally this would be async)
                // For this example, we'll just echo the message instantly
                self.taskStore.updateTask(id: taskId, state: .working)
                
                // Generate a simple echo response
                let responseText = "Echo Agent: You said '\(messageContent)'. Hello from the A2A protocol example!"
                
                // Complete the task with a response message
                let responseMessage = Message(
                    role: "agent",
                    parts: [
                        MessagePart(
                            type: "text",
                            text: responseText
                        )
                    ]
                )
                
                self.taskStore.updateTask(id: taskId, state: .completed, response: responseMessage)
                
                // Return the task ID per A2A protocol
                let response = TaskSendResponse(taskId: taskId)
                return try HBResponse(status: .ok, body: .json(response))
            } catch {
                let errorResponse = ErrorResponse(error: error.localizedDescription, code: 400)
                return try! HBResponse(status: .badRequest, body: .json(errorResponse))
            }
        }
        
        // A2A protocol endpoint to get task status and results
        app.router.post("/tasks/get") { request -> HBResponse in
            do {
                let taskRequest = try request.decode(as: TaskGetRequest.self)
                let taskId = taskRequest.taskId
                
                guard let task = self.taskStore.getTask(id: taskId) else {
                    let errorResponse = ErrorResponse(error: "Task not found", code: 404)
                    return try HBResponse(status: .notFound, body: .json(errorResponse))
                }
                
                var response = TaskGetResponse(
                    taskId: taskId,
                    state: task.state.rawValue,
                    message: nil
                )
                
                // If the task is completed, include the response message
                if task.state == .completed, let responseMessage = task.response {
                    response = TaskGetResponse(
                        taskId: taskId, 
                        state: task.state.rawValue,
                        message: responseMessage
                    )
                }
                
                return try HBResponse(status: .ok, body: .json(response))
            } catch {
                let errorResponse = ErrorResponse(error: error.localizedDescription, code: 400)
                return try! HBResponse(status: .badRequest, body: .json(errorResponse))
            }
        }
        
        // A2A protocol endpoint to cancel a task
        app.router.post("/tasks/cancel") { request -> HBResponse in
            do {
                let taskRequest = try request.decode(as: TaskCancelRequest.self)
                let taskId = taskRequest.taskId
                
                guard let task = self.taskStore.getTask(id: taskId) else {
                    let errorResponse = ErrorResponse(error: "Task not found", code: 404)
                    return try HBResponse(status: .notFound, body: .json(errorResponse))
                }
                
                // Only allow cancellation of tasks not already completed
                if task.state != .completed && task.state != .failed && task.state != .canceled {
                    _ = self.taskStore.updateTask(id: taskId, state: .canceled)
                }
                
                guard let updatedTask = self.taskStore.getTask(id: taskId) else {
                    let errorResponse = ErrorResponse(error: "Task not found after update", code: 404)
                    return try HBResponse(status: .notFound, body: .json(errorResponse))
                }
                
                let response = TaskGetResponse(
                    taskId: taskId,
                    state: updatedTask.state.rawValue,
                    message: nil
                )
                
                return try HBResponse(status: .ok, body: .json(response))
            } catch {
                let errorResponse = ErrorResponse(error: error.localizedDescription, code: 400)
                return try! HBResponse(status: .badRequest, body: .json(errorResponse))
            }
        }
    }
    
    func start() async throws {
        try await app.start()
        print("A2A Server started on http://localhost:8080")
    }
    
    func stop() async throws {
        try await app.stop()
        print("A2A Server stopped")
    }
}

// MARK: - Client Implementation

actor A2AClient {
    private let httpClient: HTTPClient
    private let baseURL: String
    
    init(agentURL: String) {
        self.baseURL = agentURL.hasSuffix("/") ? String(agentURL.dropLast()) : agentURL
        self.httpClient = HTTPClient(eventLoopGroupProvider: .createNew)
    }
    
    deinit {
        try? httpClient.syncShutdown()
    }
    
    func getAgentCard() async throws -> AgentCard {
        let agentCardURL = "\(baseURL)/.well-known/agent.json"
        let request = try HTTPClient.Request(url: agentCardURL, method: .GET)
        
        let response = try await httpClient.execute(request: request).get()
        
        guard response.status == .ok else {
            throw NSError(domain: "A2A", code: Int(response.status.code), userInfo: [NSLocalizedDescriptionKey: "Failed to fetch agent card"])
        }
        
        guard let body = response.body else {
            throw NSError(domain: "A2A", code: 0, userInfo: [NSLocalizedDescriptionKey: "Empty response body"])
        }
        
        let data = Data(buffer: body)
        return try JSONDecoder().decode(AgentCard.self, from: data)
    }
    
    func sendMessage(textContent: String) async throws -> Message {
        // Create a message with text content
        let message = Message(
            role: "user",
            parts: [
                MessagePart(
                    type: "text",
                    text: textContent
                )
            ]
        )
        
        // Send the task
        let taskRequest = TaskSendRequest(message: message)
        let requestData = try JSONEncoder().encode(taskRequest)
        
        var request = try HTTPClient.Request(
            url: "\(baseURL)/tasks/send",
            method: .POST
        )
        request.headers.add(name: "Content-Type", value: "application/json")
        request.body = .data(requestData)
        
        let response = try await httpClient.execute(request: request).get()
        
        guard response.status == .ok, let body = response.body else {
            throw NSError(domain: "A2A", code: Int(response.status.code), userInfo: [NSLocalizedDescriptionKey: "Failed to send task"])
        }
        
        let data = Data(buffer: body)
        let taskResponse = try JSONDecoder().decode(TaskSendResponse.self, from: data)
        let taskId = taskResponse.taskId
        
        // Poll for the task result
        while true {
            let taskGetRequest = TaskGetRequest(taskId: taskId)
            let taskGetData = try JSONEncoder().encode(taskGetRequest)
            
            var taskGetHTTPRequest = try HTTPClient.Request(
                url: "\(baseURL)/tasks/get",
                method: .POST
            )
            taskGetHTTPRequest.headers.add(name: "Content-Type", value: "application/json")
            taskGetHTTPRequest.body = .data(taskGetData)
            
            let taskGetResponse = try await httpClient.execute(request: taskGetHTTPRequest).get()
            
            guard taskGetResponse.status == .ok, let taskGetBody = taskGetResponse.body else {
                throw NSError(domain: "A2A", code: Int(taskGetResponse.status.code), userInfo: [NSLocalizedDescriptionKey: "Failed to get task status"])
            }
            
            let taskGetData = Data(buffer: taskGetBody)
            let taskStatus = try JSONDecoder().decode(TaskGetResponse.self, from: taskGetData)
            
            if taskStatus.state == TaskState.completed.rawValue, let message = taskStatus.message {
                return message
            } else if taskStatus.state == TaskState.failed.rawValue || taskStatus.state == TaskState.canceled.rawValue {
                throw NSError(domain: "A2A", code: 0, userInfo: [NSLocalizedDescriptionKey: "Task \(taskStatus.state)"])
            }
            
            // Wait before polling again
            try await Task.sleep(nanoseconds: 500_000_000) // 0.5 seconds
        }
    }
}

// MARK: - Example usage

@main
struct A2AExample {
    static func main() async throws {
        // Start the server
        let server = A2AServer()
        
        // Start the server in a separate task
        Task {
            try await server.start()
        }
        
        // Give the server time to start
        try await Task.sleep(nanoseconds: 1_000_000_000) // 1 second
        
        // Create a client and communicate with the agent
        print("Starting A2A client communication...")
        
        do {
            let client = A2AClient(agentURL: "http://localhost:8080")
            
            // Get the agent card
            let agentCard = try await client.getAgentCard()
            print("\n=== Agent Card ===")
            print("Name: \(agentCard.name)")
            print("Description: \(agentCard.description)")
            print("Skills: \(agentCard.skills.map { $0.name }.joined(separator: ", "))")
            
            // Send a message
            let userMessage = "Hello A2A agent! This is a test message."
            print("\nSending message: '\(userMessage)'")
            
            let responseMessage = try await client.sendMessage(textContent: userMessage)
            
            // Extract and print the response text
            let responseText = responseMessage.parts.first?.text ?? "No response"
            print("\nAgent response: \(responseText)")
            
        } catch {
            print("Error: \(error.localizedDescription)")
        }
        
        // Keep the server running for a short while
        print("\nServer is running. Press Ctrl+C to exit.")
        
        // Wait for a signal to exit
        try await Task.sleep(nanoseconds: 30_000_000_000) // 30 seconds
        
        // Stop the server
        try await server.stop()
    }
}

// MARK: - Extensions

extension HBRequest {
    func decode<T: Decodable>(as type: T.Type) throws -> T {
        guard let body = body.buffer else {
            throw NSError(domain: "A2A", code: 0, userInfo: [NSLocalizedDescriptionKey: "Empty request body"])
        }
        
        let data = Data(buffer: body)
        return try JSONDecoder().decode(type, from: data)
    }
}

extension Data {
    init(buffer: ByteBuffer) {
        var buffer = buffer
        self = buffer.readData(length: buffer.readableBytes) ?? Data()
    }
}