iop/apps/edge/internal/service/service_test.go

package service_test

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func TestBuildRunRequestNormalizesSessionAndTimeout(t *testing.T) {
	req, runID, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter: "cli",
		Target:  "codex",
		Prompt:  "hello",
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	if runID == "" || req.GetRunId() != runID {
		t.Fatalf("run id mismatch: runID=%q req=%q", runID, req.GetRunId())
	}
	if req.GetSessionId() != edgeservice.DefaultSessionID {
		t.Fatalf("SessionId: got %q want %q", req.GetSessionId(), edgeservice.DefaultSessionID)
	}
	if req.GetTimeoutSec() != edgeservice.DefaultTimeoutSec {
		t.Fatalf("TimeoutSec: got %d want %d", req.GetTimeoutSec(), edgeservice.DefaultTimeoutSec)
	}
	if req.GetSessionMode() != iop.RunSessionMode_RUN_SESSION_MODE_CREATE_IF_MISSING {
		t.Fatalf("SessionMode: got %v", req.GetSessionMode())
	}
}

func TestBuildNodeCommandRequest_NewTypes(t *testing.T) {
	cases := []struct {
		name    string
		cmdType iop.NodeCommandType
		prefix  string
	}{
		{"capabilities", iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "caps"},
		{"sessions", iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, "sessions"},
		{"transport", iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, "transport"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			req := edgeservice.BuildNodeCommandRequest(tc.cmdType, tc.prefix, "cli", "codex", "", 0)
			if req.GetType() != tc.cmdType {
				t.Errorf("Type: got %v want %v", req.GetType(), tc.cmdType)
			}
			if req.GetAdapter() != "cli" {
				t.Errorf("Adapter: got %q want cli", req.GetAdapter())
			}
			if req.GetTarget() != "codex" {
				t.Errorf("Target: got %q want codex", req.GetTarget())
			}
			if req.GetSessionId() != edgeservice.DefaultSessionID {
				t.Errorf("SessionId: got %q want default", req.GetSessionId())
			}
			if req.GetTimeoutSec() != edgeservice.DefaultTimeoutSec {
				t.Errorf("TimeoutSec: got %d want default", req.GetTimeoutSec())
			}
			if req.GetRequestId() == "" {
				t.Error("RequestId: empty")
			}
		})
	}
}

func TestBuildUsageStatusRequestAndWaitTimeout(t *testing.T) {
	req := edgeservice.BuildUsageStatusRequest("cli", "antigravity", "", 0)
	if req.GetSessionId() != edgeservice.DefaultSessionID {
		t.Fatalf("SessionId: got %q want %q", req.GetSessionId(), edgeservice.DefaultSessionID)
	}
	if req.GetTimeoutSec() != edgeservice.DefaultTimeoutSec {
		t.Fatalf("TimeoutSec: got %d want %d", req.GetTimeoutSec(), edgeservice.DefaultTimeoutSec)
	}
	if got, want := edgeservice.StatusWaitTimeout(req), 35*time.Second; got != want {
		t.Fatalf("StatusWaitTimeout: got %v want %v", got, want)
	}
}

func TestBuildRunRequest_SessionAndBackground(t *testing.T) {
	req, runID, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter:    "cli",
		Target:     "codex",
		SessionID:  "session-a",
		Background: true,
		TimeoutSec: 30,
		Prompt:     "hello",
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	if runID == "" {
		t.Fatal("expected non-empty runID")
	}
	if req.GetSessionId() != "session-a" {
		t.Errorf("SessionId: got %q want %q", req.GetSessionId(), "session-a")
	}
	if !req.GetBackground() {
		t.Error("Background: expected true")
	}
}

func TestResolveNode_RequiresExplicitSelectionForMultipleNodes(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
	svc := edgeservice.New(reg, nil)

	if _, err := svc.ResolveNode(""); err == nil {
		t.Error("expected error for implicit resolve with multiple nodes")
	}
	if e, err := svc.ResolveNode("node-1"); err != nil || e.NodeID != "node-1" {
		t.Errorf("failed explicit resolve: %v", err)
	}
}

func TestBuildRunRequestDoesNotInjectConsoleSource(t *testing.T) {
	req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter:  "cli",
		Target:   "codex",
		Prompt:   "hello",
		Metadata: nil,
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	if req.GetMetadata() == nil {
		t.Fatal("expected non-nil metadata map")
	}
	if v, ok := req.GetMetadata()["source"]; ok {
		t.Errorf("BuildRunRequest should not inject source, got source=%q", v)
	}
}

func TestBuildRunRequestCopiesMetadata(t *testing.T) {
	req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter:  "cli",
		Target:   "codex",
		Prompt:   "hello",
		Metadata: map[string]string{"request_id": "req-001", "x-custom": "value"},
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	if got := req.GetMetadata()["request_id"]; got != "req-001" {
		t.Errorf("request_id: got %q, want %q", got, "req-001")
	}
	if got := req.GetMetadata()["x-custom"]; got != "value" {
		t.Errorf("x-custom: got %q, want %q", got, "value")
	}
}

