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

package a2a_test

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"go.uber.org/zap"

	"iop/apps/edge/internal/input/a2a"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

// fakeService is a minimal runService implementation for testing.
type fakeService struct {
	submitFn func(context.Context, edgeservice.SubmitRunRequest) (edgeservice.RunResult, error)
	cancelFn func(context.Context, edgeservice.CancelRunRequest) (edgeservice.CommandResult, error)
}

func (f *fakeService) SubmitRun(ctx context.Context, req edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
	if f.submitFn != nil {
		return f.submitFn(ctx, req)
	}
	return nil, nil
}

func (f *fakeService) CancelRun(ctx context.Context, req edgeservice.CancelRunRequest) (edgeservice.CommandResult, error) {
	if f.cancelFn != nil {
		return f.cancelFn(ctx, req)
	}
	return edgeservice.CommandResult{}, nil
}

// newTestServer builds a disabled-by-default a2a.Server for handler-level testing.
// We call the handlers directly via httptest rather than starting a real listener.
func newTestServer(cfg config.EdgeA2AConf, svc runServiceIface) *a2a.Server {
	return a2a.NewServer(cfg, svc, zap.NewNop())
}

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

func rpcPost(t *testing.T, handler http.Handler, body any) *httptest.ResponseRecorder {
	t.Helper()
	b, err := json.Marshal(body)
	if err != nil {
		t.Fatalf("marshal body: %v", err)
	}
	req := httptest.NewRequest(http.MethodPost, "/a2a", bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	handler.ServeHTTP(w, req)
	return w
}

func decodeRPCResponse(t *testing.T, w *httptest.ResponseRecorder) a2a.JSONRPCResponse {
	t.Helper()
	var resp a2a.JSONRPCResponse
	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	return resp
}

func TestHandleMessageSendDispatchesRun(t *testing.T) {
	dispatched := false
	svc := &fakeService{
		submitFn: func(_ context.Context, req edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
			dispatched = true
			if req.Metadata["source"] != "a2a" {
				t.Errorf("expected source=a2a, got %q", req.Metadata["source"])
			}
			if req.ModelGroupKey != "" {
				t.Errorf("expected empty ModelGroupKey, got %q", req.ModelGroupKey)
			}
			return feedHandle("run-1", &iop.RunEvent{RunId: "run-1", Type: "complete"}), nil
		},
	}

	cfg := config.EdgeA2AConf{Enabled: true, Listen: "127.0.0.1:0", Path: "/a2a", TimeoutSec: 5}
	srv := a2a.NewServer(cfg, svc, zap.NewNop())

	mux := http.NewServeMux()
	mux.HandleFunc("/a2a", srv.RPCHandlerForTest())

	sendBody := a2a.JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "message/send",
		Params: mustMarshal(t, a2a.MessageSendParams{
			Message: a2a.Message{
				Role:  "user",
				Parts: []a2a.Part{{Type: "text", Text: "hi"}},
			},
			Configuration: &a2a.SendConfig{Blocking: boolPtr(true)},
		}),
		ID: 1,
	}
	w := rpcPost(t, mux, sendBody)
	if w.Code != http.StatusOK {
		t.Fatalf("message/send status: %d body: %s", w.Code, w.Body.String())
	}
	resp := decodeRPCResponse(t, w)
	if resp.Error != nil {
		t.Fatalf("message/send error: %+v", resp.Error)
	}
	if !dispatched {
		t.Error("SubmitRun was not called")
	}
}

