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

package service

import (
	"context"
	"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"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

func TestRunHandleCloseIsIdempotent(t *testing.T) {
	calls := 0
	handle := &RunHandle{
		close: func() {
			calls++
		},
	}

	handle.Close()
	handle.Close()

	if calls != 1 {
		t.Fatalf("close called %d times, want 1", calls)
	}
}

func waitForCondition(t *testing.T, cond func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal(msg)
}

func inflightRunCount(q *modelQueueManager) int {
	q.mu.Lock()
	defer q.mu.Unlock()
	return len(q.inflightByRun)
}

// TestRouteProviderTunnelFrameDoesNotPublishToBus verifies that tunnel frames
// reach only the request-bound tunnel subscriber and never the run event bus
// fanout (negative test for the passthrough/no-sideband boundary).
func TestRouteProviderTunnelFrameDoesNotPublishToBus(t *testing.T) {
	bus := edgeevents.NewBus()
	svc := New(edgenode.NewRegistry(), bus)

	runCh, unsubRun := bus.SubscribeAllRuns(8)
	defer unsubRun()

	tunnelCh, unsubTunnel := svc.tunnels.subscribe("tunnel-neg-1", 8)
	defer unsubTunnel()

	frame := &iop.ProviderTunnelFrame{
		RunId:    "run-neg-1",
		TunnelId: "tunnel-neg-1",
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
		Body:     []byte("data: raw provider bytes\n\n"),
	}
	svc.RouteProviderTunnelFrame(frame)

	select {
	case got := <-tunnelCh:
		if string(got.GetBody()) != string(frame.GetBody()) {
			t.Fatalf("tunnel frame body mismatch: %q", got.GetBody())
		}
	case <-time.After(time.Second):
		t.Fatal("tunnel subscriber did not receive the frame")
	}

	select {
	case e := <-runCh:
		t.Fatalf("tunnel frame leaked into run event bus: %+v", e)
	case <-time.After(100 * time.Millisecond):
	}
}

// TestRouteProviderTunnelFrameWithoutSubscriberIsDropped verifies routing a
// frame with no subscriber neither blocks nor panics.
func TestRouteProviderTunnelFrameWithoutSubscriberIsDropped(t *testing.T) {
	svc := New(edgenode.NewRegistry(), nil)
	svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{TunnelId: "no-sub"})
}

type providerTunnelTestEnv struct {
	svc        *Service
	nodeClient *toki.TcpClient
	capturedMu sync.Mutex
	captured   *iop.ProviderTunnelRequest
}

func (e *providerTunnelTestEnv) capturedRequest() *iop.ProviderTunnelRequest {
	e.capturedMu.Lock()
	defer e.capturedMu.Unlock()
	return e.captured
}

// newProviderTunnelTestEnv wires a provider-pool service to a fake node over
// net.Pipe and captures the ProviderTunnelRequest the node receives.
func newProviderTunnelTestEnv(t *testing.T) *providerTunnelTestEnv {
	t.Helper()
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	env := &providerTunnelTestEnv{nodeClient: nodeClient}
	toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
		env.capturedMu.Lock()
		env.captured = req
		env.capturedMu.Unlock()
	})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pool",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-vllm-01",
				Adapter:  "vllm-gpu",
				Type:     "vllm",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 1,
			},
		},
	})

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

	svc := New(reg, edgeevents.NewBus())
	svc.SetNodeStore(store)
	svc.SetModelCatalog([]config.ModelCatalogEntry{
		{ID: "qwen3.6:35b", Providers: map[string]string{"prov-vllm-01": "served-qwen"}},
	})
	env.svc = svc
	return env
}

