package service_test import ( "context" "errors" "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 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) } } // 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) } }