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

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