// TestSubmitProviderTunnelProviderPoolSendsRequestAndReleasesSlotOnEnd covers
// the provider-pool tunnel path end to end: admission rewrites adapter/target,
// BuildBody receives the served target, ordered frames reach the handle, and
// the admission slot is released on the END frame.
func TestSubmitProviderTunnelProviderPoolSendsRequestAndReleasesSlotOnEnd(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	var builtTarget string
	handle, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{
		RunID:         "run-tunnel-001",
		ModelGroupKey: "qwen3.6:35b",
		Method:        "POST",
		Path:          "/v1/chat/completions",
		BuildBody: func(target string) ([]byte, error) {
			builtTarget = target
			return []byte(`{"model":"` + target + `","stream":true}`), nil
		},
		Stream:       true,
		ProviderPool: true,
	})
	if err != nil {
		t.Fatalf("SubmitProviderTunnel: %v", err)
	}
	defer handle.Close()

	if builtTarget != "served-qwen" {
		t.Errorf("BuildBody target: got %q, want served-qwen", builtTarget)
	}
	dispatch := handle.Dispatch()
	if dispatch.Adapter != "vllm-gpu" || dispatch.Target != "served-qwen" {
		t.Errorf("dispatch adapter/target: got %q/%q", dispatch.Adapter, dispatch.Target)
	}

	waitForCondition(t, func() bool { return env.capturedRequest() != nil },
		"fake node did not receive ProviderTunnelRequest")
	captured := env.capturedRequest()
	if captured.GetAdapter() != "vllm-gpu" || captured.GetTarget() != "served-qwen" {
		t.Errorf("wire adapter/target: got %q/%q", captured.GetAdapter(), captured.GetTarget())
	}
	if !strings.Contains(string(captured.GetBody()), `"model":"served-qwen"`) {
		t.Errorf("wire body missing served model: %s", captured.GetBody())
	}
	if inflightRunCount(svc.queue) != 1 {
		t.Fatalf("expected 1 inflight run after dispatch, got %d", inflightRunCount(svc.queue))
	}

	tunnelID := captured.GetTunnelId()
	frames := []*iop.ProviderTunnelFrame{
		{RunId: captured.GetRunId(), TunnelId: tunnelID, Sequence: 0, Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200},
		{RunId: captured.GetRunId(), TunnelId: tunnelID, Sequence: 1, Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte("data: chunk-1\n\n")},
		{RunId: captured.GetRunId(), TunnelId: tunnelID, Sequence: 2, Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true},
	}
	for _, f := range frames {
		svc.RouteProviderTunnelFrame(f)
	}

	stream := handle.Stream()
	for i, want := range frames {
		select {
		case got := <-stream.Frames:
			if got.GetSequence() != want.GetSequence() || got.GetKind() != want.GetKind() {
				t.Fatalf("frame %d: got seq=%d kind=%v", i, got.GetSequence(), got.GetKind())
			}
		case <-time.After(time.Second):
			t.Fatalf("frame %d not delivered", i)
		}
	}
	select {
	case _, open := <-stream.Frames:
		if open {
			t.Fatal("expected frame channel to close after END")
		}
	case <-time.After(time.Second):
		t.Fatal("frame channel not closed after END")
	}

	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after END frame")
}

// TestSubmitProviderTunnelCloseReleasesSlot verifies the cancel path: closing
// the handle before a terminal frame releases the provider-pool slot.
func TestSubmitProviderTunnelCloseReleasesSlot(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	handle, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{
		RunID:         "run-tunnel-cancel",
		ModelGroupKey: "qwen3.6:35b",
		Method:        "POST",
		Path:          "/v1/chat/completions",
		Body:          []byte(`{"model":"qwen3.6:35b"}`),
		ProviderPool:  true,
	})
	if err != nil {
		t.Fatalf("SubmitProviderTunnel: %v", err)
	}
	if inflightRunCount(svc.queue) != 1 {
		t.Fatalf("expected 1 inflight run after dispatch, got %d", inflightRunCount(svc.queue))
	}

	handle.Close()

	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after Close")
}

// TestProviderTunnelRouterBackpressureBlocksWithoutDrop verifies the router
// blocks the producer while the subscriber buffer is full instead of dropping
// frames, and delivers every frame in order once the subscriber drains.
func TestProviderTunnelRouterBackpressureBlocksWithoutDrop(t *testing.T) {
	svc := New(edgenode.NewRegistry(), nil)
	ch, unsub := svc.tunnels.subscribe("tunnel-bp", 2)
	defer unsub()

	const total = 6
	routed := make(chan struct{})
	go func() {
		defer close(routed)
		for i := 0; i < total; i++ {
			svc.tunnels.route(&iop.ProviderTunnelFrame{
				TunnelId: "tunnel-bp",
				Sequence: int64(i),
				Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
			})
		}
	}()

	// With buffer 2 the producer cannot finish all routes before the
	// subscriber drains: full buffer must block, not drop.
	select {
	case <-routed:
		t.Fatal("route calls finished while the subscriber buffer was full; expected blocking backpressure")
	case <-time.After(100 * time.Millisecond):
	}

	for i := 0; i < total; i++ {
		select {
		case f := <-ch:
			if f.GetSequence() != int64(i) {
				t.Fatalf("frame %d: got sequence %d; frames must stay ordered and lossless under backpressure", i, f.GetSequence())
			}
		case <-time.After(time.Second):
			t.Fatalf("frame %d not delivered after draining", i)
		}
	}

	select {
	case <-routed:
	case <-time.After(time.Second):
		t.Fatal("producer still blocked after all frames were drained")
	}
}

