diff --git a/services/nvpair-engine-manager/README.md b/services/nvpair-engine-manager/README.md index ced4bd404..8ca398cc2 100644 --- a/services/nvpair-engine-manager/README.md +++ b/services/nvpair-engine-manager/README.md @@ -353,3 +353,17 @@ termination, console hiding) are the only build-tagged Go Shuts down on stdin EOF (parent closed the pipe), `SIGINT`/`SIGTERM`, or a `shutdown` JSON-RPC request — stopping any running engines first so none are orphaned. + +That engine sweep (`StopAll`) runs **exactly once per process**, and every +caller returns only once it has finished. An ordinary quit asks for it twice: +the broker sends `engine:prepare-shutdown` from its own teardown, and then +closing stdin reaches the EOF path. Each is the right trigger for a different +way of being shut down, so both stay. + +Later callers **wait** for the sweep in flight rather than returning early: +returning early would report engines stopped before they were, and the broker +would close stdin and exit while they were still running. They do not repeat +the sweep, which would re-pay the readiness probe for every adopted engine whose +stop is declined. They do retry any engine whose stop failed through a process +handle or a stop command, such as an `lms server stop` that errored, because +another attempt can succeed where a declined stop cannot. diff --git a/services/nvpair-engine-manager/executor.go b/services/nvpair-engine-manager/executor.go index f71584570..2b5c03f85 100644 --- a/services/nvpair-engine-manager/executor.go +++ b/services/nvpair-engine-manager/executor.go @@ -127,8 +127,17 @@ type Executor struct { reservedPort atomic.Int32 // StopAll is terminal for an Executor; the gate closes its start/snapshot race. shuttingDown atomic.Bool - mu sync.Mutex - engines map[string]*engineState + // stopAllOnce collapses the shutdown entry points into a single sweep. + // See StopAll. + stopAllOnce sync.Once + // failedStops holds the engines whose stop failed through a process handle + // or a stop command, which later StopAll callers retry. The sweep writes it + // without retryMu, since stopAllOnce returns to no caller before the sweep + // finishes; retryMu serializes the retries that follow. + retryMu sync.Mutex + failedStops []string + mu sync.Mutex + engines map[string]*engineState } func NewExecutor(reg *Registry, reporter *Reporter, emit func(string, any), baseDir string) *Executor { diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index e1b2c53c2..96ef0a959 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -751,7 +751,36 @@ func (e *Executor) Restart(ctx context.Context, engine string) error { // StopAll rejects new starts and terminates every running engine during // shutdown without changing the user's saved ON/OFF intent. +// +// Two entry points ask for it on an ordinary quit — the +// engine:prepare-shutdown handler, which the broker calls from its own +// teardown, and the stdin-EOF path in Run, which follows when the broker then +// closes stdin — because each is the right trigger for a different way of +// being shut down: a broker ordering its teardown, and a parent that simply +// severed the pipe. +// +// One sweep runs per process, and every caller returns only once it has +// finished. A caller that returned while the sweep was still running would +// report engines stopped before they were: the broker would close stdin, +// engine-manager would exit, and the engines it launched would be orphaned. +// Repeating the sweep would re-pay doStop's readiness probe for every adopted +// engine whose stop it declines, each time. +// +// A later caller does retry the engines whose stop failed through a process +// handle or a stop command, such as an `lms server stop` that errored, because +// another attempt can succeed where a declined stop cannot. func (e *Executor) StopAll() { + swept := false + e.stopAllOnce.Do(func() { + swept = true + e.stopAllNow() + }) + if !swept { + e.retryFailedStops() + } +} + +func (e *Executor) stopAllNow() { e.shuttingDown.Store(true) e.mu.Lock() names := make([]string, 0, len(e.engines)) @@ -759,24 +788,47 @@ func (e *Executor) StopAll() { names = append(names, n) } e.mu.Unlock() - // Stop engines concurrently. Each doStop can wait up to the manifest grace - // (default 5s) for its engine to exit, so a sequential sweep would take - // sum(grace) across engines and let a slow first engine delay — or, if the - // parent severs us mid-sweep, entirely prevent — every later engine from - // even being signaled. Running them in parallel bounds the whole sweep to - // ~one grace and guarantees every engine gets its stop signal immediately. - // Each engine still serializes on its own opMu, so this is safe. + e.failedStops = e.stopEngines(names) +} + +// retryFailedStops tries again to stop the engines whose stop failed in the +// sweep. Callers take turns, so two retries never stop one engine at once. +func (e *Executor) retryFailedStops() { + e.retryMu.Lock() + defer e.retryMu.Unlock() + if len(e.failedStops) == 0 { + return + } + e.failedStops = e.stopEngines(e.failedStops) +} + +// stopEngines stops the named engines and returns those whose stop failed +// through a process handle or a stop command, the failures a later attempt can +// fix. Without either, doStop declines a foreign listener or has already +// escalated against PAIR's own orphan, and a retry would only repeat that. +// +// Engines stop concurrently. Each doStop can wait up to the manifest grace +// (default 5s) for its engine to exit, so a sequential sweep would take +// sum(grace) across engines and let a slow first engine delay — or, if the +// parent severs us mid-sweep, entirely prevent — every later engine from even +// being signaled. Running them in parallel bounds the whole sweep to ~one grace +// and guarantees every engine gets its stop signal immediately. Each engine +// still serializes on its own opMu, so this is safe. +func (e *Executor) stopEngines(names []string) []string { var wg sync.WaitGroup - for _, n := range names { + // Each goroutine writes only its own index, so no lock is needed. + retry := make([]bool, len(names)) + for i, n := range names { st, err := e.state(n) if err != nil { continue } wg.Add(1) - go func(st *engineState, n string) { + go func(i int, st *engineState, n string) { defer wg.Done() st.mu.Lock() cancel := st.startCancel + retryable := st.plat.Runtime.modeOrDefault() == "command" || st.proc != nil st.mu.Unlock() if cancel != nil { cancel() @@ -785,10 +837,18 @@ func (e *Executor) StopAll() { defer st.opMu.Unlock() if err := e.doStop(st, n); err != nil { slog.Warn("engine stop during shutdown failed", "engine", n, "err", err) + retry[i] = retryable } - }(st, n) + }(i, st, n) } wg.Wait() + var failed []string + for i, n := range names { + if retry[i] { + failed = append(failed, n) + } + } + return failed } func (e *Executor) runHealth(ctx context.Context, st *engineState, engine string, port int, gen int64) { diff --git a/services/nvpair-engine-manager/shutdown_test.go b/services/nvpair-engine-manager/shutdown_test.go index 6170b8237..d30c56f9c 100644 --- a/services/nvpair-engine-manager/shutdown_test.go +++ b/services/nvpair-engine-manager/shutdown_test.go @@ -9,13 +9,17 @@ import ( "errors" "net" "net/http" + "net/http/httptest" + "net/url" "os" "path/filepath" "strconv" "strings" + "sync" "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -50,6 +54,179 @@ func TestStopAll_StopsEveryEngine(t *testing.T) { require.ErrorIs(t, ex.Start(ctx, names[0]), context.Canceled, "start after StopAll error") } +// adoptedEngineOnLivePort seeds an engine that this process did not launch and +// cannot stop: running, adopted, no owned proc, and a real listener on its port +// so doStop's readiness probe can never confirm it went away. That is the one +// shape whose stop is refused, leaving st.running true — so a repeated sweep +// does the whole thing again instead of hitting the !st.running early return. +// Its bin path is outside the managed install dir, so the port-reclaim branch +// declines it rather than terminating a stranger's process. +func adoptedEngineOnLivePort(t *testing.T, ex *Executor) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err, "listen") + t.Cleanup(func() { _ = ln.Close() }) + + st, err := ex.state("fake") + require.NoError(t, err) + st.mu.Lock() + st.running = true + st.adopted = true + st.proc = nil + st.port = ln.Addr().(*net.TCPAddr).Port + st.binPath = filepath.Join(t.TempDir(), "someone-elses-engine") + st.mu.Unlock() +} + +// TestStopAll_RunsOneSweepPerProcess is the quit-latency guard. An ordinary quit +// asks for StopAll twice — the broker's engine:prepare-shutdown during its +// teardown, and the stdin-EOF path in Run once the broker closes stdin. An +// engine whose stop can only be declined stays running, so each repeat sweep +// would re-pay doStop's readiness probe for it. +func TestStopAll_RunsOneSweepPerProcess(t *testing.T) { + ex := newTestExecutor(t, testEngineManifest(fakeEngineBin)) + adoptedEngineOnLivePort(t, ex) + + first := time.Now() + ex.StopAll() + sweep := time.Since(first) + require.GreaterOrEqual(t, sweep, 500*time.Millisecond, + "the first sweep was too fast to have probed, so this no longer covers the repeat") + + // The engine is still there and still un-stoppable, so an unguarded sweep + // would spend the probe again. + repeat := time.Now() + ex.StopAll() + ex.StopAll() + assert.LessOrEqual(t, time.Since(repeat), sweep/2, "two further StopAll calls must not repeat the sweep") +} + +// TestStopAll_WaitsForTheSweepInFlight is the other half, and the reason this +// is a joining guard rather than a plain skip. A caller that returned while the +// first sweep was still running would report engines stopped before they were: +// the broker would close stdin, engine-manager would exit, and the engines it +// launched would be left running with nothing owning them. +func TestStopAll_WaitsForTheSweepInFlight(t *testing.T) { + ex := newTestExecutor(t, testEngineManifest(fakeEngineBin)) + adoptedEngineOnLivePort(t, ex) + st, err := ex.state("fake") + require.NoError(t, err) + + // Holding the engine's opMu parks the sweep before doStop, so it cannot + // finish until the test lets it, however the goroutines are scheduled. + st.opMu.Lock() + releaseSweep := sync.OnceFunc(st.opMu.Unlock) + t.Cleanup(releaseSweep) + sweepDone := make(chan struct{}) + go func() { + defer close(sweepDone) + ex.StopAll() + }() + + // Join behind the first caller only once it owns the sweep: only the sweep + // sets shuttingDown, and it does so before stopping anything. + require.Eventually(t, ex.shuttingDown.Load, 5*time.Second, 5*time.Millisecond, + "the first StopAll caller never started the sweep") + secondDone := make(chan struct{}) + go func() { + defer close(secondDone) + ex.StopAll() + }() + + require.Never(t, func() bool { return isClosed(secondDone) }, 200*time.Millisecond, 10*time.Millisecond, + "the second caller returned while the first sweep was still running") + + releaseSweep() + require.Eventually(t, func() bool { return isClosed(sweepDone) }, 10*time.Second, 10*time.Millisecond, + "the first StopAll caller never returned once the sweep was released") + require.Eventually(t, func() bool { return isClosed(secondDone) }, 10*time.Second, 10*time.Millisecond, + "the second StopAll caller never returned once the sweep was released") +} + +// isClosed reports whether done has been closed, without blocking. +func isClosed(done <-chan struct{}) bool { + select { + case <-done: + return true + default: + return false + } +} + +// commandDaemon stands in for a command-mode engine's daemon: it serves the +// readiness endpoint once startMarker exists and closes its listener once +// stopMarker does, as a daemon told to stop by its control CLI would. It +// returns the daemon's port. +func commandDaemon(t *testing.T, startMarker, stopMarker string) int { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if fileExists(startMarker) && !fileExists(stopMarker) { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusServiceUnavailable) + })) + t.Cleanup(srv.Close) + stopWatch := make(chan struct{}) + t.Cleanup(func() { close(stopWatch) }) + go func() { + for { + select { + case <-stopWatch: + return + case <-time.After(50 * time.Millisecond): + if fileExists(stopMarker) { + srv.Close() + return + } + } + } + }() + u, err := url.Parse(srv.URL) + require.NoError(t, err, "parse daemon URL") + port, err := strconv.Atoi(u.Port()) + require.NoError(t, err, "parse daemon port") + return port +} + +// TestStopAll_RetriesAFailedStopCommand checks that a later StopAll caller +// retries an engine whose stop command failed during the sweep, so a stop +// command that errors once does not leave the engine running after quit. +func TestStopAll_RetriesAFailedStopCommand(t *testing.T) { + dir := t.TempDir() + startMarker := filepath.Join(dir, "started") + stopAttempt := filepath.Join(dir, "stop-attempted") + stopMarker := filepath.Join(dir, "stopped") + port := commandDaemon(t, startMarker, stopMarker) + ex := newTestExecutor(t, &Manifest{ + Engine: "daemon", DisplayName: "Daemon Engine", ManifestVersion: 1, + Platforms: map[string]Platform{ + hostKey(): { + Detect: []string{fakeEngineBin}, + Runtime: Runtime{ + Mode: "command", + Port: port, + Start: [][]string{{fakeEngineBin, "touch", startMarker}}, + Stop: &StopSpec{Cmd: []string{fakeEngineBin, "failoncethentouch", stopAttempt, stopMarker}}, + Ready: &Probe{HTTP: "http://127.0.0.1:{port}/", Status: http.StatusOK, TimeoutS: 10}, + }, + }, + }, + }) + require.NoError(t, ex.Start(context.Background(), "daemon")) + ex.StopAll() + st, err := ex.Status("daemon") + require.NoError(t, err) + require.True(t, st.Running, "the failed stop must leave the engine running after the sweep") + + ex.StopAll() + + assert.FileExists(t, stopMarker, "the later StopAll must run the stop command again") + st, err = ex.Status("daemon") + require.NoError(t, err) + assert.False(t, st.Running, "the retried stop must leave the engine stopped") +} + func TestStopAll_CancelsCommandStartWithoutError(t *testing.T) { pidFile := filepath.Join(t.TempDir(), "fake.pid") t.Setenv("FAKE_PID_FILE", pidFile) diff --git a/services/nvpair-engine-manager/spec.md b/services/nvpair-engine-manager/spec.md index c18f172f3..ac2ec2c5f 100644 --- a/services/nvpair-engine-manager/spec.md +++ b/services/nvpair-engine-manager/spec.md @@ -157,6 +157,7 @@ Requests (caller → service): | `engine:remote-delete-model` | `{ node, engine, model }` | the remote action result | | `engine:remote-start` | `{ node, engine, port? }` | `EngineStatus` from the remote node (manifest `runtime.bind`; no per-call bind on the remote path) | | `engine:remote-stop` | `{ node, engine }` | `EngineStatus` from the remote node | +| `engine:prepare-shutdown` | — | `null` once the shutdown sweep, and any retry of a failed stop, has finished; per-engine stop failures are logged, not returned | | `shutdown` | — | `null` | | `log/set-level` | `{ level }` | `{ level }` | @@ -248,7 +249,16 @@ The operator starts it: `engine:start {engine:"ollama"}` resolves the manifest r The operator stops it: `engine:stop {engine:"ollama"}` signals a process the service owns. For an **adopted** engine (no owned process), it resolves the PID bound to the port and terminates it only when that process is running the binary we manage — reclaiming an orphan a prior run left on our own managed port; a genuinely foreign listener (a different image on a different port) is declined with an error naming its PID and image. A user-initiated `stop` records the OFF intent regardless (even when the RPC returns an error), so the health loop and restore-on-restart don't flip the engine back on — clients must not treat a stop error as proof the OFF choice was discarded. The cluster `ec` stop endpoint shares this semantics and may return HTTP 500 while OFF is persisted. -The operator pulls a model: `engine:action {engine:"ollama", action:"pull_model", params:{name:"llama3.2"}}` issues the manifest-declared `POST 127.0.0.1:{port}/api/pull`. Because the action is `pull_model`, the request is routed through the streaming pull path (not the buffered `engine:action` reader): each `/api/pull` status line is emitted as an `engine:pull-progress` notification — so a local pull shows live download progress just like a remote pull's `engine:remote-progress` — and the request settles with the pull's terminal result line. Frames are coalesced (only a change in `stage` or `percent` is emitted) so a chatty engine that streams many byte-progress lines per layer doesn't flood subscribers. Streaming Ollama and llama.cpp pulls use a 30-minute inactivity watchdog that is refreshed only when a layer/file completed-byte count advances, allowing active downloads to exceed 30 minutes without letting duplicate progress or heartbeat frames keep a stalled pull alive. The engine's terminal `{"status":"success"}` surfaces as a `stage:"success"` frame; a **failed** pull emits a terminal `stage:"error", percent:-1, message:` frame in addition to the JSON-RPC error, so a UI whose synchronous call already timed out on a long download still converges off "pulling". A CLI-driven pull (LM Studio's `lms get`) has no line-level progress, so it emits one `stage:"pulling"` marker and retains the fixed 30-minute action timeout before returning the command's result. On `shutdown` (or stdin EOF) the service stops every running engine first, so none are orphaned. +The operator pulls a model: `engine:action {engine:"ollama", action:"pull_model", params:{name:"llama3.2"}}` issues the manifest-declared `POST 127.0.0.1:{port}/api/pull`. +Because the action is `pull_model`, the request is routed through the streaming pull path (not the buffered `engine:action` reader): each `/api/pull` status line is emitted as an `engine:pull-progress` notification — so a local pull shows live download progress just like a remote pull's `engine:remote-progress` — and the request settles with the pull's terminal result line. +Frames are coalesced (only a change in `stage` or `percent` is emitted) so a chatty engine that streams many byte-progress lines per layer doesn't flood subscribers. +Streaming Ollama and llama.cpp pulls use a 30-minute inactivity watchdog that is refreshed only when a layer/file completed-byte count advances, allowing active downloads to exceed 30 minutes without letting duplicate progress or heartbeat frames keep a stalled pull alive. +The engine's terminal `{"status":"success"}` surfaces as a `stage:"success"` frame; a **failed** pull emits a terminal `stage:"error", percent:-1, message:` frame in addition to the JSON-RPC error, so a UI whose synchronous call already timed out on a long download still converges off "pulling". +A CLI-driven pull (LM Studio's `lms get`) has no line-level progress, so it emits one `stage:"pulling"` marker and retains the fixed 30-minute action timeout before returning the command's result. +On `shutdown` (or stdin EOF) the service stops every running engine first, so none are orphaned. +`engine:prepare-shutdown` requests that same sweep while the service is still fully up, which is how the parent drains engines before it starts tearing the worker tree down. +The sweep runs **once per process** no matter how many of these paths ask for it, and a caller that arrives while it is already running waits for it to finish — responding earlier would report engines stopped while they were still running, and the parent would exit and orphan them. +A later caller does not repeat the sweep; it retries only an engine whose stop failed through a process handle or a stop command, not an adopted engine whose stop was declined. For the `llamacpp-models-sse` adapter, the monitored model must match `params.model`. Subscribe before starting, and bound the `POST /models` diff --git a/services/nvpair-engine-manager/testdata/fakeengine/main.go b/services/nvpair-engine-manager/testdata/fakeengine/main.go index 9f12ba821..e58973ce8 100644 --- a/services/nvpair-engine-manager/testdata/fakeengine/main.go +++ b/services/nvpair-engine-manager/testdata/fakeengine/main.go @@ -130,6 +130,17 @@ func main() { } fmt.Fprintf(os.Stderr, "Error: Failed to resolve artifact %q: The artifact does not exist or you do not have permission to read it\n", strings.ToLower(arg)) os.Exit(1) + case "failoncethentouch": // fail the first call, then write a marker — for stop-retry tests + // args: failoncethentouch + if len(os.Args) < 4 { + os.Exit(2) + } + if _, err := os.Stat(os.Args[2]); err != nil { + _ = os.WriteFile(os.Args[2], []byte("1"), 0o644) + os.Exit(1) + } + _ = os.WriteFile(os.Args[3], []byte("ok"), 0o644) + return case "noserve": // run but never bind — used to test readiness timeout time.Sleep(time.Hour) return diff --git a/services/nvpair-ui-broker/README.md b/services/nvpair-ui-broker/README.md index 37f8fc0d7..e80ff55d9 100644 --- a/services/nvpair-ui-broker/README.md +++ b/services/nvpair-ui-broker/README.md @@ -554,6 +554,8 @@ If no cluster-manager is being supervised (or it has exited), the broker replies Acknowledges, then terminates the broker. The proxy is stopped first, then engine-manager is asked to stop its engines (`engine:prepare-shutdown`), then every running worker subprocess is torn down by closing its stdin and waiting for it to exit — no grace timer, no force-kill (each worker bounds its own shutdown). +The broker's teardown sends `engine:prepare-shutdown`, and closing engine-manager's stdin afterwards reaches the same engine sweep. Engine-manager sweeps its engines once; later callers wait for that sweep rather than starting another, and retry only an engine whose stop failed through a process handle or a stop command. + ```json {"jsonrpc":"2.0","id":4,"method":"shutdown"} ``` diff --git a/services/tests/broker_supervision_test.go b/services/tests/broker_supervision_test.go index 5aad1cc88..8a7d56ad2 100644 --- a/services/tests/broker_supervision_test.go +++ b/services/tests/broker_supervision_test.go @@ -81,6 +81,7 @@ func startBrokerWithDirsAndEnv(t *testing.T, configDir, clusterDir string, extra "XDG_CONFIG_HOME="+configDir, "APPDATA="+configDir, "LOCALAPPDATA="+configDir, + "USERPROFILE="+configDir, ) cmd.Env = append(cmd.Env, extraEnv...) diff --git a/services/tests/shutdown_latency_test.go b/services/tests/shutdown_latency_test.go new file mode 100644 index 000000000..a50725b27 --- /dev/null +++ b/services/tests/shutdown_latency_test.go @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package tests + +import ( + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "regexp" + "runtime" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "nvpair-shared/jsonrpc" +) + +// foreignStopDeclinedRe matches engine-manager's warning for the foreign +// engine, whose stop it can only decline. One sweep logs it exactly once, so +// counting the line counts the sweeps. It names the engine because a real +// engine on the host can be adopted and declined in the same sweep. +var foreignStopDeclinedRe = regexp.MustCompile(`engine stop during shutdown failed.*engine=foreign`) + +// startForeignEngine serves an engine readiness endpoint on a free port without +// engine-manager having launched it. That is what an externally-managed engine +// looks like from the inside: the readiness probe succeeds, so the engine is +// adopted as running, but nothing engine-manager owns can stop it, so its stop +// is declined and it stays marked running. +func startForeignEngine(t *testing.T) int { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err, "listen for foreign engine") + srv := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"models":[]}`)) + })} + go func() { _ = srv.Serve(ln) }() + t.Cleanup(func() { _ = srv.Close() }) + + port := ln.Addr().(*net.TCPAddr).Port + t.Logf("foreign engine listening on 127.0.0.1:%d", port) + return port +} + +// writeForeignEngineManifest declares a process-mode engine pinned to port with +// an HTTP readiness probe. A process-mode engine is reconciled from its +// readiness endpoint alone, so it is adopted without any installed binary — +// which is the point: nothing here is ours to stop. +// +// The runtime goes in a host platform block, not just at the top level: an +// engine with no block for this os/arch has no per-engine state, so it would +// never be reconciled at all. +func writeForeignEngineManifest(t *testing.T, configDir string, port int) { + t.Helper() + runtimeSpec := map[string]any{ + "mode": "process", + "bin": filepath.Join(t.TempDir(), "not-ours"), + "port": port, + "ready": map[string]any{ + "http": fmt.Sprintf("http://127.0.0.1:%d/", port), + "status": http.StatusOK, + "timeout_s": 10, + }, + "stop": map[string]any{"signal": "term", "grace_s": 3}, + } + manifest := map[string]any{ + "engine": "foreign", + "display_name": "Foreign Engine", + "manifest_version": 1, + "platforms": map[string]any{ + runtime.GOOS + "/" + runtime.GOARCH: map[string]any{"runtime": runtimeSpec}, + }, + } + data, err := json.MarshalIndent(manifest, "", " ") + require.NoError(t, err, "encode foreign engine manifest") + // The per-user manifest dir differs by platform, and the broker points every + // worker at one disposable root, so write both candidates rather than + // branching on GOOS: the Windows %LocalAppData% / Linux $XDG_CONFIG_HOME + // layout, and the macOS Application Support layout under $HOME. + for _, dir := range []string{ + filepath.Join(configDir, "Nvidia Corporation", "Personal AI Router", "engines"), + filepath.Join(configDir, "Library", "Application Support", "Nvidia Corporation", "Personal AI Router", "engines"), + } { + require.NoError(t, os.MkdirAll(dir, 0o755), "create engine manifest dir") + require.NoError(t, os.WriteFile(filepath.Join(dir, "foreign.json"), data, 0o644), "write foreign engine manifest") + } +} + +// TestBrokerQuit_SweepsEnginesOnce is the end-to-end guard for quit latency, driving +// the real broker and engine-manager through the sequence the desktop and the +// TUI use on quit: the broker's shutdown request, then close stdin. +// +// Two separate paths ask engine-manager to stop its engines on that sequence — +// the broker's engine:prepare-shutdown during its teardown, and +// engine-manager's stdin-EOF path once the broker closes its stdin. Each is the +// right trigger for a different way of being shut down, so both stay; the sweep +// behind them must run once. +// +// This counts sweeps rather than milliseconds so it cannot flake on a loaded +// CI machine. An engine whose stop can only be declined stays marked running, +// so every repeat sweep would re-pay its readiness probe on the quit path. +func TestBrokerQuit_SweepsEnginesOnce(t *testing.T) { + configDir := t.TempDir() + port := startForeignEngine(t) + writeForeignEngineManifest(t, configDir, port) + + stdin, msgs, stderrLines, cleanup := startBrokerWithConfigDir(t, configDir, + "--engine-manager-path", engineMgrBin) + defer cleanup() + + waitForMethod(t, msgs, "app:ready", 15*time.Second) + + // engine:get-installed reconciles every engine's presence, which is what + // adopts the foreign listener as running. Without it there is nothing for a + // sweep to decline and the test would pass while covering nothing. + sendReq(t, stdin, 98, "engine:get-installed") + require.Nil(t, waitForResponseID(t, msgs, 98, 15*time.Second).Error, "engine:get-installed errored") + + sendReq(t, stdin, 99, "shutdown") + sweeps := quitAndCountStderr(t, msgs, stderrLines, foreignStopDeclinedRe, stdin, 99, 30*time.Second) + + require.NotZero(t, sweeps, "engine-manager never declined a stop; the foreign engine was not adopted, so this exercised nothing") + assert.Equal(t, 1, sweeps, "each repeat sweep re-pays the stop probe and is what made quitting slow") +} + +// quitAndCountStderr completes a client's quit sequence and counts stderr lines +// matching re until the broker exits. It closes stdin once the shutdown reply +// (replyID) has arrived. Both streams are drained throughout, so the broker +// never blocks writing to either. +func quitAndCountStderr(t *testing.T, msgs <-chan jsonrpc.Message, lines <-chan string, re *regexp.Regexp, stdin io.Closer, replyID int, timeout time.Duration) int { + t.Helper() + replied := false + count := 0 + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + select { + case msg, ok := <-msgs: + if !ok { + msgs = nil + continue + } + if replied || msg.Method != "" || !idEquals(t, msg.ID, replyID) { + continue + } + require.Nil(t, msg.Error, "shutdown errored") + replied = true + require.NoError(t, stdin.Close(), "close broker stdin") + case line, ok := <-lines: + if !ok { + require.True(t, replied, "broker exited before replying to shutdown; counted %d sweep(s)", count) + return count + } + if re.MatchString(line) { + count++ + t.Logf("declined stop #%d: %s", count, line) + } + case <-timer.C: + require.FailNowf(t, "broker did not exit", "timed out after %s; replied=%t, counted %d sweep(s)", timeout, replied, count) + } + } +}