func TestBuildRunRequestPreservesResponsesMetadataKeys(t *testing.T) {
	req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter: "ollama",
		Target:  "llama",
		Prompt:  "test",
		Metadata: map[string]string{
			"request_id": "req-001",
			"task_id":    "task-123",
		},
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	for _, key := range []string{"request_id", "task_id"} {
		if req.GetMetadata()[key] == "" {
			t.Errorf("metadata key %q not preserved", key)
		}
	}
}

func TestBuildRunRequestPreservesWorkspace(t *testing.T) {
	req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter:   "cli",
		Target:    "codex",
		Workspace: "/config/workspace/iop",
		Prompt:    "test",
		Metadata:  map[string]string{"request_id": "req-001"},
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	if req.GetWorkspace() != "/config/workspace/iop" {
		t.Fatalf("Workspace: got %q", req.GetWorkspace())
	}
	if _, ok := req.GetMetadata()["workspace"]; ok {
		t.Fatal("workspace should be preserved as RunRequest.Workspace, not metadata")
	}
}

func TestListNodesReturnsSnapshots(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alpha"})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
	svc := edgeservice.New(reg, nil)

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 2 {
		t.Fatalf("ListNodeSnapshots: got %d want 2", len(snaps))
	}
	byID := map[string]edgeservice.NodeSnapshot{}
	for _, s := range snaps {
		byID[s.NodeID] = s
	}
	if got := byID["node-1"]; got.Alias != "alpha" || got.Label != "node0" {
		t.Errorf("node-1 snapshot: %+v, want Label=node0", got)
	}
	if got := byID["node-2"]; got.Alias != "" || got.Label != "node1" {
		t.Errorf("node-2 snapshot: %+v, want Label=node1", got)
	}
}

func TestResolveNodeSnapshotReturnsDTO(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alpha"})
	svc := edgeservice.New(reg, nil)

	snap, err := svc.ResolveNodeSnapshot("alpha")
	if err != nil {
		t.Fatalf("ResolveNodeSnapshot: %v", err)
	}
	if snap.NodeID != "node-1" || snap.Label != "node0" {
		t.Errorf("snapshot: %+v", snap)
	}
}

func TestSubmitRunReturnsDispatchMetadata(t *testing.T) {
	dispatch := edgeservice.RunDispatch{
		RunID:         "run-x",
		NodeID:        "node-1",
		NodeLabel:     "alpha",
		ModelGroupKey: "codex-model",
		Adapter:       "cli",
		Target:        "codex",
		SessionID:     "session-a",
		Background:    true,
		TimeoutSec:    30,
	}
	handle := &edgeservice.RunHandle{RunDispatch: dispatch}
	if handle.RunID != dispatch.RunID || handle.NodeLabel != dispatch.NodeLabel {
		t.Fatalf("RunHandle does not expose embedded RunDispatch fields: %+v", handle)
	}
	if handle.ModelGroupKey != "codex-model" {
		t.Errorf("RunHandle ModelGroupKey: got %q want codex-model", handle.ModelGroupKey)
	}
	if handle.Background != true || handle.TimeoutSec != 30 {
		t.Errorf("RunHandle dispatch fields wrong: %+v", handle)
	}
}

func TestBuildCancelRunRequest(t *testing.T) {
	cases := []struct {
		name          string
		req           edgeservice.CancelRunRequest
		wantSessionID string
	}{
		{
			name: "explicit session",
			req: edgeservice.CancelRunRequest{
				RunID:     "run-123",
				Adapter:   "cli",
				Target:    "claude",
				SessionID: "session-a",
			},
			wantSessionID: "session-a",
		},
		{
			name: "empty session normalizes to default",
			req: edgeservice.CancelRunRequest{
				RunID:     "run-456",
				Adapter:   "ollama",
				Target:    "llama3",
				SessionID: "",
			},
			wantSessionID: edgeservice.DefaultSessionID,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			proto := edgeservice.BuildCancelRunRequest(tc.req)
			if proto.GetRunId() != tc.req.RunID {
				t.Errorf("RunId: got %q want %q", proto.GetRunId(), tc.req.RunID)
			}
			if proto.GetAdapter() != tc.req.Adapter {
				t.Errorf("Adapter: got %q want %q", proto.GetAdapter(), tc.req.Adapter)
			}
			if proto.GetTarget() != tc.req.Target {
				t.Errorf("Target: got %q want %q", proto.GetTarget(), tc.req.Target)
			}
			if proto.GetSessionId() != tc.wantSessionID {
				t.Errorf("SessionId: got %q want %q", proto.GetSessionId(), tc.wantSessionID)
			}
			if proto.GetAction() != iop.CancelAction_CANCEL_ACTION_CANCEL_RUN {
				t.Errorf("Action: got %v want CANCEL_ACTION_CANCEL_RUN", proto.GetAction())
			}
		})
	}
}

func TestResolveNode_AllowsSingleNodeFallback(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
	svc := edgeservice.New(reg, nil)

	if e, err := svc.ResolveNode(""); err != nil || e.NodeID != "node-1" {
		t.Errorf("expected single node fallback, got error: %v", err)
	}
}