func TestHandleGetTaskReturnsStoredTask(t *testing.T) {
	runID := "run-get-1"
	svc := &fakeService{
		submitFn: func(_ context.Context, _ edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
			return feedHandle(runID,
				&iop.RunEvent{RunId: runID, Type: "delta", Delta: "result text"},
				&iop.RunEvent{RunId: runID, Type: "complete"},
			), nil
		},
	}
	cfg := config.EdgeA2AConf{Enabled: false, Path: "/a2a", TimeoutSec: 5}
	srv := a2a.NewServer(cfg, svc, zap.NewNop())

	// Inject a task by calling message/send handler directly via httptest.
	sendBody := a2a.JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "message/send",
		Params: mustMarshal(t, a2a.MessageSendParams{
			Message: a2a.Message{
				Role:  "user",
				Parts: []a2a.Part{{Type: "text", Text: "hello"}},
			},
			Configuration: &a2a.SendConfig{Blocking: boolPtr(true)},
		}),
		ID: 1,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/a2a", srv.RPCHandlerForTest())
	w := rpcPost(t, mux, sendBody)
	if w.Code != http.StatusOK {
		t.Fatalf("message/send status: %d body: %s", w.Code, w.Body.String())
	}
	sendResp := decodeRPCResponse(t, w)
	if sendResp.Error != nil {
		t.Fatalf("message/send error: %+v", sendResp.Error)
	}

	// Extract task ID from result.
	taskMap, ok := sendResp.Result.(map[string]any)
	if !ok {
		t.Fatalf("expected map result, got %T", sendResp.Result)
	}
	taskID, _ := taskMap["id"].(string)
	if taskID == "" {
		t.Fatal("task id is empty")
	}

	// Now tasks/get.
	getBody := a2a.JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "tasks/get",
		Params:  mustMarshal(t, a2a.TaskQueryParams{ID: taskID}),
		ID:      2,
	}
	w2 := rpcPost(t, mux, getBody)
	if w2.Code != http.StatusOK {
		t.Fatalf("tasks/get status: %d", w2.Code)
	}
	getResp := decodeRPCResponse(t, w2)
	if getResp.Error != nil {
		t.Fatalf("tasks/get error: %+v", getResp.Error)
	}
	taskMap2, _ := getResp.Result.(map[string]any)
	status, _ := taskMap2["status"].(map[string]any)
	state, _ := status["state"].(string)
	if state != "completed" {
		t.Errorf("expected completed, got %q", state)
	}
}

func TestHandleCancelTask(t *testing.T) {
	runID := "run-cancel-1"
	cancelCalled := false
	svc := &fakeService{
		submitFn: func(_ context.Context, _ edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
			// Return a handle that stays working (never completes) for non-blocking.
			events := make(chan *iop.RunEvent, 1)
			nodeEvents := make(chan *iop.EdgeNodeEvent, 1)
			return &edgeservice.RunHandle{
				RunDispatch: edgeservice.RunDispatch{RunID: runID, TimeoutSec: 60},
				RunStream:   edgeservice.RunStream{Events: events, NodeEvents: nodeEvents},
			}, nil
		},
		cancelFn: func(_ context.Context, req edgeservice.CancelRunRequest) (edgeservice.CommandResult, error) {
			cancelCalled = true
			if req.RunID != runID {
				t.Errorf("expected RunID=%q, got %q", runID, req.RunID)
			}
			return edgeservice.CommandResult{}, nil
		},
	}
	cfg := config.EdgeA2AConf{Enabled: false, Path: "/a2a", TimeoutSec: 60}
	srv := a2a.NewServer(cfg, svc, zap.NewNop())

	mux := http.NewServeMux()
	mux.HandleFunc("/a2a", srv.RPCHandlerForTest())

	// Submit non-blocking so task stays in working state.
	sendBody := a2a.JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "message/send",
		Params: mustMarshal(t, a2a.MessageSendParams{
			Message: a2a.Message{
				Role:  "user",
				Parts: []a2a.Part{{Type: "text", Text: "work"}},
			},
			Configuration: &a2a.SendConfig{Blocking: boolPtr(false)},
		}),
		ID: 1,
	}
	w := rpcPost(t, mux, sendBody)
	if w.Code != http.StatusOK {
		t.Fatalf("message/send: %d %s", w.Code, w.Body.String())
	}
	sendResp := decodeRPCResponse(t, w)
	if sendResp.Error != nil {
		t.Fatalf("message/send error: %+v", sendResp.Error)
	}
	taskMap, _ := sendResp.Result.(map[string]any)
	taskID, _ := taskMap["id"].(string)

	// Cancel it.
	cancelBody := a2a.JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "tasks/cancel",
		Params:  mustMarshal(t, a2a.TaskIDParams{ID: taskID}),
		ID:      2,
	}
	w2 := rpcPost(t, mux, cancelBody)
	if w2.Code != http.StatusOK {
		t.Fatalf("tasks/cancel: %d %s", w2.Code, w2.Body.String())
	}
	cancelResp := decodeRPCResponse(t, w2)
	if cancelResp.Error != nil {
		t.Fatalf("tasks/cancel error: %+v", cancelResp.Error)
	}
	if !cancelCalled {
		t.Error("CancelRun was not called")
	}
}

