iop/apps/node/internal/node/node_test.go

package node_test

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/node"
	"iop/apps/node/internal/runtime"
	"iop/apps/node/internal/store"
	"iop/apps/node/internal/transport"
	iop "iop/proto/gen/iop"
)

// --- test doubles ---

type countingAdapter struct {
	executeCalls int32
}

func (a *countingAdapter) Name() string { return "test" }
func (a *countingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "test", Targets: []string{"v1"}, MaxConcurrency: 1}, nil
}
func (a *countingAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	atomic.AddInt32(&a.executeCalls, 1)
	return nil
}

// failingAdapter returns a fixed error from Execute.
type failingAdapter struct{ err error }

func (a *failingAdapter) Name() string { return "failing" }
func (a *failingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "failing"}, nil
}
func (a *failingAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	return a.err
}

// blockingAdapter blocks until its context is cancelled, then returns ErrRunCancelled.
type blockingAdapter struct {
	started chan struct{}
	done    chan struct{}
}

func newBlockingAdapter() *blockingAdapter {
	return &blockingAdapter{started: make(chan struct{}), done: make(chan struct{})}
}

func (a *blockingAdapter) Name() string { return "blocking" }
func (a *blockingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "blocking"}, nil
}
func (a *blockingAdapter) Execute(ctx context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	close(a.started)
	<-ctx.Done()
	close(a.done)
	return runtime.ErrRunCancelled
}

// terminatingAdapter implements SessionTerminator.
type terminatingAdapter struct {
	terminateCalls int32
	lastTarget     string
	lastSessionID  string
}

func (a *terminatingAdapter) Name() string { return "terminating" }
func (a *terminatingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "terminating"}, nil
}
func (a *terminatingAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	return nil
}
func (a *terminatingAdapter) TerminateSession(_ context.Context, target, sessionID string) error {
	atomic.AddInt32(&a.terminateCalls, 1)
	a.lastTarget = target
	a.lastSessionID = sessionID
	return nil
}

type commandAdapter struct {
	lastReq runtime.CommandRequest
}

func (a *commandAdapter) Name() string { return "command" }
func (a *commandAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: "command"}, nil
}
func (a *commandAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
	return nil
}
func (a *commandAdapter) HandleCommand(ctx context.Context, req runtime.CommandRequest) (runtime.CommandResponse, error) {
	a.lastReq = req
	if req.Type != runtime.CommandTypeUsageStatus {
		return runtime.CommandResponse{}, errors.New("command missing")
	}
	return runtime.CommandResponse{
		RequestID: req.RequestID,
		Type:      req.Type,
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionID: req.SessionID,
		UsageStatus: &runtime.AgentUsageStatus{
			RawOutput: "success",
		},
	}, nil
}

type fixedRouter struct {
	adapterName string
	adapter     runtime.Adapter
	adapters    map[string]runtime.Adapter
}

func (r *fixedRouter) Resolve(_ context.Context, req runtime.RunRequest) (runtime.ExecutionSpec, error) {
	return runtime.ExecutionSpec{
		RunID:       req.RunID,
		Adapter:     r.adapterName,
		Target:      req.Target,
		SessionID:   req.SessionID,
		SessionMode: req.SessionMode,
		Background:  req.Background,
		Workspace:   req.Workspace,
		Policy:      req.Policy,
		Input:       req.Input,
		TimeoutSec:  req.TimeoutSec,
		Metadata:    req.Metadata,
	}, nil
}

func (r *fixedRouter) ResolveAdapter(ctx context.Context, req runtime.RunRequest) (runtime.ExecutionSpec, runtime.Adapter, error) {
	spec, err := r.Resolve(ctx, req)
	if err != nil {
		return runtime.ExecutionSpec{}, nil, err
	}
	a, ok := r.adapters[spec.Adapter]
	if !ok {
		return runtime.ExecutionSpec{}, nil, fmt.Errorf("router: adapter %q not found", spec.Adapter)
	}
	return spec, a, nil
}

func (r *fixedRouter) GetAdapter(adapterName string) (runtime.Adapter, bool) {
	a, ok := r.adapters[adapterName]
	return a, ok
}

type errorRouter struct{ err error }

func (r *errorRouter) Resolve(_ context.Context, _ runtime.RunRequest) (runtime.ExecutionSpec, error) {
	return runtime.ExecutionSpec{}, r.err
}

func (r *errorRouter) ResolveAdapter(_ context.Context, _ runtime.RunRequest) (runtime.ExecutionSpec, runtime.Adapter, error) {
	return runtime.ExecutionSpec{}, nil, r.err
}

func (r *errorRouter) GetAdapter(_ string) (runtime.Adapter, bool) {
	return nil, false
}

