nomadcode/services/core/internal/adapters/a2a/client.go

package a2a

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"sync/atomic"
	"time"

	"github.com/nomadcode/nomadcode-core/internal/agent"
)

const (
	defaultTimeoutSec = 300
	jsonRPCVersion    = "2.0"
)

type Config struct {
	URL        string
	Token      string
	TimeoutSec int
}

type Client struct {
	cfg        Config
	httpClient *http.Client
	logger     *slog.Logger
	nextID     atomic.Uint64
}

func NewClient(cfg Config, logger *slog.Logger) *Client {
	if cfg.TimeoutSec <= 0 {
		cfg.TimeoutSec = defaultTimeoutSec
	}
	return &Client{
		cfg:        cfg,
		httpClient: &http.Client{Timeout: time.Duration(cfg.TimeoutSec) * time.Second},
		logger:     logger,
	}
}

func (c *Client) SendMessage(ctx context.Context, input agent.SendMessageInput) (agent.SendMessageResult, error) {
	if strings.TrimSpace(input.Text) == "" {
		return agent.SendMessageResult{}, fmt.Errorf("a2a message text is required")
	}

	messageID := newMessageID()
	msg := agent.Message{
		Kind:      "message",
		Role:      "user",
		MessageID: messageID,
		ContextID: input.ContextID,
		TaskID:    input.TaskID,
		Parts: []agent.Part{{
			Kind: "text",
			Text: input.Text,
		}},
	}

	params := messageSendParams{
		Message: msg,
		Configuration: &messageSendConfiguration{
			AcceptedOutputModes: input.AcceptedOutputModes,
			HistoryLength:       input.HistoryLength,
			Blocking:            input.Blocking,
		},
		Metadata: input.Metadata,
	}
	if params.Configuration.empty() {
		params.Configuration = nil
	}

	raw, err := c.call(ctx, "message/send", params)
	if err != nil {
		return agent.SendMessageResult{}, err
	}
	result, err := decodeMessageSendResult(raw)
	if err != nil {
		return agent.SendMessageResult{}, err
	}
	result.Raw = raw
	return result, nil
}

func (c *Client) GetTask(ctx context.Context, input agent.GetTaskInput) (agent.Task, error) {
	if strings.TrimSpace(input.TaskID) == "" {
		return agent.Task{}, fmt.Errorf("a2a task id is required")
	}
	raw, err := c.call(ctx, "tasks/get", taskIDParams{
		ID:            input.TaskID,
		HistoryLength: input.HistoryLength,
		Metadata:      input.Metadata,
	})
	if err != nil {
		return agent.Task{}, err
	}
	var task agent.Task
	if err := json.Unmarshal(raw, &task); err != nil {
		return agent.Task{}, fmt.Errorf("decode a2a task: %w", err)
	}
	return task, nil
}

func (c *Client) CancelTask(ctx context.Context, input agent.CancelTaskInput) (agent.Task, error) {
	if strings.TrimSpace(input.TaskID) == "" {
		return agent.Task{}, fmt.Errorf("a2a task id is required")
	}
	raw, err := c.call(ctx, "tasks/cancel", taskIDParams{
		ID:       input.TaskID,
		Metadata: input.Metadata,
	})
	if err != nil {
		return agent.Task{}, err
	}
	var task agent.Task
	if err := json.Unmarshal(raw, &task); err != nil {
		return agent.Task{}, fmt.Errorf("decode canceled a2a task: %w", err)
	}
	return task, nil
}

type messageSendParams struct {
	Message       agent.Message             `json:"message"`
	Configuration *messageSendConfiguration `json:"configuration,omitempty"`
	Metadata      map[string]any            `json:"metadata,omitempty"`
}

type messageSendConfiguration struct {
	AcceptedOutputModes []string `json:"acceptedOutputModes,omitempty"`
	HistoryLength       int      `json:"historyLength,omitempty"`
	Blocking            bool     `json:"blocking,omitempty"`
}

func (c *messageSendConfiguration) empty() bool {
	return c == nil || (len(c.AcceptedOutputModes) == 0 && c.HistoryLength == 0 && !c.Blocking)
}

type taskIDParams struct {
	ID            string         `json:"id"`
	HistoryLength int            `json:"historyLength,omitempty"`
	Metadata      map[string]any `json:"metadata,omitempty"`
}