// TestProviderTunnelRouterCloseUnblocksBlockedRoute verifies unsubscribe (the
// handle Close/cancel path) releases a route call blocked on a full buffer, so
// an abandoned tunnel can never wedge the transport frame reader.
func TestProviderTunnelRouterCloseUnblocksBlockedRoute(t *testing.T) {
	svc := New(edgenode.NewRegistry(), nil)
	_, unsub := svc.tunnels.subscribe("tunnel-close", 1)

	frame := func(seq int64) *iop.ProviderTunnelFrame {
		return &iop.ProviderTunnelFrame{
			TunnelId: "tunnel-close",
			Sequence: seq,
			Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
		}
	}

	blocked := make(chan struct{})
	go func() {
		defer close(blocked)
		svc.tunnels.route(frame(0)) // fills the buffer
		svc.tunnels.route(frame(1)) // blocks until close releases it
	}()

	select {
	case <-blocked:
		t.Fatal("second route must block on the full subscriber buffer")
	case <-time.After(100 * time.Millisecond):
	}

	unsub()

	select {
	case <-blocked:
	case <-time.After(time.Second):
		t.Fatal("unsubscribe did not release the blocked route call")
	}

	// A late frame for the closed tunnel is dropped without blocking.
	svc.tunnels.route(frame(2))
}

// TestProviderTunnelReleaseExactlyOnceAcrossEndAndClose verifies the tracked
// provider-pool slot (capacity 1) is released exactly once when both the
// terminal END frame and handle Close fire, and that the freed slot admits the
// next tunnel immediately instead of queueing.
func TestProviderTunnelReleaseExactlyOnceAcrossEndAndClose(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	submit := func(runID string) ProviderTunnelResult {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		handle, err := svc.SubmitProviderTunnel(ctx, SubmitProviderTunnelRequest{
			RunID:         runID,
			ModelGroupKey: "qwen3.6:35b",
			Method:        "POST",
			Path:          "/v1/chat/completions",
			Body:          []byte(`{"model":"qwen3.6:35b"}`),
			ProviderPool:  true,
		})
		if err != nil {
			t.Fatalf("SubmitProviderTunnel(%s): %v", runID, err)
		}
		return handle
	}

	handle1 := submit("run-release-1")
	waitForCondition(t, func() bool {
		captured := env.capturedRequest()
		return captured != nil && captured.GetRunId() == "run-release-1"
	}, "fake node did not receive the first ProviderTunnelRequest")
	captured := env.capturedRequest()

	svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{
		RunId:    captured.GetRunId(),
		TunnelId: captured.GetTunnelId(),
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
		End:      true,
	})
	stream := handle1.Stream()
	for {
		if _, open := <-stream.Frames; !open {
			break
		}
	}
	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after END frame")

	// Close after the terminal frame must stay a no-op release.
	handle1.Close()
	if got := inflightRunCount(svc.queue); got != 0 {
		t.Fatalf("inflight count after END+Close: got %d, want 0", got)
	}

	// The capacity-1 slot is free again: the next tunnel is admitted without
	// waiting on the queue (a leaked slot would block this submit).
	handle2 := submit("run-release-2")
	defer handle2.Close()
	waitForCondition(t, func() bool {
		captured := env.capturedRequest()
		return captured != nil && captured.GetRunId() == "run-release-2"
	}, "fake node did not receive the second ProviderTunnelRequest after slot release")
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("inflight count after second dispatch: got %d, want 1", got)
	}
}

