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