func makeNode(t *testing.T, rtr runtime.Router) (*node.Node, *store.Store) {
	t.Helper()
	st, err := store.New(":memory:", zap.NewNop())
	if err != nil {
		t.Fatalf("store: %v", err)
	}
	t.Cleanup(func() { _ = st.Close() })
	return node.New("test-node", rtr, st, io.Discard, zap.NewNop()), st
}

// --- existing tests (updated) ---

func TestOnRunRequest_RouterError(t *testing.T) {
	n, _ := makeNode(t, &errorRouter{err: errors.New("boom")})

	err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1"})
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "node: resolve:") {
		t.Fatalf("expected resolve prefix, got %v", err)
	}
}

func TestOnRunRequest_AdapterNotFound(t *testing.T) {
	router := &fixedRouter{adapterName: "missing", adapters: make(map[string]runtime.Adapter)}
	n, _ := makeNode(t, router)

	err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1"})
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "not found") {
		t.Fatalf("expected adapter lookup error, got %v", err)
	}
}

func TestOnRunRequest_Success(t *testing.T) {
	adapter := &countingAdapter{}
	router := &fixedRouter{adapterName: "test", adapters: make(map[string]runtime.Adapter)}
	router.adapters["test"] = adapter
	n, st := makeNode(t, router)

	err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
		RunId:   "run-1",
		Adapter: "test",
		Target:  "v1",
	})
	if err != nil {
		t.Fatalf("run request: %v", err)
	}
	if got := atomic.LoadInt32(&adapter.executeCalls); got != 1 {
		t.Fatalf("expected 1 execute call, got %d", got)
	}

	run, err := st.GetRun(context.Background(), "run-1")
	if err != nil {
		t.Fatalf("get run: %v", err)
	}
	if run == nil {
		t.Fatal("expected persisted run record")
	}
	if run.Status != "completed" {
		t.Fatalf("expected completed status, got %q", run.Status)
	}
}

// --- new tests ---

func TestOnRunRequest_ForegroundAdapterErrorReturned(t *testing.T) {
	adapter := &failingAdapter{err: errors.New("adapter boom")}
	router := &fixedRouter{adapterName: "failing", adapters: make(map[string]runtime.Adapter)}
	router.adapters["failing"] = adapter
	n, st := makeNode(t, router)

	err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
		RunId:   "run-fail",
		Adapter: "failing",
		Target:  "v1",
	})
	if err == nil {
		t.Fatal("expected error from OnRunRequest")
	}
	if !strings.Contains(err.Error(), "adapter boom") {
		t.Fatalf("expected adapter boom in error, got %v", err)
	}

	run, err := st.GetRun(context.Background(), "run-fail")
	if err != nil {
		t.Fatalf("get run: %v", err)
	}
	if run == nil {
		t.Fatal("expected persisted run record")
	}
	if run.Status != "failed" {
		t.Fatalf("expected failed status, got %q", run.Status)
	}
	if !strings.Contains(run.Error, "adapter boom") {
		t.Fatalf("expected adapter boom in stored error, got %q", run.Error)
	}
}

func TestOnRunRequest_ForegroundCancelReturnedAndStored(t *testing.T) {
	adapter := &failingAdapter{err: runtime.ErrRunCancelled}
	router := &fixedRouter{adapterName: "failing", adapters: make(map[string]runtime.Adapter)}
	router.adapters["failing"] = adapter
	n, st := makeNode(t, router)

	err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
		RunId:   "run-cancel-fg",
		Adapter: "failing",
		Target:  "v1",
	})
	if !errors.Is(err, runtime.ErrRunCancelled) {
		t.Fatalf("expected ErrRunCancelled, got %v", err)
	}
	run, err := st.GetRun(context.Background(), "run-cancel-fg")
	if err != nil {
		t.Fatalf("get run: %v", err)
	}
	if run == nil || run.Status != "cancelled" {
		t.Fatalf("expected cancelled status, got %+v", run)
	}
}

func TestOnRunRequest_BackgroundReturnsBeforeAdapterCompletes(t *testing.T) {
	ba := newBlockingAdapter()
	router := &fixedRouter{adapterName: "blocking", adapters: make(map[string]runtime.Adapter)}
	router.adapters["blocking"] = ba
	n, st := makeNode(t, router)

	err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
		RunId:      "run-bg",
		Adapter:    "blocking",
		Target:     "v1",
		Background: true,
	})
	if err != nil {
		t.Fatalf("OnRunRequest: %v", err)
	}

	// OnRunRequest should return before adapter finishes.
	select {
	case <-ba.started:
	case <-time.After(2 * time.Second):
		t.Fatal("adapter never started")
	}

	// Cancel via OnCancel so the blocking adapter finishes.
	if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{RunId: "run-bg"}); err != nil {
		t.Fatalf("OnCancel: %v", err)
	}

	select {
	case <-ba.done:
	case <-time.After(2 * time.Second):
		t.Fatal("adapter did not finish after cancel")
	}

	// Wait briefly for completeRun goroutine to update the store.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		run, err := st.GetRun(context.Background(), "run-bg")
		if err != nil {
			t.Fatalf("get run: %v", err)
		}
		if run != nil && run.Status == "cancelled" {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	run, _ := st.GetRun(context.Background(), "run-bg")
	t.Fatalf("expected cancelled status, got %q", run.Status)
}