// TestSubmitProviderTunnelErrorFrameReleasesSlot verifies the ERROR frame is
// terminal and releases the provider-pool slot like END.
func TestSubmitProviderTunnelErrorFrameReleasesSlot(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	handle, err := svc.SubmitProviderTunnel(context.Background(), SubmitProviderTunnelRequest{
		RunID:         "run-tunnel-error",
		ModelGroupKey: "qwen3.6:35b",
		Method:        "POST",
		Path:          "/v1/chat/completions",
		Body:          []byte(`{"model":"qwen3.6:35b"}`),
		ProviderPool:  true,
	})
	if err != nil {
		t.Fatalf("SubmitProviderTunnel: %v", err)
	}
	defer handle.Close()

	waitForCondition(t, func() bool { return env.capturedRequest() != nil },
		"fake node did not receive ProviderTunnelRequest")
	captured := env.capturedRequest()

	svc.RouteProviderTunnelFrame(&iop.ProviderTunnelFrame{
		RunId:    captured.GetRunId(),
		TunnelId: captured.GetTunnelId(),
		Kind:     iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
		Error:    "provider unavailable",
	})

	stream := handle.Stream()
	select {
	case got := <-stream.Frames:
		if got.GetKind() != iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR {
			t.Fatalf("expected ERROR frame, got %v", got.GetKind())
		}
	case <-time.After(time.Second):
		t.Fatal("ERROR frame not delivered")
	}

	waitForCondition(t, func() bool { return inflightRunCount(svc.queue) == 0 },
		"admission slot not released after ERROR frame")
}

// TestSubmitProviderPoolSelectsTunnelProviderAndReleasesSlot verifies that
// SubmitProviderPool with an OpenAI-compatible provider (tunnel path) performs
// a single admission, sends a ProviderTunnelRequest, and releases the slot on
// tunnel END frame.
func TestSubmitProviderPoolSelectsTunnelProviderAndReleasesSlot(t *testing.T) {
	env := newProviderTunnelTestEnv(t)
	svc := env.svc

	// Catalog entry with a vLLM provider (tunnel path).
	catalog := []config.ModelCatalogEntry{
		{ID: "qwen3.6:35b", Providers: map[string]string{"prov-vllm-01": "served-qwen"}},
	}
	svc.SetModelCatalog(catalog)

	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: "qwen3.6:35b",
			ProviderPool:  true,
			Background:    true,
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil result")
	}
	if result.Path != ProviderPoolPathTunnel {
		t.Fatalf("expected path=tunnel, got %q", result.Path)
	}
	if result.Tunnel == nil {
		t.Fatal("expected tunnel handle for tunnel path")
	}
	if result.Run != nil {
		t.Fatal("expected nil run for tunnel path")
	}

	// Verify inflight after dispatch.
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("expected 1 inflight after dispatch, got %d", got)
	}
}

// TestSubmitProviderPoolSelectsNormalizedProviderAndDoesNotOpenTunnel verifies
// that SubmitProviderPool with an Ollama provider (normalized path) performs
// a single admission, sends a RunRequest (not ProviderTunnelRequest), and
// the tunnel handle is nil.
func TestSubmitProviderPoolSelectsNormalizedProviderAndDoesNotOpenTunnel(t *testing.T) {
	// Build a provider-pool service with an Ollama provider.
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

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

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

	// Capture the RunRequest received by the fake node.
	var capturedRunReq *iop.RunRequest
	var capturedMu sync.Mutex
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		capturedMu.Lock()
		capturedRunReq = req
		capturedMu.Unlock()
	})

	// Build NodeStore with an Ollama provider (normalized path).
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ollama",
		Runtime: config.RuntimeConf{Concurrency: 2},
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-ollama-01",
				Adapter:  "ollama",
				Type:     "ollama",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Build catalog with Ollama provider.
	catalog := []config.ModelCatalogEntry{
		{ID: "qwen3.6:35b", Providers: map[string]string{"prov-ollama-01": "served-qwen"}},
	}

	// Build registry with the fake node.
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-ollama",
		LifecycleState: edgenode.LifecycleConnected,
		Client:         edgeClient,
	})

	// Create Service with queue and catalog.
	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// SubmitProviderPool with Ollama provider → normalized path.
	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: "qwen3.6:35b",
			ProviderPool:  true,
			Background:    true,
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil result")
	}
	if result.Path != ProviderPoolPathNormalized {
		t.Fatalf("expected path=normalized, got %q", result.Path)
	}
	if result.Run == nil {
		t.Fatal("expected run handle for normalized path")
	}
	if result.Tunnel != nil {
		t.Fatal("expected nil tunnel for normalized path")
	}

	// Wait for the fake node to receive the RunRequest.
	time.Sleep(50 * time.Millisecond)

	capturedMu.Lock()
	defer capturedMu.Unlock()

	if capturedRunReq == nil {
		t.Fatal("no RunRequest captured from fake node; SubmitProviderPool did not send")
	}

	// Verify that the adapter and target were rewritten from the provider-pool candidate.
	if capturedRunReq.GetAdapter() != "ollama" {
		t.Errorf("adapter: got %q, want %q", capturedRunReq.GetAdapter(), "ollama")
	}
	if capturedRunReq.GetTarget() != "served-qwen" {
		t.Errorf("target: got %q, want %q", capturedRunReq.GetTarget(), "served-qwen")
	}

	// Verify inflight after dispatch.
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("expected 1 inflight after dispatch, got %d", got)
	}
}

