iop/apps/edge/internal/input/a2a/server.go

package a2a

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
)

type runService interface {
	SubmitRun(context.Context, edgeservice.SubmitRunRequest) (edgeservice.RunResult, error)
	CancelRun(context.Context, edgeservice.CancelRunRequest) (edgeservice.CommandResult, error)
}

type Server struct {
	cfg    config.EdgeA2AConf
	svc    runService
	tasks  *TaskStore
	logger *zap.Logger
	server *http.Server
}

func NewServer(cfg config.EdgeA2AConf, svc runService, logger *zap.Logger) *Server {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &Server{cfg: cfg, svc: svc, tasks: newTaskStore(), logger: logger}
}

func (s *Server) Enabled() bool {
	return s != nil && s.cfg.Enabled
}

func (s *Server) Start(ctx context.Context) error {
	if !s.Enabled() {
		return nil
	}
	if s.cfg.Listen == "" {
		s.cfg.Listen = "0.0.0.0:8081"
	}
	path := s.cfg.Path
	if path == "" {
		path = "/a2a"
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/agent.json", s.handleAgentCard)
	mux.HandleFunc(path, s.withAuth(s.handleRPC))

	s.server = &http.Server{
		Addr:              s.cfg.Listen,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}
	ln, err := net.Listen("tcp", s.cfg.Listen)
	if err != nil {
		return fmt.Errorf("a2a server listen %s: %w", s.cfg.Listen, err)
	}
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = s.server.Shutdown(shutdownCtx)
	}()
	go func() {
		s.logger.Info("a2a server listening", zap.String("addr", s.cfg.Listen), zap.String("path", path))
		if err := s.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			s.logger.Warn("a2a server exited", zap.Error(err))
		}
	}()
	return nil
}

func (s *Server) Stop(ctx context.Context) error {
	if s == nil || s.server == nil {
		return nil
	}
	return s.server.Shutdown(ctx)
}

// RPCHandlerForTest returns the auth-wrapped JSON-RPC handler for unit testing
// without starting a real listener.
func (s *Server) RPCHandlerForTest() http.HandlerFunc {
	return s.withAuth(s.handleRPC)
}

func (s *Server) withAuth(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if s.cfg.BearerToken != "" {
			auth := r.Header.Get("Authorization")
			expected := "Bearer " + s.cfg.BearerToken
			if auth != expected {
				writeRPCError(w, nil, http.StatusUnauthorized, ErrCodeInvalidRequest, "unauthorized")
				return
			}
		}
		next(w, r)
	}
}

func (s *Server) handleAgentCard(w http.ResponseWriter, _ *http.Request) {
	card := map[string]any{
		"name":        "IOP Edge A2A Agent",
		"description": "IOP Edge A2A input surface",
		"url":         fmt.Sprintf("http://%s%s", s.cfg.Listen, s.cfg.Path),
		"version":     "1.0.0",
		"capabilities": map[string]any{
			"streaming": false,
		},
	}
	writeJSON(w, http.StatusOK, card)
}

func (s *Server) handleRPC(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeRPCError(w, nil, http.StatusMethodNotAllowed, ErrCodeInvalidRequest, "method not allowed")
		return
	}
	defer r.Body.Close()

	var req JSONRPCRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeRPCError(w, nil, http.StatusOK, ErrCodeParseError, "parse error")
		return
	}
	if req.JSONRPC != "2.0" {
		writeRPCError(w, req.ID, http.StatusOK, ErrCodeInvalidRequest, "jsonrpc must be \"2.0\"")
		return
	}

	switch req.Method {
	case "message/send":
		s.handleMessageSend(w, r, req)
	case "tasks/get":
		s.handleTasksGet(w, req)
	case "tasks/cancel":
		s.handleTasksCancel(w, r, req)
	default:
		writeRPCError(w, req.ID, http.StatusOK, ErrCodeMethodNotFound, "method not found: "+req.Method)
	}
}