func TestListNodeSnapshotsWithConfig(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alpha"})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:    "node-1",
		Alias: "alpha",
		Token: "tok-1",
		Runtime: config.RuntimeConf{
			Concurrency: 3,
		},
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama", Enabled: true, BaseURL: "http://localhost:11434", ContextSize: 2048},
			},
		},
	})

	svc := edgeservice.New(reg, nil)
	svc.SetNodeStore(store)

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("ListNodeSnapshots: got %d want 1", len(snaps))
	}
	snap := snaps[0]
	if snap.NodeID != "node-1" || snap.Label != "node0" {
		t.Errorf("snapshot basic info wrong: %+v", snap)
	}
	if snap.Config == nil {
		t.Fatal("expected snap.Config to be populated")
	}
	if snap.Config.Runtime == nil || snap.Config.Runtime.Concurrency != 3 {
		t.Errorf("unexpected runtime config: %+v", snap.Config.Runtime)
	}
	var ollama *iop.AdapterConfig
	for _, a := range snap.Config.Adapters {
		if a.Type == "ollama" {
			ollama = a
			break
		}
	}
	if ollama == nil {
		t.Fatal("expected ollama adapter config")
	}
	if !ollama.Enabled || ollama.GetOllama().GetBaseUrl() != "http://localhost:11434" {
		t.Errorf("unexpected ollama config: %+v", ollama)
	}
}

func TestListNodeSnapshotsIncludesAgentKind(t *testing.T) {
	reg := edgenode.NewRegistry()
	// generic node: kind/lifecycle defaults are applied by Register.
	reg.Register(&edgenode.NodeEntry{NodeID: "node-generic"})
	// Explicit generic node: kind and a non-default lifecycle must pass through unchanged.
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-explicit",
		AgentKind:      config.AgentKindGenericNode,
		LifecycleState: "registering",
	})
	svc := edgeservice.New(reg, nil)

	byID := map[string]edgeservice.NodeSnapshot{}
	for _, s := range svc.ListNodeSnapshots() {
		byID[s.NodeID] = s
	}

	generic, ok := byID["node-generic"]
	if !ok {
		t.Fatal("node-generic snapshot missing")
	}
	if generic.AgentKind != config.AgentKindGenericNode {
		t.Errorf("generic AgentKind: got %q want %q", generic.AgentKind, config.AgentKindGenericNode)
	}
	if generic.LifecycleState != edgenode.LifecycleConnected {
		t.Errorf("generic LifecycleState: got %q want %q", generic.LifecycleState, edgenode.LifecycleConnected)
	}

	explicit, ok := byID["node-explicit"]
	if !ok {
		t.Fatal("node-explicit snapshot missing")
	}
	if explicit.AgentKind != config.AgentKindGenericNode {
		t.Errorf("explicit AgentKind: got %q want %q", explicit.AgentKind, config.AgentKindGenericNode)
	}
	if explicit.LifecycleState != "registering" {
		t.Errorf("explicit LifecycleState: got %q want %q", explicit.LifecycleState, "registering")
	}
}

func TestResolveNodeSnapshotCarriesAgentKind(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:    "node-explicit",
		Alias:     "explicit-1",
		AgentKind: config.AgentKindGenericNode,
	})
	svc := edgeservice.New(reg, nil)

	snap, err := svc.ResolveNodeSnapshot("explicit-1")
	if err != nil {
		t.Fatalf("ResolveNodeSnapshot: %v", err)
	}
	if snap.AgentKind != config.AgentKindGenericNode {
		t.Errorf("AgentKind: got %q want %q", snap.AgentKind, config.AgentKindGenericNode)
	}
	if snap.LifecycleState != edgenode.LifecycleConnected {
		t.Errorf("LifecycleState: got %q want %q", snap.LifecycleState, edgenode.LifecycleConnected)
	}
}

func TestServiceCapabilitiesAndDomainAgents(t *testing.T) {
	for _, state := range []string{"", "connected", "deploying", "running", "busy", "error", "failed", "unknown-state"} {
		t.Run(fmt.Sprintf("state-%s", state), func(t *testing.T) {
			reg := edgenode.NewRegistry()
			reg.Register(&edgenode.NodeEntry{
				NodeID:    "node-1",
				Alias:     "gen-node",
				AgentKind: config.AgentKindGenericNode,
			})
			reg.Register(&edgenode.NodeEntry{
				NodeID:         "node-2",
				Alias:          "node2",
				AgentKind:      config.AgentKindGenericNode,
				LifecycleState: state,
			})
			svc := edgeservice.New(reg, nil)

			caps := svc.GetCapabilities()
			if len(caps) != 1 {
				t.Fatalf("expected 1 capability, got %d", len(caps))
			}
			capMap := map[string]*iop.EdgeCapabilitySummary{}
			for _, c := range caps {
				capMap[c.Kind] = c
			}
			if c := capMap["run"]; c == nil || !c.Available || c.Status != "ready" {
				t.Errorf("invalid run capability: %+v", c)
			}

			agents := svc.GetDomainAgents()
			if len(agents) != 0 {
				t.Errorf("expected 0 domain agents, got %d", len(agents))
			}
		})
	}
}