// TestSubmitProviderPoolUsesSingleAdmissionForMixedCandidates verifies that
// SubmitProviderPool performs exactly one queue admission when presented with
// mixed tunnel+normalized candidates (e.g., vLLM + Ollama). The selected
// candidate (by in-flight/priority rotation) determines the path.
func TestSubmitProviderPoolUsesSingleAdmissionForMixedCandidates(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

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

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

	var capturedRunReq *iop.RunRequest
	var capturedMu sync.Mutex
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		capturedMu.Lock()
		capturedRunReq = req
		capturedMu.Unlock()
	})

	// Build NodeStore with two providers: vLLM (tunnel) and Ollama (normalized).
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-mixed",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-vllm",
				Adapter:  "vllm-gpu",
				Type:     "vllm",
				Models:   []string{"served-model"},
				Health:   "available",
				Capacity: 2,
			},
			{
				ID:       "prov-ollama",
				Adapter:  "ollama",
				Type:     "ollama",
				Models:   []string{"served-model"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Catalog with both providers.
	catalog := []config.ModelCatalogEntry{
		{ID: "mixed-model", Providers: map[string]string{"prov-vllm": "served-model", "prov-ollama": "served-model"}},
	}

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

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

	// SubmitProviderPool → should do exactly one admission.
	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: "mixed-model",
			ProviderPool:  true,
			Background:    true,
		},
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil result")
	}

	// The selected provider depends on admission order (vLLM or Ollama by
	// providerID alphabetical tie-break). Verify only one request was sent.
	time.Sleep(50 * time.Millisecond)

	capturedMu.Lock()
	defer capturedMu.Unlock()

	if capturedRunReq == nil {
		t.Fatal("expected exactly one RunRequest for normalized path; admission must be single")
	}

	// Verify the request target matches the catalog served model.
	if capturedRunReq.GetTarget() != "served-model" {
		t.Errorf("target: got %q, want %q", capturedRunReq.GetTarget(), "served-model")
	}

	// Verify inflight is exactly 1 (single admission).
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("expected 1 inflight (single admission), got %d", got)
	}
}

