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",
				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")
}

// 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")
}