func TestServiceExecuteCommand(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:    "node-1",
		Alias:     "deployer",
		AgentKind: config.AgentKindGenericNode,
	})
	svc := edgeservice.New(reg, nil)

	t.Run("health.check", func(t *testing.T) {
		var events []*iop.EdgeCommandEvent
		onEvent := func(e *iop.EdgeCommandEvent) {
			events = append(events, e)
		}
		req := &iop.EdgeCommandRequest{
			CommandId: "cmd-hc",
			Operation: "health.check",
		}
		resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
		if err != nil {
			t.Fatalf("ExecuteCommand error: %v", err)
		}
		if resp.Status != "completed" {
			t.Errorf("status: got %q want completed", resp.Status)
		}
		if len(events) != 2 {
			t.Fatalf("expected 2 events, got %d", len(events))
		}
		if events[0].Phase != "started" || events[1].Phase != "completed" {
			t.Errorf("unexpected event sequence: %+v", events)
		}
	})

	t.Run("agent.status found", func(t *testing.T) {
		var events []*iop.EdgeCommandEvent
		onEvent := func(e *iop.EdgeCommandEvent) {
			events = append(events, e)
		}
		req := &iop.EdgeCommandRequest{
			CommandId:      "cmd-status",
			Operation:      "agent.status",
			TargetSelector: "deployer",
		}
		resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
		if err != nil {
			t.Fatalf("ExecuteCommand error: %v", err)
		}
		if resp.Status != "completed" {
			t.Errorf("status: got %q want completed", resp.Status)
		}
		if len(events) != 2 {
			t.Fatalf("expected 2 events, got %d", len(events))
		}
		if events[0].Phase != "started" || events[1].Phase != "completed" {
			t.Errorf("unexpected event sequence: %+v", events)
		}
	})

	t.Run("agent.status missing", func(t *testing.T) {
		req := &iop.EdgeCommandRequest{
			CommandId:      "cmd-status",
			Operation:      "agent.status",
			TargetSelector: "missing-agent",
		}
		resp, err := svc.ExecuteCommand(context.Background(), req, func(*iop.EdgeCommandEvent) {})
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if resp.Status != "error" || resp.Error == "" {
			t.Errorf("expected error status, got: %+v", resp)
		}
	})

	t.Run("agent.command does not fake success on nil client", func(t *testing.T) {
		var events []*iop.EdgeCommandEvent
		onEvent := func(e *iop.EdgeCommandEvent) {
			events = append(events, e)
		}
		req := &iop.EdgeCommandRequest{
			CommandId:      "cmd-run",
			Operation:      "agent.command",
			TargetSelector: "deployer",
			Parameters: map[string]string{
				"command": "capabilities",
			},
		}
		resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if resp.Status != "error" {
			t.Errorf("expected error status, got: %q", resp.Status)
		}
		if len(events) < 2 {
			t.Fatalf("expected at least started and failed events, got %d", len(events))
		}
		if events[0].Phase != "started" || events[len(events)-1].Phase != "failed" {
			t.Errorf("unexpected event sequence: %+v", events)
		}
	})

	t.Run("agent.command missing command param", func(t *testing.T) {
		req := &iop.EdgeCommandRequest{
			CommandId:      "cmd-run",
			Operation:      "agent.command",
			TargetSelector: "deployer",
		}
		resp, err := svc.ExecuteCommand(context.Background(), req, func(*iop.EdgeCommandEvent) {})
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if resp.Status != "unsupported" {
			t.Errorf("expected unsupported status, got %q", resp.Status)
		}
	})

	t.Run("unsupported operation", func(t *testing.T) {
		req := &iop.EdgeCommandRequest{
			CommandId: "cmd-bad",
			Operation: "bad.op",
		}
		resp, err := svc.ExecuteCommand(context.Background(), req, func(*iop.EdgeCommandEvent) {})
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if resp.Status != "unsupported" {
			t.Errorf("status: got %q want unsupported", resp.Status)
		}
	})
}

func TestServiceExecuteCommandHealthCheckDerivesFromStatus(t *testing.T) {
	t.Run("with nodes summarizes registry and capabilities", func(t *testing.T) {
		reg := edgenode.NewRegistry()
		reg.Register(&edgenode.NodeEntry{NodeID: "node-1", AgentKind: config.AgentKindGenericNode})
		reg.Register(&edgenode.NodeEntry{NodeID: "node-2", AgentKind: config.AgentKindGenericNode})
		svc := edgeservice.New(reg, nil)

		// HealthStatus is the Edge-owned boundary the command must read from.
		status := svc.HealthStatus()
		if len(status.Nodes) != 2 {
			t.Fatalf("HealthStatus nodes: got %d want 2", len(status.Nodes))
		}
		if len(status.Capabilities) != 1 {
			t.Fatalf("HealthStatus capabilities: got %d want 1", len(status.Capabilities))
		}

		var events []*iop.EdgeCommandEvent
		resp, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
			CommandId: "cmd-hc",
			Operation: "health.check",
		}, func(e *iop.EdgeCommandEvent) { events = append(events, e) })
		if err != nil {
			t.Fatalf("ExecuteCommand error: %v", err)
		}
		if resp.Status != "completed" {
			t.Errorf("status: got %q want completed", resp.Status)
		}
		if !strings.Contains(resp.Summary, "2 node(s)") || !strings.Contains(resp.Summary, "1 capability(ies)") {
			t.Errorf("summary not derived from status: got %q", resp.Summary)
		}
		if len(events) != 2 || events[1].Phase != "completed" {
			t.Fatalf("unexpected events: %+v", events)
		}
		if events[1].Summary != resp.Summary {
			t.Errorf("completed event summary %q != response summary %q", events[1].Summary, resp.Summary)
		}
	})

	t.Run("without nodes reports empty status", func(t *testing.T) {
		svc := edgeservice.New(edgenode.NewRegistry(), nil)
		resp, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
			CommandId: "cmd-hc",
			Operation: "health.check",
		}, func(*iop.EdgeCommandEvent) {})
		if err != nil {
			t.Fatalf("ExecuteCommand error: %v", err)
		}
		if !strings.Contains(resp.Summary, "0 node(s)") || !strings.Contains(resp.Summary, "0 capability(ies)") {
			t.Errorf("empty status summary wrong: got %q", resp.Summary)
		}
	})
}