// TestSubmitProviderPoolDispatchInfoObservation verifies that the provider-pool
// one-shot dispatch carries selected provider id, provider type, and execution
// path in RunDispatch on both tunnel and normalized paths (SURFACE_OBS-1).
func TestSubmitProviderPoolDispatchInfoObservation(t *testing.T) {
	for _, tc := range []struct {
		name            string
		providerType    string
		wantPath        providerPoolPath
		wantExecPath    string
	}{
		{
			name:         "vllm_tunnel",
			providerType: "vllm",
			wantPath:     ProviderPoolPathTunnel,
			wantExecPath: "provider_tunnel",
		},
		{
			name:         "ollama_normalized",
			providerType: "ollama",
			wantPath:     ProviderPoolPathNormalized,
			wantExecPath: "normalized",
		},
		{
			name:         "cli_normalized",
			providerType: "cli",
			wantPath:     ProviderPoolPathNormalized,
			wantExecPath: "normalized",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			edgeConn, nodeConn := net.Pipe()
			t.Cleanup(func() {
				_ = edgeConn.Close()
				_ = nodeConn.Close()
			})

			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.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
					m := &iop.ProviderTunnelRequest{}
					return m, proto.Unmarshal(b, m)
				},
			}

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

			store := edgenode.NewNodeStore()
			store.Add(&edgenode.NodeRecord{
				ID:      "node-obs",
				Runtime: config.RuntimeConf{Concurrency: 2},
				Adapters: config.AdaptersConf{
					VllmInstances:       []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
					OllamaInstances:     []config.OllamaInstanceConf{{Name: "ollama", Enabled: true}},
					OpenAICompatInstances: []config.OpenAICompatInstanceConf{{Name: "openai", Enabled: true}},
				},
				Providers: []config.NodeProviderConf{
					{
						ID:       "prov-obs",
						Adapter:  "vllm-gpu",
						Type:     tc.providerType,
						Models:   []string{"served-model"},
						Health:   "available",
						Capacity: 2,
					},
				},
			})

			catalog := []config.ModelCatalogEntry{
				{ID: "obs-model", Providers: map[string]string{"prov-obs": "served-model"}},
			}

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

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

			result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
				Run: SubmitRunRequest{
					ModelGroupKey: "obs-model",
					ProviderPool:  true,
					Background:    true,
				},
			})
			if err != nil {
				t.Fatalf("SubmitProviderPool: %v", err)
			}
			if result == nil {
				t.Fatal("expected non-nil result")
			}
			if result.Path != tc.wantPath {
				t.Fatalf("path: got %q, want %q", result.Path, tc.wantPath)
			}

			disp := result.DispatchInfo
			if disp.ProviderID != "prov-obs" {
				t.Errorf("provider_id: got %q, want %q", disp.ProviderID, "prov-obs")
			}
			if disp.ProviderType != tc.providerType {
				t.Errorf("provider_type: got %q, want %q", disp.ProviderType, tc.providerType)
			}
			if disp.ExecutionPath != tc.wantExecPath {
				t.Errorf("execution_path: got %q, want %q", disp.ExecutionPath, tc.wantExecPath)
			}
			if disp.QueueReason != "dispatched" {
				t.Errorf("queue_reason: got %q, want %q", disp.QueueReason, "dispatched")
			}
			if disp.Adapter != "vllm-gpu" {
				t.Errorf("adapter: got %q, want %q", disp.Adapter, "vllm-gpu")
			}
			if disp.Target != "served-model" {
				t.Errorf("target: got %q, want %q", disp.Target, "served-model")
			}

			// Also verify tunnel/normalized handle DispatchInfo matches.
			switch result.Path {
			case ProviderPoolPathTunnel:
				if result.Tunnel == nil {
					t.Fatal("expected tunnel handle for tunnel path")
				}
				hDisp := result.Tunnel.Dispatch()
				if hDisp.ProviderID != "prov-obs" {
					t.Errorf("tunnel handle provider_id: got %q, want %q", hDisp.ProviderID, "prov-obs")
				}
				if hDisp.ProviderType != tc.providerType {
					t.Errorf("tunnel handle provider_type: got %q, want %q", hDisp.ProviderType, tc.providerType)
				}
				if hDisp.ExecutionPath != tc.wantExecPath {
					t.Errorf("tunnel handle execution_path: got %q, want %q", hDisp.ExecutionPath, tc.wantExecPath)
				}
				// Close releases the inflight slot.
				result.Tunnel.Close()
				if got := inflightRunCount(svc.queue); got != 0 {
					t.Fatalf("expected 0 inflight after tunnel close, got %d", got)
				}
			case ProviderPoolPathNormalized:
				if result.Run == nil {
					t.Fatal("expected run handle for normalized path")
				}
				hDisp := result.Run.Dispatch()
				if hDisp.ProviderID != "prov-obs" {
					t.Errorf("run handle provider_id: got %q, want %q", hDisp.ProviderID, "prov-obs")
				}
				if hDisp.ProviderType != tc.providerType {
					t.Errorf("run handle provider_type: got %q, want %q", hDisp.ProviderType, tc.providerType)
				}
				if hDisp.ExecutionPath != tc.wantExecPath {
					t.Errorf("run handle execution_path: got %q, want %q", hDisp.ExecutionPath, tc.wantExecPath)
				}
				// Inflight should be 1 for normalized (run hasn't completed yet).
				if got := inflightRunCount(svc.queue); got != 1 {
					t.Fatalf("expected 1 inflight for normalized, got %d", got)
				}
				result.Run.Close()
			}
		})
	}
}