func TestRejectsBadAuth(t *testing.T) {
	svc := &fakeService{}
	cfg := config.EdgeA2AConf{Enabled: false, Path: "/a2a", BearerToken: "secret"}
	srv := a2a.NewServer(cfg, svc, zap.NewNop())

	mux := http.NewServeMux()
	mux.HandleFunc("/a2a", srv.RPCHandlerForTest())

	body := a2a.JSONRPCRequest{JSONRPC: "2.0", Method: "message/send", ID: 1}
	b, _ := json.Marshal(body)
	req := httptest.NewRequest(http.MethodPost, "/a2a", bytes.NewReader(b))
	req.Header.Set("Authorization", "Bearer wrong")
	w := httptest.NewRecorder()
	mux.ServeHTTP(w, req)
	if w.Code != http.StatusUnauthorized {
		t.Errorf("expected 401, got %d", w.Code)
	}
}

func TestRejectsUnknownMethod(t *testing.T) {
	svc := &fakeService{}
	cfg := config.EdgeA2AConf{Enabled: false, Path: "/a2a"}
	srv := a2a.NewServer(cfg, svc, zap.NewNop())

	mux := http.NewServeMux()
	mux.HandleFunc("/a2a", srv.RPCHandlerForTest())

	body := a2a.JSONRPCRequest{JSONRPC: "2.0", Method: "unknown/method", ID: 1}
	w := rpcPost(t, mux, body)
	resp := decodeRPCResponse(t, w)
	if resp.Error == nil {
		t.Fatal("expected error for unknown method")
	}
	if resp.Error.Code != a2a.ErrCodeMethodNotFound {
		t.Errorf("expected code %d, got %d", a2a.ErrCodeMethodNotFound, resp.Error.Code)
	}
}

// TestBlockingDefault verifies the four cases for the blocking field:
//   - no configuration → blocking true
//   - configuration: {} (empty object, Blocking nil) → blocking true
//   - configuration: {"blocking": true} → blocking true
//   - configuration: {"blocking": false} → blocking false
func TestBlockingDefault(t *testing.T) {
	cases := []struct {
		name      string
		cfg       *a2a.SendConfig
		wantBlock bool
	}{
		{"no configuration", nil, true},
		{"empty configuration", &a2a.SendConfig{}, true},
		{"explicit true", &a2a.SendConfig{Blocking: boolPtr(true)}, true},
		{"explicit false", &a2a.SendConfig{Blocking: boolPtr(false)}, false},
	}

	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			events := make(chan *iop.RunEvent, 2)
			nodeEvents := make(chan *iop.EdgeNodeEvent, 1)
			var dispatched bool
			svc := &fakeService{
				submitFn: func(_ context.Context, req edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
					dispatched = true
					if req.ModelGroupKey != "" {
						t.Errorf("expected empty ModelGroupKey, got %q", req.ModelGroupKey)
					}
					got := req.Metadata["blocking"]
					want := "true"
					if !tc.wantBlock {
						want = "false"
					}
					if got != want {
						t.Errorf("blocking metadata: got %q want %q", got, want)
					}
					if tc.wantBlock {
						events <- &iop.RunEvent{RunId: "r", Type: "complete"}
					}
					return &edgeservice.RunHandle{
						RunDispatch: edgeservice.RunDispatch{RunID: "r", TimeoutSec: 5},
						RunStream:   edgeservice.RunStream{Events: events, NodeEvents: nodeEvents},
					}, nil
				},
			}
			cfg := config.EdgeA2AConf{Enabled: false, Path: "/a2a", TimeoutSec: 5}
			srv := a2a.NewServer(cfg, svc, zap.NewNop())
			mux := http.NewServeMux()
			mux.HandleFunc("/a2a", srv.RPCHandlerForTest())

			body := a2a.JSONRPCRequest{
				JSONRPC: "2.0",
				Method:  "message/send",
				Params: mustMarshal(t, a2a.MessageSendParams{
					Message:       a2a.Message{Role: "user", Parts: []a2a.Part{{Type: "text", Text: "hi"}}},
					Configuration: tc.cfg,
				}),
				ID: 1,
			}
			w := rpcPost(t, mux, body)
			if w.Code != http.StatusOK {
				t.Fatalf("status: %d body: %s", w.Code, w.Body.String())
			}
			resp := decodeRPCResponse(t, w)
			if resp.Error != nil {
				t.Fatalf("rpc error: %+v", resp.Error)
			}
			if !dispatched {
				t.Error("SubmitRun was not called")
			}
		})
	}
}

func mustMarshal(t *testing.T, v any) json.RawMessage {
	t.Helper()
	b, err := json.Marshal(v)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	return b
}

func boolPtr(b bool) *bool { return &b }