type jsonRPCRequest struct {
	JSONRPC string `json:"jsonrpc"`
	ID      string `json:"id"`
	Method  string `json:"method"`
	Params  any    `json:"params,omitempty"`
}

type jsonRPCResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      string          `json:"id"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   *jsonRPCError   `json:"error,omitempty"`
}

type jsonRPCError struct {
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

func (c *Client) call(ctx context.Context, method string, params any) (json.RawMessage, error) {
	endpoint, err := endpointURL(c.cfg.URL)
	if err != nil {
		return nil, err
	}

	requestID := c.requestID()
	payload, err := json.Marshal(jsonRPCRequest{
		JSONRPC: jsonRPCVersion,
		ID:      requestID,
		Method:  method,
		Params:  params,
	})
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	if c.cfg.Token != "" {
		req.Header.Set("Authorization", "Bearer "+c.cfg.Token)
	}

	if c.logger != nil {
		c.logger.Info("a2a json-rpc request", "method", method, "endpoint", endpoint)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("a2a request failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}

	var rpcResp jsonRPCResponse
	if err := json.Unmarshal(body, &rpcResp); err != nil {
		return nil, fmt.Errorf("decode a2a json-rpc response: %w", err)
	}
	if rpcResp.Error != nil {
		return nil, fmt.Errorf("a2a json-rpc error %d: %s", rpcResp.Error.Code, rpcResp.Error.Message)
	}
	if rpcResp.ID != "" && rpcResp.ID != requestID {
		return nil, fmt.Errorf("a2a response id mismatch: got %q want %q", rpcResp.ID, requestID)
	}
	if len(rpcResp.Result) == 0 {
		return nil, fmt.Errorf("a2a response result is empty")
	}
	return rpcResp.Result, nil
}

func (c *Client) requestID() string {
	return fmt.Sprintf("nomadcode-%d", c.nextID.Add(1))
}

func decodeMessageSendResult(raw json.RawMessage) (agent.SendMessageResult, error) {
	var envelope struct {
		Kind string `json:"kind"`
	}
	_ = json.Unmarshal(raw, &envelope)

	switch envelope.Kind {
	case "message":
		var msg agent.Message
		if err := json.Unmarshal(raw, &msg); err != nil {
			return agent.SendMessageResult{}, fmt.Errorf("decode a2a message: %w", err)
		}
		return agent.SendMessageResult{Message: &msg}, nil
	case "task":
		var task agent.Task
		if err := json.Unmarshal(raw, &task); err != nil {
			return agent.SendMessageResult{}, fmt.Errorf("decode a2a task: %w", err)
		}
		return agent.SendMessageResult{Task: &task}, nil
	}

	var probe map[string]json.RawMessage
	if err := json.Unmarshal(raw, &probe); err != nil {
		return agent.SendMessageResult{}, fmt.Errorf("decode a2a result: %w", err)
	}
	if _, ok := probe["status"]; ok {
		var task agent.Task
		if err := json.Unmarshal(raw, &task); err != nil {
			return agent.SendMessageResult{}, fmt.Errorf("decode a2a task: %w", err)
		}
		return agent.SendMessageResult{Task: &task}, nil
	}
	if _, ok := probe["role"]; ok {
		var msg agent.Message
		if err := json.Unmarshal(raw, &msg); err != nil {
			return agent.SendMessageResult{}, fmt.Errorf("decode a2a message: %w", err)
		}
		return agent.SendMessageResult{Message: &msg}, nil
	}
	return agent.SendMessageResult{}, fmt.Errorf("a2a result is neither message nor task")
}

func endpointURL(rawURL string) (string, error) {
	rawURL = strings.TrimSpace(rawURL)
	if rawURL == "" {
		return "", fmt.Errorf("a2a endpoint URL is required")
	}
	parsed, err := url.Parse(rawURL)
	if err != nil {
		return "", err
	}
	if parsed.Scheme == "" || parsed.Host == "" {
		return "", fmt.Errorf("a2a endpoint URL must include scheme and host")
	}
	return parsed.String(), nil
}

func newMessageID() string {
	var b [16]byte
	if _, err := rand.Read(b[:]); err == nil {
		return "msg-" + hex.EncodeToString(b[:])
	}
	return fmt.Sprintf("msg-%d", time.Now().UnixNano())
}