func TestServiceExecuteCommandAgentCommandCapabilitiesSuccess(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
		if req.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES {
			return nil, fmt.Errorf("unexpected command type: %v", req.GetType())
		}
		if req.GetAdapter() != "cli" || req.GetTarget() != "codex" || req.GetSessionId() != "sess-success" {
			return nil, fmt.Errorf("unexpected adapter/target/session: %s/%s/%s", req.GetAdapter(), req.GetTarget(), req.GetSessionId())
		}
		return &iop.NodeCommandResponse{
			RequestId: req.RequestId,
			Type:      req.Type,
			Result:    map[string]string{"caps": "cli,codex"},
		}, nil
	})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:    "node-success-1",
		Alias:     "deployer",
		AgentKind: config.AgentKindGenericNode,
		Client:    edgeClient,
	})
	svc := edgeservice.New(reg, nil)

	var events []*iop.EdgeCommandEvent
	onEvent := func(e *iop.EdgeCommandEvent) {
		events = append(events, e)
	}

	req := &iop.EdgeCommandRequest{
		CommandId:      "cmd-success-1",
		Operation:      "agent.command",
		TargetSelector: "deployer",
		Parameters: map[string]string{
			"command":    "capabilities",
			"adapter":    "cli",
			"target":     "codex",
			"session_id": "sess-success",
		},
	}

	resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
	if err != nil {
		t.Fatalf("ExecuteCommand: %v", err)
	}

	if resp.Status != "completed" {
		t.Errorf("status: got %q want completed", resp.Status)
	}
	if !strings.Contains(resp.Summary, "caps=cli,codex") {
		t.Errorf("summary: got %q, expected containing caps=cli,codex", resp.Summary)
	}

	if len(events) < 2 {
		t.Fatalf("expected at least 2 events, got %d", len(events))
	}
	if events[0].Phase != "started" || events[len(events)-1].Phase != "completed" {
		t.Errorf("unexpected event sequence: %+v", events)
	}
	if events[len(events)-1].Summary != "Command capabilities completed successfully" {
		t.Errorf("unexpected event summary: got %q, expected 'Command capabilities completed successfully'", events[len(events)-1].Summary)
	}
}

func TestUsageStatusPreservesAdapterIdentity(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
		return &iop.NodeCommandResponse{
			RequestId:   req.RequestId,
			Type:        req.Type,
			Adapter:     req.Adapter,
			Target:      req.Target,
			SessionId:   req.SessionId,
			UsageStatus: &iop.AgentUsageStatus{RawOutput: "ok"},
		}, nil
	})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID: "node-uid",
		Alias:  "svc-node",
		Client: edgeClient,
	})
	svc := edgeservice.New(reg, nil)

	result, err := svc.UsageStatus(context.Background(), edgeservice.UsageStatusRequest{
		NodeRef:   "node-uid",
		Adapter:   "cli@1",
		Target:    "claude",
		SessionID: "sess-u",
	})
	if err != nil {
		t.Fatalf("UsageStatus: %v", err)
	}
	if result.Adapter != "cli@1" {
		t.Errorf("Adapter: got %q want %q", result.Adapter, "cli@1")
	}
	if result.Target != "claude" {
		t.Errorf("Target: got %q want %q", result.Target, "claude")
	}
	if result.SessionID != "sess-u" {
		t.Errorf("SessionID: got %q want %q", result.SessionID, "sess-u")
	}
	if result.UsageStatus == nil || result.UsageStatus.RawOutput != "ok" {
		t.Errorf("UsageStatus: got %+v", result.UsageStatus)
	}
}

