iop/apps/node/internal/node/command_test.go

package node_test

import (
	"context"
	"errors"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"iop/apps/node/internal/transport"
	runtime "iop/packages/go/execution"
	iop "iop/proto/gen/iop"
)

type providerCommandAdapter struct {
	mu       sync.Mutex
	commands []runtime.CommandType
	runs     []runtime.ExecutionSpec
	started  chan struct{}
	release  chan struct{}
	probe    runtime.ProviderProbeResult
	probeErr error
	probes   int
}

func (a *providerCommandAdapter) Name() string { return "provider" }
func (a *providerCommandAdapter) Capabilities(context.Context) (runtime.Capabilities, error) {
	return runtime.Capabilities{AdapterName: a.Name(), InstanceKey: "provider-instance", Targets: []string{"model"}, MaxConcurrency: 2}, nil
}
func (a *providerCommandAdapter) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
	a.mu.Lock()
	a.runs = append(a.runs, spec)
	a.mu.Unlock()
	if a.started != nil {
		select {
		case a.started <- struct{}{}:
		default:
		}
		select {
		case <-ctx.Done():
			return runtime.ErrRunCancelled
		case <-a.release:
		}
	}
	return nil
}
func (a *providerCommandAdapter) HandleCommand(_ context.Context, req runtime.CommandRequest) (runtime.CommandResponse, error) {
	a.mu.Lock()
	a.commands = append(a.commands, req.Type)
	a.mu.Unlock()
	return runtime.CommandResponse{
		RequestID: req.RequestID, Type: req.Type, Adapter: req.Adapter,
		Target: req.Target, SessionID: req.SessionID, Result: map[string]string{"status": "ok"},
	}, nil
}
func (a *providerCommandAdapter) ProbeProvider(_ context.Context, target string) (runtime.ProviderProbeResult, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.probes++
	result := a.probe
	if result.InstanceKey == "" {
		result.InstanceKey = "provider-instance"
	}
	if result.Target == "" {
		result.Target = target
	}
	return result, a.probeErr
}

func TestCapabilitiesHealthEvidence(t *testing.T) {
	t.Run("exact available evidence is session sequenced", func(t *testing.T) {
		adapter := &providerCommandAdapter{probe: runtime.ProviderProbeResult{
			AdapterName: "provider", Target: "model", Status: runtime.ProviderStatusAvailable,
		}}
		router := &fixedRouter{adapterName: "provider", adapters: map[string]runtime.Provider{"provider": adapter}}
		n, _ := makeNode(t, router)
		sess := &transport.Session{}

		for wantSeq := uint64(1); wantSeq <= 2; wantSeq++ {
			resp, err := n.OnCommandRequest(context.Background(), sess, &iop.NodeCommandRequest{
				RequestId: "caps", Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
				Adapter: "provider", Target: "model",
			})
			if err != nil || resp.GetError() != "" {
				t.Fatalf("response=%v err=%v", resp, err)
			}
			result := resp.GetResult()
			if result["adapter_key"] != "provider" || result["target"] != "model" || result["provider_status"] != "available" {
				t.Fatalf("unstable exact evidence: %#v", result)
			}
			if result["health_observation_seq"] != strconv.FormatUint(wantSeq, 10) {
				t.Fatalf("sequence=%q, want %d", result["health_observation_seq"], wantSeq)
			}
			if len(resp.GetProviderSnapshots()) != 1 || resp.GetProviderSnapshots()[0].GetHealth() != "available" {
				t.Fatalf("provider snapshot did not carry normalized health: %#v", resp.GetProviderSnapshots())
			}
		}
		if adapter.probes != 2 {
			t.Fatalf("probe calls=%d, want 2", adapter.probes)
		}
	})

	for _, tc := range []struct {
		name       string
		result     runtime.ProviderProbeResult
		err        error
		wantStatus string
	}{
		{name: "transport error", result: runtime.ProviderProbeResult{AdapterName: "provider", Target: "model", Status: runtime.ProviderStatusAvailable}, err: errors.New("probe failed"), wantStatus: "unknown"},
		{name: "identity mismatch", result: runtime.ProviderProbeResult{AdapterName: "other", Target: "model", Status: runtime.ProviderStatusAvailable}, wantStatus: "unknown"},
		{name: "unknown", result: runtime.ProviderProbeResult{AdapterName: "provider", Target: "model", Status: runtime.ProviderStatusUnknown}, wantStatus: "unknown"},
		{name: "exact unavailable", result: runtime.ProviderProbeResult{AdapterName: "provider", Target: "model", Status: runtime.ProviderStatusUnavailable}, wantStatus: "unavailable"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			adapter := &providerCommandAdapter{probe: tc.result, probeErr: tc.err}
			router := &fixedRouter{adapterName: "provider", adapters: map[string]runtime.Provider{"provider": adapter}}
			n, _ := makeNode(t, router)
			resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
				RequestId: "caps", Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
				Adapter: "provider", Target: "model",
			})
			if err != nil || resp.GetError() != "" {
				t.Fatalf("response=%v err=%v", resp, err)
			}
			if got := resp.GetResult()["provider_status"]; got != tc.wantStatus {
				t.Fatalf("provider_status=%q, want %q", got, tc.wantStatus)
			}
			if resp.GetResult()["health_observation_seq"] != "1" {
				t.Fatalf("missing Session sequence: %#v", resp.GetResult())
			}
		})
	}
}

