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

package a2a_test

import (
	"context"
	"net/http"
	"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"
)

// feedHandle builds a RunHandle pre-loaded with the given events.
func feedHandle(runID string, events ...*iop.RunEvent) *edgeservice.RunHandle {
	ch := make(chan *iop.RunEvent, len(events))
	for _, e := range events {
		ch <- e
	}
	nodeEvents := make(chan *iop.EdgeNodeEvent, 1)
	return &edgeservice.RunHandle{
		RunDispatch: edgeservice.RunDispatch{RunID: runID, TimeoutSec: 5},
		RunStream:   edgeservice.RunStream{Events: ch, NodeEvents: nodeEvents},
	}
}

func checkBlockingResult(t *testing.T, handle edgeservice.RunResult, check func(map[string]any)) {
	t.Helper()
	svc := &fakeService{
		submitFn: func(_ context.Context, _ edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
			return handle, 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: "ping"}},
			},
			Configuration: &a2a.SendConfig{Blocking: boolPtr(true)},
		}),
		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)
	}
	task, ok := resp.Result.(map[string]any)
	if !ok {
		t.Fatalf("expected map result, got %T", resp.Result)
	}
	check(task)
}

func TestTaskStoreMapsRunEventsToCompletedTask(t *testing.T) {
	runID := "run-ts-1"
	checkBlockingResult(t,
		feedHandle(runID,
			&iop.RunEvent{RunId: runID, Type: "delta", Delta: "hello "},
			&iop.RunEvent{RunId: runID, Type: "delta", Delta: "world"},
			&iop.RunEvent{RunId: runID, Type: "complete"},
		),
		func(task map[string]any) {
			status := task["status"].(map[string]any)
			if status["state"] != "completed" {
				t.Errorf("expected completed, got %q", status["state"])
			}
			artifacts, ok := task["artifacts"].([]any)
			if !ok || len(artifacts) == 0 {
				t.Fatal("expected non-empty artifacts")
			}
			parts := artifacts[0].(map[string]any)["parts"].([]any)
			text := parts[0].(map[string]any)["text"].(string)
			if text != "hello world" {
				t.Errorf("expected %q, got %q", "hello world", text)
			}
		},
	)
}

func TestTaskStoreMapsErrorToFailedTask(t *testing.T) {
	runID := "run-ts-2"
	checkBlockingResult(t,
		feedHandle(runID,
			&iop.RunEvent{RunId: runID, Type: "error", Error: "something broke"},
		),
		func(task map[string]any) {
			status := task["status"].(map[string]any)
			if status["state"] != "failed" {
				t.Errorf("expected failed, got %q", status["state"])
			}
		},
	)
}

func TestTaskStoreMapsClosedRunStreamToFailedTask(t *testing.T) {
	runID := "run-ts-closed"
	events := make(chan *iop.RunEvent)
	close(events)
	handle := &edgeservice.RunHandle{
		RunDispatch: edgeservice.RunDispatch{RunID: runID, TimeoutSec: 60},
		RunStream: edgeservice.RunStream{
			Events:     events,
			NodeEvents: make(chan *iop.EdgeNodeEvent),
		},
	}

	checkBlockingResult(t, handle, func(task map[string]any) {
		status := task["status"].(map[string]any)
		if status["state"] != "failed" {
			t.Errorf("expected failed, got %q", status["state"])
		}
	})
}

func TestMessageSendBlockingReturnsCompletedTask(t *testing.T) {
	runID := "run-ts-3"
	checkBlockingResult(t,
		feedHandle(runID,
			&iop.RunEvent{RunId: runID, Type: "delta", Delta: "done"},
			&iop.RunEvent{RunId: runID, Type: "complete"},
		),
		func(task map[string]any) {
			status := task["status"].(map[string]any)
			if status["state"] != "completed" {
				t.Errorf("expected completed, got %q", status["state"])
			}
		},
	)
}

func TestCancelTaskSendsCancelRun(t *testing.T) {
	runID := "run-ts-cancel"
	cancelCalled := false

	events := make(chan *iop.RunEvent, 1)
	nodeEvents := make(chan *iop.EdgeNodeEvent, 1)
	handle := &edgeservice.RunHandle{
		RunDispatch: edgeservice.RunDispatch{RunID: runID, TimeoutSec: 60},
		RunStream:   edgeservice.RunStream{Events: events, NodeEvents: nodeEvents},
	}

	svc := &fakeService{
		submitFn: func(_ context.Context, _ edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
			return handle, 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())

	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)
	sendResp := decodeRPCResponse(t, w)
	if sendResp.Error != nil {
		t.Fatalf("message/send: %+v", sendResp.Error)
	}
	taskID := sendResp.Result.(map[string]any)["id"].(string)

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