func TestServiceExecuteCommandAgentCommandCapabilitiesNodeError(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
		if req.GetSessionId() != "sess-error" {
			return nil, fmt.Errorf("unexpected session_id: %s", req.GetSessionId())
		}
		return &iop.NodeCommandResponse{
			RequestId: req.RequestId,
			Type:      req.Type,
			Error:     "node failure",
		}, nil
	})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:    "node-error-1",
		Alias:     "deployer",
		AgentKind: config.AgentKindGenericNode,
		Client:    edgeClient,
	})
	svc := edgeservice.New(reg, nil)

	var events []*iop.EdgeCommandEvent
	onEvent := func(e *iop.EdgeCommandEvent) {
		events = append(events, e)
	}

	req := &iop.EdgeCommandRequest{
		CommandId:      "cmd-error-1",
		Operation:      "agent.command",
		TargetSelector: "deployer",
		Parameters: map[string]string{
			"command":    "capabilities",
			"adapter":    "cli",
			"target":     "codex",
			"session_id": "sess-error",
		},
	}

	resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
	if err != nil {
		t.Fatalf("ExecuteCommand: %v", err)
	}

	if resp.Status != "error" {
		t.Errorf("status: got %q want error", resp.Status)
	}
	if !strings.Contains(resp.Error, "node reported error: node failure") {
		t.Errorf("error: got %q, expected containing node failure", resp.Error)
	}

	if len(events) < 2 {
		t.Fatalf("expected at least 2 events, got %d", len(events))
	}
	if events[0].Phase != "started" || events[len(events)-1].Phase != "failed" {
		t.Errorf("unexpected event sequence: %+v", events)
	}
	expectedSummary := "Command capabilities failed: node reported error: node failure"
	if events[len(events)-1].Summary != expectedSummary {
		t.Errorf("unexpected event summary: got %q, expected %q", events[len(events)-1].Summary, expectedSummary)
	}
}

func TestServiceCapabilitiesPreservesProviderSnapshots(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
		if req.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES {
			return nil, fmt.Errorf("unexpected command type: %v", req.GetType())
		}
		return &iop.NodeCommandResponse{
			RequestId: req.RequestId,
			Type:      req.Type,
			Result:    map[string]string{"caps": "cli,codex"},
			ProviderSnapshots: []*iop.ProviderSnapshot{
				{
					Adapter:  "cli",
					Status:   "available",
					Capacity: 4,
					InFlight: 2,
					Queued:   3,
				},
			},
		}, nil
	})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:    "node-caps-1",
		Alias:     "deployer",
		AgentKind: config.AgentKindGenericNode,
		Client:    edgeClient,
	})
	svc := edgeservice.New(reg, nil)

	view, err := svc.Capabilities(context.Background(), edgeservice.NodeCommandRequestSpec{
		NodeRef:   "deployer",
		Adapter:   "cli",
		Target:    "codex",
		SessionID: "sess-caps",
	})
	if err != nil {
		t.Fatalf("Capabilities: %v", err)
	}

	if len(view.ProviderSnapshots) != 1 {
		t.Fatalf("expected 1 provider snapshot, got %d", len(view.ProviderSnapshots))
	}
	snap := view.ProviderSnapshots[0]
	if snap.Adapter != "cli" || snap.Status != "available" || snap.Capacity != 4 || snap.InFlight != 2 || snap.Queued != 3 {
		t.Fatalf("unexpected snapshot contents: %+v", snap)
	}
}

// TestSubmitRunModelQueueDispatchesQueuedRunAfterDisconnectToLiveNode verifies
// that after a node disconnect, a queued run is dispatched to a remaining live
// node once that node's slot becomes available.
func TestSubmitRunModelQueueDispatchesQueuedRunAfterDisconnectToLiveNode(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn1, nodeConn1 := net.Pipe()
	edgeConn2, nodeConn2 := net.Pipe()
	defer edgeConn1.Close()
	defer nodeConn1.Close()
	defer edgeConn2.Close()
	defer nodeConn2.Close()

	edgeClient1 := toki.NewTcpClient(edgeConn1, 0, 0, parserMap)
	edgeClient2 := toki.NewTcpClient(edgeConn2, 0, 0, parserMap)
	nodeClient1 := toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	nodeClient2 := toki.NewTcpClient(nodeConn2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "dc-node-1", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "dc-node-2", Client: edgeClient2})

	catalog := []config.ModelCatalogEntry{
		{
			ID: "dc-group",
			Providers: map[string]string{
				"prov-dc-1": "dc-model",
				"prov-dc-2": "dc-model",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "dc-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-dc-1", Adapter: "mock", Models: []string{"dc-model"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID: "dc-node-2",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-dc-2", Adapter: "mock", Models: []string{"dc-model"}, Health: "available", Capacity: 1},
		},
	})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// Submit two background runs to fill both nodes.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "dc-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer res1.Close()

	res2, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "dc-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run2: %v", err)
	}
	defer res2.Close()

	dispatch1 := res1.Dispatch()
	dispatch2 := res2.Dispatch()
	if dispatch1.NodeID == "" || dispatch2.NodeID == "" || dispatch1.NodeID == dispatch2.NodeID {
		t.Fatalf("expected two different nodes; got %q and %q", dispatch1.NodeID, dispatch2.NodeID)
	}

	// Start run3 in a goroutine — will queue because both nodes are full.
	var (
		res3 edgeservice.RunResult
		err3 error
		wg   sync.WaitGroup
	)
	wg.Add(1)
	go func() {
		defer wg.Done()
		res3, err3 = svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
			ModelGroupKey: "dc-group",
			ProviderPool:  true,
			Background:    true,
		})
	}()

	// Allow run3 to enter the queue.
	time.Sleep(30 * time.Millisecond)

	// Identify which node is nd1 (to disconnect) and nd2 (to receive run3).
	nd1NodeID := dispatch1.NodeID
	nd2RunID := dispatch2.RunID
	nd2NodeID := dispatch2.NodeID

	// nd1 disconnects — its candidate is removed from run3's queue item.
	bus.PublishNode(&iop.EdgeNodeEvent{NodeId: nd1NodeID, Type: "node.disconnected"})

	// nd2's run terminates — now nd2 has capacity, run3 dispatches to nd2.
	bus.PublishRun(&iop.RunEvent{RunId: nd2RunID, Type: "complete"})

	wg.Wait()

	if err3 != nil {
		t.Fatalf("run3 error: %v", err3)
	}
	if res3 == nil {
		t.Fatal("run3: expected non-nil result")
	}
	defer res3.Close()

	if got := res3.Dispatch().NodeID; got != nd2NodeID {
		t.Errorf("run3 dispatched to %q, want %q (the live node; nd1=%q disconnected)", got, nd2NodeID, nd1NodeID)
	}
}