func TestNodeCommandProviderAllowlist(t *testing.T) {
	adapter := &providerCommandAdapter{}
	router := &fixedRouter{adapterName: "provider", adapters: map[string]runtime.Provider{"provider": adapter}}
	n, _ := makeNode(t, router)

	for _, commandType := range []iop.NodeCommandType{
		iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
		iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS,
		iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API,
	} {
		resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
			RequestId: "request", Type: commandType, Adapter: "provider", Target: "model",
		})
		if err != nil || resp.GetError() != "" {
			t.Fatalf("command %v: response=%v error=%v", commandType, resp, err)
		}
	}
}

func TestNodeCommandRejectsUnknownAgentAndShell(t *testing.T) {
	router := &fixedRouter{adapterName: "missing", adapters: map[string]runtime.Provider{}}
	n, _ := makeNode(t, router)
	for _, removedType := range []iop.NodeCommandType{100, 101} {
		resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
			RequestId: "removed", Type: removedType, Adapter: "missing",
		})
		if err != nil || !strings.Contains(resp.GetError(), "unsupported command type") {
			t.Fatalf("removed command %d: response=%v error=%v", removedType, resp, err)
		}
	}
}

func TestSessionIDIsCorrelationOnly(t *testing.T) {
	adapter := &providerCommandAdapter{}
	router := &fixedRouter{adapterName: "provider", adapters: map[string]runtime.Provider{"provider": adapter}}
	n, _ := makeNode(t, router)
	for _, runID := range []string{"run-one", "run-two"} {
		if err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
			RunId: runID, Adapter: "provider", Target: "model", SessionId: "same-correlation",
		}); err != nil {
			t.Fatal(err)
		}
	}
	adapter.mu.Lock()
	defer adapter.mu.Unlock()
	if len(adapter.runs) != 2 || adapter.runs[0].RunID == adapter.runs[1].RunID || adapter.runs[0].SessionID != adapter.runs[1].SessionID {
		t.Fatalf("runs = %#v", adapter.runs)
	}
}

func TestCancelByRunIDOnly(t *testing.T) {
	adapter := &providerCommandAdapter{started: make(chan struct{}, 1), release: make(chan struct{})}
	router := &fixedRouter{adapterName: "provider", adapters: map[string]runtime.Provider{"provider": adapter}}
	n, _ := makeNode(t, router)
	if err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
		RunId: "cancel-me", Adapter: "provider", Target: "model", Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	select {
	case <-adapter.started:
	case <-time.After(2 * time.Second):
		t.Fatal("provider did not start")
	}
	if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{RunId: "cancel-me"}); err != nil {
		t.Fatal(err)
	}
	if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{}); err == nil {
		t.Fatal("empty run id was accepted")
	}
}