func (s *Server) handleMessageSend(w http.ResponseWriter, r *http.Request, rpcReq JSONRPCRequest) {
	var params MessageSendParams
	if err := json.Unmarshal(rpcReq.Params, &params); err != nil {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeInvalidParams, "invalid params")
		return
	}

	prompt := extractPrompt(params.Message)
	if strings.TrimSpace(prompt) == "" {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeInvalidParams, "message parts are empty")
		return
	}

	// blocking defaults to true when configuration is absent or Blocking field is omitted.
	blocking := params.Configuration == nil || params.Configuration.Blocking == nil || *params.Configuration.Blocking

	handle, err := s.svc.SubmitRun(r.Context(), edgeservice.SubmitRunRequest{
		NodeRef:    s.cfg.NodeRef,
		Adapter:    s.resolveAdapter(),
		Target:     s.resolveTarget(),
		SessionID:  s.resolveSessionID(),
		Prompt:     prompt,
		TimeoutSec: s.resolveTimeoutSec(),
		Metadata: map[string]string{
			"source":   "a2a",
			"blocking": fmt.Sprintf("%t", blocking),
		},
	})
	if err != nil {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeInvalidRequest, err.Error())
		return
	}

	taskID := s.tasks.create(handle.Dispatch().RunID, params.Message)

	var snapshot *Task
	if blocking {
		snapshot = s.tasks.collectBlocking(taskID, handle)
	} else {
		s.tasks.collectBackground(taskID, handle)
		// Get a snapshot after background drain is started; background goroutine
		// owns the internal pointer from this point, so we never share it.
		snapshot, _ = s.tasks.get(taskID)
	}

	writeRPCResult(w, rpcReq.ID, snapshot)
}

func (s *Server) handleTasksGet(w http.ResponseWriter, rpcReq JSONRPCRequest) {
	var params TaskQueryParams
	if err := json.Unmarshal(rpcReq.Params, &params); err != nil {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeInvalidParams, "invalid params")
		return
	}
	task, ok := s.tasks.get(params.ID)
	if !ok {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeTaskNotFound, "task not found: "+params.ID)
		return
	}
	writeRPCResult(w, rpcReq.ID, task)
}

func (s *Server) handleTasksCancel(w http.ResponseWriter, r *http.Request, rpcReq JSONRPCRequest) {
	var params TaskIDParams
	if err := json.Unmarshal(rpcReq.Params, &params); err != nil {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeInvalidParams, "invalid params")
		return
	}
	task, ok := s.tasks.get(params.ID)
	if !ok {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeTaskNotFound, "task not found: "+params.ID)
		return
	}
	if task.Status.State != StateWorking {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeNotCancelable, "task is not in working state")
		return
	}
	_, err := s.svc.CancelRun(r.Context(), edgeservice.CancelRunRequest{
		NodeRef:   s.cfg.NodeRef,
		RunID:     params.ID,
		Adapter:   s.resolveAdapter(),
		Target:    s.resolveTarget(),
		SessionID: s.resolveSessionID(),
	})
	if err != nil {
		writeRPCError(w, rpcReq.ID, http.StatusOK, ErrCodeInvalidRequest, err.Error())
		return
	}
	writeRPCResult(w, rpcReq.ID, map[string]string{"id": params.ID})
}

func (s *Server) resolveAdapter() string {
	if s.cfg.Adapter != "" {
		return s.cfg.Adapter
	}
	return "cli"
}

func (s *Server) resolveTarget() string {
	return s.cfg.Target
}

func (s *Server) resolveSessionID() string {
	if s.cfg.SessionID != "" {
		return s.cfg.SessionID
	}
	return edgeservice.DefaultSessionID
}

func (s *Server) resolveTimeoutSec() int {
	if s.cfg.TimeoutSec > 0 {
		return s.cfg.TimeoutSec
	}
	return edgeservice.DefaultTimeoutSec
}

func extractPrompt(msg Message) string {
	var b strings.Builder
	for _, p := range msg.Parts {
		if p.Type == "text" && p.Text != "" {
			if b.Len() > 0 {
				b.WriteString("\n")
			}
			b.WriteString(p.Text)
		}
	}
	return b.String()
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeRPCResult(w http.ResponseWriter, id any, result any) {
	writeJSON(w, http.StatusOK, JSONRPCResponse{
		JSONRPC: "2.0",
		Result:  result,
		ID:      id,
	})
}

func writeRPCError(w http.ResponseWriter, id any, httpStatus, code int, message string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(httpStatus)
	_ = json.NewEncoder(w).Encode(JSONRPCResponse{
		JSONRPC: "2.0",
		Error:   &JSONRPCError{Code: code, Message: message},
		ID:      id,
	})
}