// TestSubmitRunModelQueueContextCancelRemovesQueuedItem verifies that cancelling
// the SubmitRun context removes the item from the queue and returns context.Canceled.
func TestSubmitRunModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "ctx-node-1", Client: edgeClient})

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "ctx-group",
			Providers: map[string]string{"prov-ctx-1": "ctx-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "ctx-node-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-ctx-1", Adapter: "mock", Models: []string{"ctx-model"}, Health: "available", Capacity: 1},
		},
	})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// Fill the node capacity with run1.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "ctx-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1: %v", err)
	}
	defer res1.Close()

	// Start run2 with a cancellable context (will queue).
	ctx, cancel := context.WithCancel(context.Background())
	var run2Err error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		_, run2Err = svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
			ModelGroupKey: "ctx-group",
			ProviderPool:  true,
			Background:    true,
		})
	}()

	// Give run2 time to enter the queue.
	time.Sleep(30 * time.Millisecond)
	cancel()

	wg.Wait()

	if run2Err == nil {
		t.Fatal("expected error from cancelled context, got nil")
	}
	if !errors.Is(run2Err, context.Canceled) {
		t.Errorf("expected context.Canceled, got: %v", run2Err)
	}
}

// TestSubmitRunModelQueueDispatchesAcrossNodes verifies that concurrent
// SubmitRun calls with the same ModelGroupKey are dispatched across multiple
// nodes when each node has capacity 1.
func TestSubmitRunModelQueueDispatchesAcrossNodes(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeCommandResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn1, nodeConn1 := net.Pipe()
	edgeConn2, nodeConn2 := net.Pipe()
	defer edgeConn1.Close()
	defer nodeConn1.Close()
	defer edgeConn2.Close()
	defer nodeConn2.Close()

	edgeClient1 := toki.NewTcpClient(edgeConn1, 0, 0, parserMap)
	edgeClient2 := toki.NewTcpClient(edgeConn2, 0, 0, parserMap)

	// Node-side clients consume RunRequest messages without responding.
	nodeClient1 := toki.NewTcpClient(nodeConn1, 0, 0, parserMap)
	nodeClient2 := toki.NewTcpClient(nodeConn2, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient1.Communicator, func(*iop.RunRequest) {})
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient2.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "mq-node-1", Client: edgeClient1})
	reg.Register(&edgenode.NodeEntry{NodeID: "mq-node-2", Client: edgeClient2})

	catalog := []config.ModelCatalogEntry{
		{
			ID: "test-group",
			Providers: map[string]string{
				"prov-mq-1": "served-model",
				"prov-mq-2": "served-model",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:    "mq-node-1",
		Token: "tok-1",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-mq-1", Adapter: "mock", Models: []string{"served-model"}, Health: "available", Capacity: 1},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID:    "mq-node-2",
		Token: "tok-2",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "mock", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-mq-2", Adapter: "mock", Models: []string{"served-model"}, Health: "available", Capacity: 1},
		},
	})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	var (
		mu      sync.Mutex
		nodeIDs []string
		errs    []error
	)
	var wg sync.WaitGroup
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			res, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
				ModelGroupKey: "test-group",
				ProviderPool:  true,
				Background:    true,
			})
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, err)
				return
			}
			nodeIDs = append(nodeIDs, res.Dispatch().NodeID)
			res.Close()
		}()
	}
	wg.Wait()

	for _, err := range errs {
		t.Fatalf("SubmitRun error: %v", err)
	}

	if len(nodeIDs) != 2 {
		t.Fatalf("expected 2 dispatches, got %d", len(nodeIDs))
	}

	usedNodes := map[string]int{}
	for _, id := range nodeIDs {
		usedNodes[id]++
	}
	if usedNodes["mq-node-1"] != 1 || usedNodes["mq-node-2"] != 1 {
		t.Errorf("expected each node used exactly once: %v", usedNodes)
	}
}