func TestOnCancel_CancelsRunViaRunManager(t *testing.T) {
	ba := newBlockingAdapter()
	router := &fixedRouter{adapterName: "blocking", adapters: make(map[string]runtime.Adapter)}
	router.adapters["blocking"] = ba
	n, _ := makeNode(t, router)

	go func() {
		_ = n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
			RunId:   "run-cancel",
			Adapter: "blocking",
			Target:  "v1",
		})
	}()

	<-ba.started

	if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{RunId: "run-cancel"}); err != nil {
		t.Fatalf("OnCancel: %v", err)
	}

	select {
	case <-ba.done:
	case <-time.After(2 * time.Second):
		t.Fatal("adapter did not finish after cancel")
	}
}

func TestOnCancel_TerminatesAdapterSession(t *testing.T) {
	ta := &terminatingAdapter{}
	router := &fixedRouter{adapterName: "terminating", adapters: make(map[string]runtime.Adapter)}
	router.adapters["terminating"] = ta
	n, _ := makeNode(t, router)

	err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{
		RunId:     "",
		Adapter:   "terminating",
		Target:    "codex",
		SessionId: "session-a",
		Action:    iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
	})
	if err != nil {
		t.Fatalf("OnCancel terminate: %v", err)
	}
	if atomic.LoadInt32(&ta.terminateCalls) != 1 {
		t.Fatal("expected TerminateSession to be called once")
	}
	if ta.lastTarget != "codex" || ta.lastSessionID != "session-a" {
		t.Fatalf("unexpected terminate args: target=%q session=%q", ta.lastTarget, ta.lastSessionID)
	}
}

func TestOnCommandRequest_Success(t *testing.T) {
	ca := &commandAdapter{}
	router := &fixedRouter{adapterName: "command", adapters: make(map[string]runtime.Adapter)}
	router.adapters["command"] = ca
	n, _ := makeNode(t, router)

	resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
		RequestId: "req-1",
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:   "command",
		Target:    "codex",
		SessionId: "default",
	})
	if err != nil {
		t.Fatalf("OnCommandRequest: %v", err)
	}
	if resp.GetError() != "" {
		t.Fatalf("expected no error, got %q", resp.GetError())
	}
	if resp.GetUsageStatus().GetRawOutput() != "success" {
		t.Fatalf("expected success raw output, got %q", resp.GetUsageStatus().GetRawOutput())
	}
	if ca.lastReq.Type != runtime.CommandTypeUsageStatus || ca.lastReq.Target != "codex" {
		t.Fatalf("unexpected last req: %+v", ca.lastReq)
	}
}

func TestOnCommandRequest_MissingAdapter(t *testing.T) {
	router := &fixedRouter{adapterName: "missing", adapters: make(map[string]runtime.Adapter)}
	n, _ := makeNode(t, router)

	resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
		RequestId: "req-1",
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:   "missing",
	})
	if err != nil {
		t.Fatalf("OnCommandRequest: %v", err)
	}
	if !strings.Contains(resp.GetError(), "not found") {
		t.Fatalf("expected not found error, got %q", resp.GetError())
	}
}

func TestOnCommandRequest_NotSupported(t *testing.T) {
	ta := &terminatingAdapter{}
	router := &fixedRouter{adapterName: "terminating", adapters: make(map[string]runtime.Adapter)}
	router.adapters["terminating"] = ta
	n, _ := makeNode(t, router)

	resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
		RequestId: "req-1",
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:   "terminating",
	})
	if err != nil {
		t.Fatalf("OnCommandRequest: %v", err)
	}
	if !strings.Contains(resp.GetError(), "does not support commands") {
		t.Fatalf("expected not supported error, got %q", resp.GetError())
	}
}

func TestOnCommandRequest_AdapterError(t *testing.T) {
	ca := &commandAdapter{}
	router := &fixedRouter{adapterName: "command", adapters: make(map[string]runtime.Adapter)}
	router.adapters["command"] = ca
	n, _ := makeNode(t, router)

	resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
		RequestId: "req-1",
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_UNSPECIFIED,
		Adapter:   "command",
	})
	if err != nil {
		t.Fatalf("OnCommandRequest: %v", err)
	}
	if !strings.Contains(resp.GetError(), "command missing") {
		t.Fatalf("expected command missing error, got %q", resp.GetError())
	}
}