func TestSubmitRunModelQueueUsesProviderInstancePolicy(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "prov-node-1", Client: edgeClient})

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "local-group",
			Providers: map[string]string{"prov-local-1": "local-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-1",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-local-1", Adapter: "ollama-local", Models: []string{"local-model"}, Health: "available", Capacity: 2, MaxQueue: 3},
		},
	})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "local-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	res2, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "local-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run2 error: %v", err)
	}
	defer res2.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err3 := svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
		ModelGroupKey: "local-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err3 == nil {
		t.Fatal("expected third run to block and timeout (capacity=2 exceeded)")
	}
	if !errors.Is(err3, context.DeadlineExceeded) {
		t.Errorf("expected deadline exceeded, got: %v", err3)
	}
}

func TestSubmitRunModelQueueFallsBackToProviderInstancePolicy(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "prov-node-fallback", Client: edgeClient})

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "fallback-group",
			Providers: map[string]string{"prov-fallback-1": "fallback-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-fallback",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			// capacity=2, maxQueue=3 come from provider config (provider-pool always uses provider policy).
			{ID: "prov-fallback-1", Adapter: "ollama-local", Models: []string{"fallback-model"}, Health: "available", Capacity: 2, MaxQueue: 3},
		},
	})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "fallback-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	res2, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "fallback-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run2 error: %v", err)
	}
	defer res2.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err3 := svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
		ModelGroupKey: "fallback-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err3 == nil {
		t.Fatal("expected third run to block and timeout (capacity=2 exceeded)")
	}
	if !errors.Is(err3, context.DeadlineExceeded) {
		t.Errorf("expected deadline exceeded, got: %v", err3)
	}
}

func TestSubmitRunModelQueueUsesRouteQueueTimeout(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "prov-node-timeout", Client: edgeClient})

	catalog := []config.ModelCatalogEntry{
		{
			ID:        "timeout-group",
			Providers: map[string]string{"prov-timeout-1": "timeout-model"},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-timeout",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			// QueueTimeoutMS=10 comes from provider config; provider-pool ignores request-level QueueTimeoutMS.
			{ID: "prov-timeout-1", Adapter: "ollama-local", Models: []string{"timeout-model"}, Health: "available", Capacity: 1, MaxQueue: 5, QueueTimeoutMS: 10},
		},
	})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// run 1: fills capacity; provider policy sets QueueTimeoutMS=10.
	res1, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "timeout-group",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	// run 2: enters queue; times out after ~10ms per provider policy.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	start := time.Now()
	_, err2 := svc.SubmitRun(ctx, edgeservice.SubmitRunRequest{
		ModelGroupKey: "timeout-group",
		ProviderPool:  true,
		Background:    true,
	})
	elapsed := time.Since(start)

	if err2 == nil {
		t.Fatal("expected second run to fail with queue timeout")
	}
	if !strings.Contains(err2.Error(), "queue timeout") {
		t.Errorf("expected queue timeout error, got: %v", err2)
	}
	if elapsed > 1*time.Second {
		t.Errorf("expected timeout to happen quickly, took: %v", elapsed)
	}
}

// TestSubmitRunLegacyModelGroupKeyGoesDirectNotQueued verifies that a request
// with ModelGroupKey but ProviderPool=false bypasses the queue and goes directly
// to the single registered node. This proves the queue admission gate on
// req.ProviderPool: without ProviderPool=true, ModelGroupKey is decorative only
// and the service falls back to direct dispatch without touching the model catalog.
func TestSubmitRunLegacyModelGroupKeyGoesDirectNotQueued(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "leg-node-1", Client: edgeClient})

	bus := edgeevents.NewBus()
	svc := edgeservice.New(reg, bus)
	// No catalog set — if ProviderPool were true, resolveProviderPoolCandidates would
	// return "not found in catalog". Since ProviderPool=false, direct dispatch succeeds.

	res, err := svc.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "legacy-group",
		Adapter:       "cli",
		Target:        "codex",
		Background:    true,
		// ProviderPool: false (default)
	})
	if err != nil {
		t.Fatalf("expected direct dispatch to succeed, got: %v", err)
	}
	defer res.Close()

	if res.Dispatch().NodeID != "leg-node-1" {
		t.Errorf("dispatch node: got %q, want leg-node-1", res.Dispatch().NodeID)
	}
}

func TestBuildRunRequestPreservesEstimatedTokensAndContextClass(t *testing.T) {
	req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
		Adapter:              "ollama",
		Target:               "llama",
		Prompt:               "test",
		EstimatedInputTokens: 150000,
		ContextClass:         "long",
		Metadata:             map[string]string{"estimated_input_tokens": "150000", "context_class": "long"},
	})
	if err != nil {
		t.Fatalf("BuildRunRequest: %v", err)
	}
	// BuildRunRequest does not carry EstimatedInputTokens/ContextClass in the protobuf,
	// so we verify that the metadata copy is preserved.
	if req.GetMetadata()["estimated_input_tokens"] != "150000" {
		t.Fatalf("metadata estimated_input_tokens: got %q", req.GetMetadata()["estimated_input_tokens"])
	}
	if req.GetMetadata()["context_class"] != "long" {
		t.Fatalf("metadata context_class: got %q", req.GetMetadata()["context_class"])
	}
}