diff --git a/managed-inference/images/llama-cpp/Dockerfile b/managed-inference/images/llama-cpp/Dockerfile index 0299f6168ae..9084fe7ceb0 100644 --- a/managed-inference/images/llama-cpp/Dockerfile +++ b/managed-inference/images/llama-cpp/Dockerfile @@ -32,6 +32,7 @@ RUN apt-get update \ curl=8.5.0-2ubuntu10.11 \ g++-14=14.2.0-4ubuntu2~24.04.1 \ gcc-14=14.2.0-4ubuntu2~24.04.1 \ + golang-go=2:1.22~2build1 \ libcurl4-openssl-dev=8.5.0-2ubuntu10.11 \ libssl-dev=3.0.13-0ubuntu3.12 \ && rm -rf /var/lib/apt/lists/* @@ -80,6 +81,23 @@ RUN cmake -S . -B build \ && find /opt/llama.cpp/licenses -type d -exec chmod 0555 '{}' + \ && find /opt/llama.cpp/licenses -type f -exec chmod 0444 '{}' + +WORKDIR /src/nemoclaw-request-guard + +COPY request-guard/go.mod request-guard/*.go ./ + +RUN go test ./... \ + && CGO_ENABLED=0 go build \ + -trimpath \ + -ldflags='-s -w -buildid=' \ + -o /opt/llama.cpp/bin/nemoclaw-llama-cpp-request-guard . \ + && mkdir -p /opt/llama.cpp/licenses/go \ + && set -- /usr/share/doc/golang-[0-9]*-go/copyright \ + && test "$#" -eq 1 \ + && cp "$1" /opt/llama.cpp/licenses/go/copyright \ + && chmod 0555 /opt/llama.cpp/bin/nemoclaw-llama-cpp-request-guard \ + && chmod 0555 /opt/llama.cpp/licenses/go \ + && chmod 0444 /opt/llama.cpp/licenses/go/copyright + FROM ${CUDA_RUNTIME_IMAGE} AS runtime ARG CUDA_DEV_IMAGE @@ -122,6 +140,7 @@ RUN apt-get update \ /usr/bin/sh COPY --from=build --chmod=0555 /opt/llama.cpp/bin/llama-server /usr/local/bin/llama-server +COPY --from=build --chmod=0555 /opt/llama.cpp/bin/nemoclaw-llama-cpp-request-guard /usr/local/bin/nemoclaw-llama-cpp-request-guard COPY --from=build --chmod=0555 /opt/llama.cpp/lib/ /opt/llama.cpp/lib/ COPY --from=build /opt/llama.cpp/licenses/ /usr/local/share/licenses/ diff --git a/managed-inference/images/llama-cpp/image.yaml b/managed-inference/images/llama-cpp/image.yaml index 447855e1c5d..5c177388983 100644 --- a/managed-inference/images/llama-cpp/image.yaml +++ b/managed-inference/images/llama-cpp/image.yaml @@ -6,6 +6,8 @@ kind: ServerImageBuild metadata: id: llama-cpp-server.v1 + annotations: + nemoclaw.nvidia.com/request-guard-state: dormant spec: repository: ghcr.io/nvidia/nemoclaw/llama-cpp-server @@ -111,6 +113,7 @@ spec: curl: 8.5.0-2ubuntu10.11 g++-14: 14.2.0-4ubuntu2~24.04.1 gcc-14: 14.2.0-4ubuntu2~24.04.1 + golang-go: 2:1.22~2build1 libcurl4-openssl-dev: 8.5.0-2ubuntu10.11 libssl-dev: 3.0.13-0ubuntu3.12 cmake: @@ -138,6 +141,8 @@ spec: requiredPaths: - /opt/llama.cpp/lib/libggml-cuda.so - /usr/local/bin/llama-server + - /usr/local/bin/nemoclaw-llama-cpp-request-guard + - /usr/local/share/licenses/go/copyright - /usr/local/share/licenses/llama.cpp/AUTHORS - /usr/local/share/licenses/llama.cpp/LICENSE forbiddenPaths: diff --git a/managed-inference/images/llama-cpp/request-guard/go.mod b/managed-inference/images/llama-cpp/request-guard/go.mod new file mode 100644 index 00000000000..607f72dd55b --- /dev/null +++ b/managed-inference/images/llama-cpp/request-guard/go.mod @@ -0,0 +1,3 @@ +module github.com/NVIDIA/NemoClaw/managed-inference/images/llama-cpp/request-guard + +go 1.22 diff --git a/managed-inference/images/llama-cpp/request-guard/main.go b/managed-inference/images/llama-cpp/request-guard/main.go new file mode 100644 index 00000000000..8c61ed81303 --- /dev/null +++ b/managed-inference/images/llama-cpp/request-guard/main.go @@ -0,0 +1,663 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "flag" + "fmt" + "io" + "net" + "net/http" + "net/http/httputil" + "net/url" + "os" + "os/exec" + "os/signal" + "strconv" + "strings" + "syscall" + "time" +) + +const ( + llamaServerPath = "/usr/local/bin/llama-server" + llamaServerAPIKeyPath = "/run/secrets/llama-cpp-api-key" + maximumBodyBytes = 64 * 1024 * 1024 + maximumHeaderBytes = 1024 * 1024 + maximumOutputTokens = 1024 * 1024 + maximumTimeoutSeconds = 24 * 60 * 60 +) + +type guardConfig struct { + listenHost string + listenPort int + upstreamHost string + upstreamPort int + maxRequestBodyBytes int64 + maxRequestHeaderBytes int + maxOutputTokens int64 + requestTimeout time.Duration + shutdownTimeout time.Duration +} + +type guardError struct { + status int + code string + message string +} + +func (e *guardError) Error() string { return e.message } + +func positiveBounded(value int64, maximum int64, name string) error { + if value < 1 || value > maximum { + return fmt.Errorf("%s must be between 1 and %d", name, maximum) + } + return nil +} + +func requireExactCommandOption(command []string, option, expected string) error { + count := 0 + for index, value := range command { + if value == option { + count++ + if index+1 >= len(command) || command[index+1] != expected { + return fmt.Errorf("llama-server %s must be %s", option, expected) + } + } + if strings.HasPrefix(value, option+"=") { + return fmt.Errorf("llama-server %s must use a separate exact value", option) + } + } + if count != 1 { + return fmt.Errorf("llama-server command must declare %s exactly once", option) + } + return nil +} + +func requireExactCommandMarker(command []string, option string) error { + count := 0 + for _, value := range command { + if value == option { + count++ + } + if strings.HasPrefix(value, option+"=") { + return fmt.Errorf("llama-server %s does not accept a value", option) + } + } + if count != 1 { + return fmt.Errorf("llama-server command must declare %s exactly once", option) + } + return nil +} + +func validateSupportedCommandOptions(command []string) error { + allowed := map[string]bool{ + "--alias": true, + "--api-key-file": true, + "--batch-size": true, + "--cache-type-k": true, + "--cache-type-v": true, + "--ctx-size": true, + "--flash-attn": true, + "--gpu-layers": true, + "--host": true, + "--metrics": false, + "--model": true, + "--no-agent": false, + "--no-mmproj": false, + "--no-slots": false, + "--no-ui": false, + "--n-predict": true, + "--parallel": true, + "--port": true, + "--sleep-idle-seconds": true, + "--timeout": true, + "--ubatch-size": true, + } + seen := make(map[string]bool, len(allowed)) + for index := 0; index < len(command); index++ { + option := command[index] + takesValue, supported := allowed[option] + if !supported { + return fmt.Errorf("llama-server option %s is not supported by the request guard", option) + } + if seen[option] { + return fmt.Errorf("llama-server command must declare %s at most once", option) + } + seen[option] = true + if !takesValue { + continue + } + index++ + if index >= len(command) || strings.HasPrefix(command[index], "--") { + return fmt.Errorf("llama-server %s requires one value", option) + } + } + if !seen["--model"] { + return errors.New("llama-server command must declare --model exactly once") + } + return nil +} + +func validateLlamaServerCommand(command []string, config guardConfig) error { + if len(command) == 0 || command[0] != llamaServerPath { + return fmt.Errorf("request guard command must start with %s", llamaServerPath) + } + if err := validateSupportedCommandOptions(command[1:]); err != nil { + return err + } + for _, required := range []struct { + option string + value string + }{ + {option: "--host", value: config.upstreamHost}, + {option: "--port", value: strconv.Itoa(config.upstreamPort)}, + {option: "--api-key-file", value: llamaServerAPIKeyPath}, + {option: "--n-predict", value: strconv.FormatInt(config.maxOutputTokens, 10)}, + } { + if err := requireExactCommandOption(command[1:], required.option, required.value); err != nil { + return err + } + } + for _, marker := range []string{"--no-agent", "--no-mmproj", "--no-slots", "--no-ui"} { + if err := requireExactCommandMarker(command[1:], marker); err != nil { + return err + } + } + return nil +} + +func parseConfig(args []string) (guardConfig, []string, error) { + var config guardConfig + var timeoutSeconds int64 + var shutdownTimeoutSeconds int64 + separator := -1 + for index, arg := range args { + if arg == "--" { + separator = index + break + } + } + if separator < 0 { + return config, nil, errors.New("request guard requires '--' before the llama-server command") + } + + flags := flag.NewFlagSet("nemoclaw-llama-cpp-request-guard", flag.ContinueOnError) + flags.SetOutput(io.Discard) + flags.StringVar(&config.listenHost, "listen-host", "", "guard listen host") + flags.IntVar(&config.listenPort, "listen-port", 0, "guard listen port") + flags.StringVar(&config.upstreamHost, "upstream-host", "", "llama-server host") + flags.IntVar(&config.upstreamPort, "upstream-port", 0, "llama-server port") + flags.Int64Var( + &config.maxRequestBodyBytes, + "max-request-body-bytes", + 0, + "maximum request body bytes", + ) + flags.IntVar( + &config.maxRequestHeaderBytes, + "max-request-header-bytes", + 0, + "maximum request header bytes", + ) + flags.Int64Var( + &config.maxOutputTokens, + "max-output-tokens", + 0, + "maximum generated tokens", + ) + flags.Int64Var( + &timeoutSeconds, + "request-timeout-seconds", + 0, + "request timeout seconds", + ) + flags.Int64Var( + &shutdownTimeoutSeconds, + "shutdown-timeout-seconds", + 0, + "graceful shutdown timeout seconds", + ) + if err := flags.Parse(args[:separator]); err != nil { + return config, nil, fmt.Errorf("invalid request guard arguments: %w", err) + } + if flags.NArg() != 0 { + return config, nil, errors.New("request guard received an argument before '--'") + } + if config.listenHost != "0.0.0.0" { + return config, nil, errors.New("request guard listen host must be 0.0.0.0") + } + if config.upstreamHost != "127.0.0.1" { + return config, nil, errors.New("request guard upstream host must be 127.0.0.1") + } + if err := positiveBounded(int64(config.listenPort), 65535, "request guard listen port"); err != nil { + return config, nil, err + } + if err := positiveBounded(int64(config.upstreamPort), 65535, "request guard upstream port"); err != nil { + return config, nil, err + } + if config.listenPort == config.upstreamPort { + return config, nil, errors.New("request guard listen and upstream ports must differ") + } + if err := positiveBounded( + config.maxRequestBodyBytes, + maximumBodyBytes, + "maximum request body bytes", + ); err != nil { + return config, nil, err + } + if err := positiveBounded( + int64(config.maxRequestHeaderBytes), + maximumHeaderBytes, + "maximum request header bytes", + ); err != nil { + return config, nil, err + } + if err := positiveBounded( + config.maxOutputTokens, + maximumOutputTokens, + "maximum output tokens", + ); err != nil { + return config, nil, err + } + if err := positiveBounded(timeoutSeconds, maximumTimeoutSeconds, "request timeout seconds"); err != nil { + return config, nil, err + } + if err := positiveBounded( + shutdownTimeoutSeconds, + maximumTimeoutSeconds, + "shutdown timeout seconds", + ); err != nil { + return config, nil, err + } + config.requestTimeout = time.Duration(timeoutSeconds) * time.Second + config.shutdownTimeout = time.Duration(shutdownTimeoutSeconds) * time.Second + + command := args[separator+1:] + if err := validateLlamaServerCommand(command, config); err != nil { + return config, nil, err + } + return config, command, nil +} + +func writeGuardError(writer http.ResponseWriter, failure *guardError) { + writer.Header().Set("Cache-Control", "no-store") + writer.Header().Set("Content-Type", "application/json") + writer.Header().Set("X-Content-Type-Options", "nosniff") + writer.WriteHeader(failure.status) + _ = json.NewEncoder(writer).Encode(map[string]any{ + "error": map[string]string{ + "code": failure.code, + "message": failure.message, + "type": "invalid_request_error", + }, + }) +} + +func boundedBody(request *http.Request, maximum int64) ([]byte, *guardError) { + if request.ContentLength > maximum { + return nil, &guardError{ + status: http.StatusRequestEntityTooLarge, + code: "request_body_too_large", + message: "Request body exceeds the declared limit.", + } + } + encoding := strings.TrimSpace(strings.ToLower(request.Header.Get("Content-Encoding"))) + if encoding != "" && encoding != "identity" { + return nil, &guardError{ + status: http.StatusUnsupportedMediaType, + code: "content_encoding_unsupported", + message: "Compressed request bodies are not supported.", + } + } + body, err := io.ReadAll(io.LimitReader(request.Body, maximum+1)) + if err != nil { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "request_body_unreadable", + message: "Request body could not be read.", + } + } + if int64(len(body)) > maximum { + return nil, &guardError{ + status: http.StatusRequestEntityTooLarge, + code: "request_body_too_large", + message: "Request body exceeds the declared limit.", + } + } + return body, nil +} + +func decodeTopLevelObject(body []byte) (map[string]json.RawMessage, *guardError) { + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + opening, err := decoder.Token() + if err != nil || opening != json.Delim('{') { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "invalid_json", + message: "Chat Completions request must be one JSON object.", + } + } + fields := make(map[string]json.RawMessage) + for decoder.More() { + keyToken, keyErr := decoder.Token() + key, ok := keyToken.(string) + if keyErr != nil || !ok { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "invalid_json", + message: "Chat Completions request must be one JSON object.", + } + } + if _, exists := fields[key]; exists { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "duplicate_json_field", + message: "Chat Completions request contains a duplicate field.", + } + } + var value json.RawMessage + if err := decoder.Decode(&value); err != nil { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "invalid_json", + message: "Chat Completions request must be one JSON object.", + } + } + fields[key] = value + } + closing, err := decoder.Token() + if err != nil || closing != json.Delim('}') { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "invalid_json", + message: "Chat Completions request must be one JSON object.", + } + } + if _, err := decoder.Token(); !errors.Is(err, io.EOF) { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "invalid_json", + message: "Chat Completions request must contain one JSON value.", + } + } + return fields, nil +} + +func parsePositiveInteger(raw json.RawMessage) (int64, bool) { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return 0, false + } + number, ok := value.(json.Number) + if !ok { + return 0, false + } + parsed, err := strconv.ParseInt(number.String(), 10, 64) + return parsed, err == nil && parsed > 0 +} + +func guardChatBody(body []byte, maximum int64) ([]byte, *guardError) { + fields, failure := decodeTopLevelObject(body) + if failure != nil { + return nil, failure + } + boundedFieldPresent := false + for _, name := range []string{"max_tokens", "max_completion_tokens", "n_predict"} { + raw, present := fields[name] + if !present { + continue + } + boundedFieldPresent = true + value, valid := parsePositiveInteger(raw) + if !valid { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "output_token_limit_invalid", + message: "Output token limit must be a positive integer.", + } + } + if value > maximum { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "output_token_limit_exceeded", + message: "Output token limit exceeds the declared limit.", + } + } + } + if boundedFieldPresent { + return body, nil + } + fields["max_tokens"] = json.RawMessage(strconv.FormatInt(maximum, 10)) + guarded, err := json.Marshal(fields) + if err != nil { + return nil, &guardError{ + status: http.StatusBadRequest, + code: "invalid_json", + message: "Chat Completions request could not be normalized.", + } + } + return guarded, nil +} + +func routeAllowed(request *http.Request) bool { + if request.URL.RawQuery != "" { + return false + } + switch request.URL.Path { + case "/v1/chat/completions": + return request.Method == http.MethodPost + case "/v1/models", "/health", "/props", "/metrics": + return request.Method == http.MethodGet + default: + return false + } +} + +func newGuardHandler(config guardConfig) (http.Handler, error) { + upstream, err := url.Parse( + fmt.Sprintf("http://%s:%d", config.upstreamHost, config.upstreamPort), + ) + if err != nil { + return nil, errors.New("request guard upstream URL is invalid") + } + proxy := httputil.NewSingleHostReverseProxy(upstream) + baseDirector := proxy.Director + proxy.Director = func(request *http.Request) { + baseDirector(request) + request.Host = upstream.Host + request.Header.Del("Forwarded") + request.Header.Del("X-Forwarded-Host") + request.Header.Del("X-Forwarded-Proto") + request.Header["X-Forwarded-For"] = nil + } + proxy.FlushInterval = -1 + proxy.ErrorHandler = func(writer http.ResponseWriter, _ *http.Request, _ error) { + writeGuardError(writer, &guardError{ + status: http.StatusBadGateway, + code: "upstream_unavailable", + message: "The managed inference server is unavailable.", + }) + } + + return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if !routeAllowed(request) { + writeGuardError(writer, &guardError{ + status: http.StatusNotFound, + code: "route_not_available", + message: "The requested server route is not available.", + }) + return + } + body, failure := boundedBody(request, config.maxRequestBodyBytes) + if failure != nil { + writeGuardError(writer, failure) + return + } + if request.URL.Path == "/v1/chat/completions" { + contentType := strings.ToLower(strings.TrimSpace(strings.Split(request.Header.Get("Content-Type"), ";")[0])) + if contentType != "application/json" { + writeGuardError(writer, &guardError{ + status: http.StatusUnsupportedMediaType, + code: "content_type_unsupported", + message: "Chat Completions request must use application/json.", + }) + return + } + body, failure = guardChatBody(body, config.maxOutputTokens) + if failure != nil { + writeGuardError(writer, failure) + return + } + } + request.Body = io.NopCloser(bytes.NewReader(body)) + request.ContentLength = int64(len(body)) + request.TransferEncoding = nil + request.Header.Set("Content-Length", strconv.Itoa(len(body))) + proxy.ServeHTTP(writer, request) + }), nil +} + +func childExitCode(state *os.ProcessState) int { + if state == nil { + return 1 + } + return state.ExitCode() +} + +func waitForChildUntil( + child *exec.Cmd, + childExited <-chan *os.ProcessState, + deadline time.Time, +) int { + remaining := time.Until(deadline) + if remaining > 0 { + timer := time.NewTimer(remaining) + defer timer.Stop() + select { + case state := <-childExited: + return childExitCode(state) + case <-timer.C: + } + } + _ = child.Process.Kill() + return childExitCode(<-childExited) +} + +func stopChildWithin( + child *exec.Cmd, + childExited <-chan *os.ProcessState, + received os.Signal, + timeout time.Duration, +) int { + deadline := time.Now().Add(timeout) + _ = child.Process.Signal(received) + return waitForChildUntil(child, childExited, deadline) +} + +func newHTTPServer(config guardConfig, handler http.Handler) *http.Server { + return &http.Server{ + Handler: handler, + ReadHeaderTimeout: config.requestTimeout, + ReadTimeout: config.requestTimeout, + WriteTimeout: config.requestTimeout, + IdleTimeout: config.requestTimeout, + MaxHeaderBytes: config.maxRequestHeaderBytes, + } +} + +func validateAPIKeyFile(path string) error { + file, err := os.Open(path) + if err != nil { + return errors.New("request guard API-key file is unavailable") + } + defer file.Close() + info, err := file.Stat() + if err != nil || !info.Mode().IsRegular() { + return errors.New("request guard API-key file is not a regular file") + } + var firstByte [1]byte + if count, err := file.Read(firstByte[:]); count != 1 || err != nil { + return errors.New("request guard API-key file is empty or unreadable") + } + return nil +} + +func run(config guardConfig, command []string) int { + if err := validateAPIKeyFile(llamaServerAPIKeyPath); err != nil { + fmt.Fprintln(os.Stderr, err.Error()) + return 1 + } + listener, err := net.Listen("tcp", net.JoinHostPort(config.listenHost, strconv.Itoa(config.listenPort))) + if err != nil { + fmt.Fprintln(os.Stderr, "request guard could not bind its declared listener") + return 1 + } + defer listener.Close() + + handler, err := newGuardHandler(config) + if err != nil { + fmt.Fprintln(os.Stderr, err.Error()) + return 1 + } + child := exec.Command(command[0], command[1:]...) + child.Stdin = os.Stdin + child.Stdout = os.Stdout + child.Stderr = os.Stderr + if err := child.Start(); err != nil { + fmt.Fprintln(os.Stderr, "request guard could not start llama-server") + return 1 + } + + server := newHTTPServer(config, handler) + serverErrors := make(chan error, 1) + go func() { + serverErrors <- server.Serve(listener) + }() + childExited := make(chan *os.ProcessState, 1) + go func() { + _ = child.Wait() + childExited <- child.ProcessState + }() + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(signals) + + select { + case state := <-childExited: + _ = server.Close() + return childExitCode(state) + case serverErr := <-serverErrors: + if !errors.Is(serverErr, http.ErrServerClosed) { + fmt.Fprintln(os.Stderr, "request guard listener stopped") + } + return stopChildWithin(child, childExited, syscall.SIGTERM, config.shutdownTimeout) + case received := <-signals: + deadline := time.Now().Add(config.shutdownTimeout) + _ = child.Process.Signal(received) + shutdownContext, cancel := context.WithDeadline( + context.Background(), + deadline, + ) + _ = server.Shutdown(shutdownContext) + cancel() + return waitForChildUntil(child, childExited, deadline) + } +} + +func main() { + config, command, err := parseConfig(os.Args[1:]) + if err != nil { + fmt.Fprintln(os.Stderr, err.Error()) + os.Exit(2) + } + os.Exit(run(config, command)) +} diff --git a/managed-inference/images/llama-cpp/request-guard/main_test.go b/managed-inference/images/llama-cpp/request-guard/main_test.go new file mode 100644 index 00000000000..89b09c1cfe9 --- /dev/null +++ b/managed-inference/images/llama-cpp/request-guard/main_test.go @@ -0,0 +1,689 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "syscall" + "testing" + "time" +) + +func testConfig() guardConfig { + return guardConfig{ + listenHost: "0.0.0.0", + listenPort: 8081, + upstreamHost: "127.0.0.1", + upstreamPort: 0, + maxRequestBodyBytes: 1024, + maxRequestHeaderBytes: 4096, + maxOutputTokens: 32, + requestTimeout: 10 * time.Second, + shutdownTimeout: 5 * time.Second, + } +} + +func configForServer(t *testing.T, upstream *httptest.Server) guardConfig { + t.Helper() + address := strings.TrimPrefix(upstream.URL, "http://") + host, portText, found := strings.Cut(address, ":") + if !found { + t.Fatal("test upstream has no port") + } + var port int + if _, err := fmt.Sscanf(portText, "%d", &port); err != nil { + t.Fatalf("parse test upstream port: %v", err) + } + config := testConfig() + config.upstreamHost = host + config.upstreamPort = port + return config +} + +func guardedServer(t *testing.T, upstream http.Handler) (*httptest.Server, *atomic.Int32) { + t.Helper() + var calls atomic.Int32 + backend := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + calls.Add(1) + upstream.ServeHTTP(writer, request) + })) + t.Cleanup(backend.Close) + handler, err := newGuardHandler(configForServer(t, backend)) + if err != nil { + t.Fatalf("create guard: %v", err) + } + guard := httptest.NewServer(handler) + t.Cleanup(guard.Close) + return guard, &calls +} + +func request(t *testing.T, method, endpoint, contentType string, body io.Reader) *http.Response { + t.Helper() + req, err := http.NewRequest(method, endpoint, body) + if err != nil { + t.Fatalf("create request: %v", err) + } + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } + response, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("send request: %v", err) + } + return response +} + +func responseCode(t *testing.T, response *http.Response) string { + t.Helper() + defer response.Body.Close() + var payload struct { + Error struct { + Code string `json:"code"` + } `json:"error"` + } + if err := json.NewDecoder(response.Body).Decode(&payload); err != nil { + t.Fatalf("decode response: %v", err) + } + return payload.Error.Code +} + +func TestParseConfigRequiresEveryDeclaredValue(t *testing.T) { + valid := []string{ + "--listen-host", "0.0.0.0", + "--listen-port", "8081", + "--upstream-host", "127.0.0.1", + "--upstream-port", "8082", + "--max-request-body-bytes", "1048576", + "--max-request-header-bytes", "32768", + "--max-output-tokens", "4096", + "--request-timeout-seconds", "900", + "--shutdown-timeout-seconds", "25", + "--", llamaServerPath, + "--model", "/models/model.gguf", + "--host", "127.0.0.1", + "--port", "8082", + "--api-key-file", llamaServerAPIKeyPath, + "--n-predict", "4096", + "--no-ui", + "--no-slots", + "--no-mmproj", + "--no-agent", + } + config, command, err := parseConfig(valid) + if err != nil { + t.Fatalf("parse valid config: %v", err) + } + if config.maxRequestBodyBytes != 1048576 || + config.maxOutputTokens != 4096 || + config.shutdownTimeout != 25*time.Second { + t.Fatalf("declared bounds were not retained: %+v", config) + } + if len(command) < 3 || command[0] != llamaServerPath { + t.Fatalf("unexpected command: %v", command) + } + + for _, remove := range []string{ + "--listen-host", + "--listen-port", + "--upstream-host", + "--upstream-port", + "--max-request-body-bytes", + "--max-request-header-bytes", + "--max-output-tokens", + "--request-timeout-seconds", + "--shutdown-timeout-seconds", + } { + candidate := append([]string(nil), valid...) + for index, value := range candidate { + if value == remove { + candidate = append(candidate[:index], candidate[index+2:]...) + break + } + } + if _, _, err := parseConfig(candidate); err == nil { + t.Fatalf("configuration without %s was accepted", remove) + } + } +} + +func TestParseConfigRejectsABypassableLlamaServerCommand(t *testing.T) { + base := []string{ + "--listen-host", "0.0.0.0", + "--listen-port", "8081", + "--upstream-host", "127.0.0.1", + "--upstream-port", "8082", + "--max-request-body-bytes", "1048576", + "--max-request-header-bytes", "32768", + "--max-output-tokens", "4096", + "--request-timeout-seconds", "900", + "--shutdown-timeout-seconds", "25", + "--", llamaServerPath, + "--host", "127.0.0.1", + "--port", "8082", + "--api-key-file", llamaServerAPIKeyPath, + "--n-predict", "4096", + "--no-ui", "--no-slots", "--no-mmproj", "--no-agent", + } + for _, mutation := range []struct { + name string + option string + replacement string + marker bool + }{ + {name: "executable", replacement: "/bin/sh"}, + {name: "host", option: "--host", replacement: "0.0.0.0"}, + {name: "api key", option: "--api-key-file", replacement: "/tmp/key"}, + {name: "token bound", option: "--n-predict", replacement: "4097"}, + {name: "disabled route", option: "--no-ui", replacement: "--metrics", marker: true}, + } { + t.Run(mutation.name, func(t *testing.T) { + candidate := append([]string(nil), base...) + separator := -1 + for index, value := range candidate { + if value == "--" { + separator = index + break + } + } + if separator < 0 || separator+1 >= len(candidate) { + t.Fatal("base command has no child command") + } + if mutation.option == "" { + candidate[separator+1] = mutation.replacement + } else { + target := -1 + for index := separator + 2; index < len(candidate); index++ { + if candidate[index] == mutation.option { + target = index + break + } + } + if target < 0 { + t.Fatalf("child option %s is absent", mutation.option) + } + if mutation.marker { + candidate[target] = mutation.replacement + } else { + if target+1 >= len(candidate) { + t.Fatalf("child option %s has no value", mutation.option) + } + candidate[target+1] = mutation.replacement + } + } + if _, _, err := parseConfig(candidate); err == nil { + t.Fatalf("child %s bypass was accepted", mutation.name) + } + }) + } + for _, extra := range []string{"--embedding", "--model", "/models/other.gguf"} { + candidate := append(append([]string(nil), base...), extra) + if _, _, err := parseConfig(candidate); err == nil { + t.Fatalf("child command with extra %s was accepted", extra) + } + } +} + +func TestAPIKeyFileMustBeReadableRegularAndNonEmpty(t *testing.T) { + root := t.TempDir() + missing := filepath.Join(root, "missing") + if err := validateAPIKeyFile(missing); err == nil { + t.Fatal("missing API-key file was accepted") + } + if err := validateAPIKeyFile(root); err == nil { + t.Fatal("API-key directory was accepted") + } + empty := filepath.Join(root, "empty") + if err := os.WriteFile(empty, nil, 0600); err != nil { + t.Fatalf("create empty API-key file: %v", err) + } + if err := validateAPIKeyFile(empty); err == nil { + t.Fatal("empty API-key file was accepted") + } + unreadable := filepath.Join(root, "unreadable") + if err := os.WriteFile(unreadable, []byte("opaque-test-key\n"), 0600); err != nil { + t.Fatalf("create unreadable API-key file: %v", err) + } + if err := os.Chmod(unreadable, 0000); err != nil { + t.Fatalf("remove API-key file read permissions: %v", err) + } + t.Run("unreadable", func(t *testing.T) { + probe, err := os.Open(unreadable) + if err == nil { + _ = probe.Close() + t.Skip("test process can read a mode-000 file") + } + if err := validateAPIKeyFile(unreadable); err == nil { + t.Fatal("unreadable API-key file was accepted") + } + }) + valid := filepath.Join(root, "valid") + if err := os.WriteFile(valid, []byte("opaque-test-key\n"), 0600); err != nil { + t.Fatalf("create API-key file: %v", err) + } + if err := validateAPIKeyFile(valid); err != nil { + t.Fatalf("valid API-key file was rejected: %v", err) + } +} + +func TestServerUsesTheDeclaredHeaderAndTimeBounds(t *testing.T) { + config := testConfig() + server := newHTTPServer(config, http.NotFoundHandler()) + if server.MaxHeaderBytes != config.maxRequestHeaderBytes { + t.Fatalf("header limit = %d", server.MaxHeaderBytes) + } + for name, actual := range map[string]time.Duration{ + "read-header": server.ReadHeaderTimeout, + "read": server.ReadTimeout, + "write": server.WriteTimeout, + "idle": server.IdleTimeout, + } { + if actual != config.requestTimeout { + t.Fatalf("%s timeout = %s", name, actual) + } + } +} + +func TestRequestGuardChildHelper(t *testing.T) { + if os.Getenv("NEMOCLAW_REQUEST_GUARD_CHILD_HELPER") != "1" { + return + } + signal.Ignore(syscall.SIGTERM) + _, _ = fmt.Fprintln(os.Stdout, "ready") + select {} +} + +func TestStopChildKillsAtTheDeclaredDeadline(t *testing.T) { + child := exec.Command(os.Args[0], "-test.run=TestRequestGuardChildHelper") + child.Env = append(os.Environ(), "NEMOCLAW_REQUEST_GUARD_CHILD_HELPER=1") + stdout, err := child.StdoutPipe() + if err != nil { + t.Fatalf("create child stdout: %v", err) + } + if err := child.Start(); err != nil { + t.Fatalf("start child: %v", err) + } + ready := bufio.NewScanner(stdout) + if !ready.Scan() || ready.Text() != "ready" { + _ = child.Process.Kill() + t.Fatalf("child did not become ready: %v", ready.Err()) + } + exited := make(chan *os.ProcessState, 1) + go func() { + _ = child.Wait() + exited <- child.ProcessState + }() + + started := time.Now() + code := stopChildWithin(child, exited, syscall.SIGTERM, 50*time.Millisecond) + if elapsed := time.Since(started); elapsed < 40*time.Millisecond || elapsed > 2*time.Second { + t.Fatalf("declared stop deadline elapsed = %s", elapsed) + } + if code == 0 { + t.Fatal("force-killed child exited successfully") + } +} + +func TestGuardRejectsOversizedBodiesBeforeUpstream(t *testing.T) { + guard, calls := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusOK) + })) + + declared := strings.NewReader(strings.Repeat("x", 1025)) + response := request(t, http.MethodPost, guard.URL+"/v1/chat/completions", "application/json", declared) + if response.StatusCode != http.StatusRequestEntityTooLarge { + t.Fatalf("declared body status = %d", response.StatusCode) + } + if code := responseCode(t, response); code != "request_body_too_large" { + t.Fatalf("declared body code = %q", code) + } + + chunkedRequest, err := http.NewRequest( + http.MethodPost, + guard.URL+"/v1/chat/completions", + bufio.NewReader(strings.NewReader(strings.Repeat("y", 1025))), + ) + if err != nil { + t.Fatalf("create chunked request: %v", err) + } + chunkedRequest.ContentLength = -1 + chunkedRequest.Header.Set("Content-Type", "application/json") + chunkedResponse, err := http.DefaultClient.Do(chunkedRequest) + if err != nil { + t.Fatalf("send chunked request: %v", err) + } + if chunkedResponse.StatusCode != http.StatusRequestEntityTooLarge { + t.Fatalf("chunked body status = %d", chunkedResponse.StatusCode) + } + if code := responseCode(t, chunkedResponse); code != "request_body_too_large" { + t.Fatalf("chunked body code = %q", code) + } + if calls.Load() != 0 { + t.Fatalf("upstream received %d oversized requests", calls.Load()) + } +} + +func TestGuardRejectsEveryOutputTokenBypass(t *testing.T) { + guard, calls := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusOK) + })) + for _, body := range []string{ + `{"model":"test","max_tokens":33}`, + `{"model":"test","max_completion_tokens":33}`, + `{"model":"test","n_predict":33}`, + `{"model":"test","max_tokens":-1}`, + `{"model":"test","max_tokens":1.5}`, + `{"model":"test","max_tokens":"32"}`, + `{"model":"test","max_tokens":32,"max_tokens":1}`, + } { + response := request( + t, + http.MethodPost, + guard.URL+"/v1/chat/completions", + "application/json", + strings.NewReader(body), + ) + if response.StatusCode != http.StatusBadRequest { + t.Fatalf("body %s status = %d", body, response.StatusCode) + } + response.Body.Close() + } + if calls.Load() != 0 { + t.Fatalf("upstream received %d denied token requests", calls.Load()) + } +} + +func TestGuardInjectsTheDeclaredLimitAndPreservesAuthorization(t *testing.T) { + var seenAuthorization string + var seenBody map[string]any + guard, calls := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + seenAuthorization = request.Header.Get("Authorization") + if err := json.NewDecoder(request.Body).Decode(&seenBody); err != nil { + t.Errorf("decode upstream body: %v", err) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(writer, `{"ok":true}`) + })) + req, err := http.NewRequest( + http.MethodPost, + guard.URL+"/v1/chat/completions", + strings.NewReader(`{"model":"test","messages":[]}`), + ) + if err != nil { + t.Fatalf("create request: %v", err) + } + req.Header.Set("Authorization", "Bearer opaque-test-value") + req.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("send request: %v", err) + } + response.Body.Close() + if response.StatusCode != http.StatusOK { + t.Fatalf("status = %d", response.StatusCode) + } + if calls.Load() != 1 || seenBody["max_tokens"] != float64(32) { + t.Fatalf("upstream calls/body = %d/%v", calls.Load(), seenBody) + } + if seenAuthorization != "Bearer opaque-test-value" { + t.Fatalf("authorization header changed: %q", seenAuthorization) + } +} + +func TestGuardPreservesBackendAuthenticationFailure(t *testing.T) { + guard, calls := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Header.Get("Authorization") == "" { + writer.WriteHeader(http.StatusUnauthorized) + return + } + writer.WriteHeader(http.StatusOK) + })) + response := request( + t, + http.MethodPost, + guard.URL+"/v1/chat/completions", + "application/json", + strings.NewReader(`{"model":"test","max_tokens":8}`), + ) + response.Body.Close() + if response.StatusCode != http.StatusUnauthorized { + t.Fatalf("unauthenticated status = %d", response.StatusCode) + } + if calls.Load() != 1 { + t.Fatalf("backend authentication calls = %d", calls.Load()) + } +} + +func TestGuardSanitizesUpstreamFailures(t *testing.T) { + upstream := httptest.NewServer(http.NotFoundHandler()) + config := configForServer(t, upstream) + upstream.Close() + handler, err := newGuardHandler(config) + if err != nil { + t.Fatalf("create guard: %v", err) + } + guard := httptest.NewServer(handler) + defer guard.Close() + + response := request(t, http.MethodGet, guard.URL+"/health", "", nil) + body, err := io.ReadAll(response.Body) + response.Body.Close() + if err != nil { + t.Fatalf("read response: %v", err) + } + if response.StatusCode != http.StatusBadGateway { + t.Fatalf("status = %d", response.StatusCode) + } + if strings.Contains(string(body), config.upstreamHost) || + strings.Contains(string(body), fmt.Sprintf("%d", config.upstreamPort)) { + t.Fatalf("upstream address leaked: %s", body) + } + if !strings.Contains(string(body), `"code":"upstream_unavailable"`) { + t.Fatalf("unexpected error body: %s", body) + } +} + +func TestGuardRejectsEncodedOrMislabeledChatBodiesBeforeUpstream(t *testing.T) { + guard, calls := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusOK) + })) + encoded, err := http.NewRequest( + http.MethodPost, + guard.URL+"/v1/chat/completions", + strings.NewReader(`{"model":"test","max_tokens":8}`), + ) + if err != nil { + t.Fatalf("create encoded request: %v", err) + } + encoded.Header.Set("Content-Encoding", "gzip") + encoded.Header.Set("Content-Type", "application/json") + encodedResponse, err := http.DefaultClient.Do(encoded) + if err != nil { + t.Fatalf("send encoded request: %v", err) + } + if encodedResponse.StatusCode != http.StatusUnsupportedMediaType { + t.Fatalf("encoded status = %d", encodedResponse.StatusCode) + } + encodedResponse.Body.Close() + + mislabeled := request( + t, + http.MethodPost, + guard.URL+"/v1/chat/completions", + "text/plain", + strings.NewReader(`{"model":"test","max_tokens":8}`), + ) + if mislabeled.StatusCode != http.StatusUnsupportedMediaType { + t.Fatalf("mislabeled status = %d", mislabeled.StatusCode) + } + mislabeled.Body.Close() + if calls.Load() != 0 { + t.Fatalf("upstream received %d encoded or mislabeled requests", calls.Load()) + } +} + +func TestGuardStreamsAllowedResponses(t *testing.T) { + release := make(chan struct{}) + defer func() { + select { + case <-release: + default: + close(release) + } + }() + guard, _ := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + flusher, ok := writer.(http.Flusher) + if !ok { + t.Error("upstream response writer cannot flush") + return + } + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(writer, "data: first\n\n") + flusher.Flush() + <-release + _, _ = io.WriteString(writer, "data: second\n\n") + })) + req, err := http.NewRequest( + http.MethodPost, + guard.URL+"/v1/chat/completions", + strings.NewReader(`{"model":"test","stream":true,"max_tokens":32}`), + ) + if err != nil { + t.Fatalf("create stream request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + client := &http.Client{Timeout: 2 * time.Second} + response, err := client.Do(req) + if err != nil { + t.Fatalf("receive streamed response headers: %v", err) + } + defer response.Body.Close() + reader := bufio.NewReader(response.Body) + first := make([]byte, len("data: first\n\n")) + if _, err := io.ReadFull(reader, first); err != nil { + t.Fatalf("read first event before upstream completes: %v", err) + } + if string(first) != "data: first\n\n" { + t.Fatalf("first event = %q", string(first)) + } + close(release) + rest, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("read remaining stream: %v", err) + } + if string(rest) != "data: second\n\n" { + t.Fatalf("remaining stream = %q", string(rest)) + } +} + +func TestGuardPropagatesCancellation(t *testing.T) { + started := make(chan struct{}) + cancelled := make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + releaseBackend := func() { releaseOnce.Do(func() { close(release) }) } + backend := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + _ = request.Body.Close() + close(started) + select { + case <-request.Context().Done(): + close(cancelled) + case <-release: + } + })) + defer func() { + releaseBackend() + backend.CloseClientConnections() + backend.Close() + }() + handler, err := newGuardHandler(configForServer(t, backend)) + if err != nil { + t.Fatalf("create guard: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + req, err := http.NewRequestWithContext( + ctx, + http.MethodPost, + "http://guard.test/v1/chat/completions", + strings.NewReader(`{"model":"test","max_tokens":8}`), + ) + if err != nil { + t.Fatalf("create guard request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + result := make(chan struct{}) + go func() { + handler.ServeHTTP(httptest.NewRecorder(), req) + close(result) + }() + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("upstream request did not start") + } + cancel() + select { + case <-cancelled: + case <-time.After(2 * time.Second): + t.Fatal("upstream request was not cancelled") + } + select { + case <-result: + case <-time.After(2 * time.Second): + t.Fatal("guard handler did not return after cancellation") + } +} + +func TestGuardBlocksUnsupportedServerSurfaces(t *testing.T) { + guard, calls := guardedServer(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusOK) + })) + for _, target := range []string{ + "/", + "/slots", + "/v1/completions", + "/v1/responses", + "/v1/embeddings", + "/v1/chat/completions?debug=true", + } { + response := request(t, http.MethodPost, guard.URL+target, "application/json", bytes.NewReader([]byte(`{}`))) + if response.StatusCode != http.StatusNotFound { + t.Fatalf("target %s status = %d", target, response.StatusCode) + } + response.Body.Close() + } + for _, probe := range []struct { + method string + target string + }{ + {method: http.MethodGet, target: "/v1/chat/completions"}, + {method: http.MethodPost, target: "/health"}, + {method: http.MethodPost, target: "/v1/models"}, + {method: http.MethodDelete, target: "/props"}, + } { + response := request(t, probe.method, guard.URL+probe.target, "", nil) + if response.StatusCode != http.StatusNotFound { + t.Fatalf("%s %s status = %d", probe.method, probe.target, response.StatusCode) + } + response.Body.Close() + } + if calls.Load() != 0 { + t.Fatalf("upstream received %d unsupported requests", calls.Load()) + } +} diff --git a/scripts/checks/export-llama-cpp-image-config.mts b/scripts/checks/export-llama-cpp-image-config.mts index 19273e66abb..5500f7e127d 100644 --- a/scripts/checks/export-llama-cpp-image-config.mts +++ b/scripts/checks/export-llama-cpp-image-config.mts @@ -18,7 +18,10 @@ import { type ServerImageManifest = { apiVersion?: unknown; kind?: unknown; - metadata?: { id?: unknown }; + metadata?: { + annotations?: { "nemoclaw.nvidia.com/request-guard-state"?: unknown }; + id?: unknown; + }; spec?: { build?: { backendDirectory?: unknown; @@ -296,7 +299,10 @@ export function loadLlamaCppImageConfig( const recipe = parseQualificationRecipe(recipeSource, recipeSchemaSource); const agentQualification = parseAgentQualificationDocument(agentQualificationSource); assertExactKeys(manifest, "manifest", ["apiVersion", "kind", "metadata", "spec"]); - assertExactKeys(manifest.metadata, "metadata", ["id"]); + assertExactKeys(manifest.metadata, "metadata", ["annotations", "id"]); + assertExactKeys(manifest.metadata?.annotations, "metadata annotations", [ + "nemoclaw.nvidia.com/request-guard-state", + ]); assertExactKeys(manifest.spec, "spec", [ "build", "cuda", @@ -391,7 +397,8 @@ export function loadLlamaCppImageConfig( if ( manifest?.apiVersion !== "nemoclaw.nvidia.com/managed-inference/v1" || manifest?.kind !== "ServerImageBuild" || - manifest?.metadata?.id !== "llama-cpp-server.v1" + manifest?.metadata?.id !== "llama-cpp-server.v1" || + manifest?.metadata?.annotations?.["nemoclaw.nvidia.com/request-guard-state"] !== "dormant" ) { throw new Error("invalid llama.cpp server image manifest identity"); } @@ -607,6 +614,7 @@ export function loadLlamaCppImageConfig( curl: "8.5.0-2ubuntu10.11", "g++-14": "14.2.0-4ubuntu2~24.04.1", "gcc-14": "14.2.0-4ubuntu2~24.04.1", + "golang-go": "2:1.22~2build1", "libcurl4-openssl-dev": "8.5.0-2ubuntu10.11", "libssl-dev": "3.0.13-0ubuntu3.12", }; @@ -623,6 +631,8 @@ export function loadLlamaCppImageConfig( const expectedRequiredPaths = [ "/opt/llama.cpp/lib/libggml-cuda.so", "/usr/local/bin/llama-server", + "/usr/local/bin/nemoclaw-llama-cpp-request-guard", + "/usr/local/share/licenses/go/copyright", "/usr/local/share/licenses/llama.cpp/AUTHORS", "/usr/local/share/licenses/llama.cpp/LICENSE", ]; diff --git a/test/llama-cpp-image.test.ts b/test/llama-cpp-image.test.ts index 4af431ad512..2fdfbc4641d 100644 --- a/test/llama-cpp-image.test.ts +++ b/test/llama-cpp-image.test.ts @@ -31,7 +31,10 @@ const exporterPath = path.join(repoRoot, "scripts", "checks", "export-llama-cpp- type ImageManifest = { apiVersion?: string; kind?: string; - metadata?: { id?: string }; + metadata?: { + annotations?: { "nemoclaw.nvidia.com/request-guard-state"?: string }; + id?: string; + }; spec?: { build?: { backendDirectory?: string; @@ -148,6 +151,18 @@ function enablePublication(source: string): string { }); } +function configureManifestAnnotations(source: string, annotations: Record): string { + const candidate = YAML.parse(source) as { metadata: { annotations?: Record } }; + candidate.metadata.annotations = annotations; + return YAML.stringify(candidate); +} + +function removeManifestAnnotations(source: string): string { + const candidate = YAML.parse(source) as { metadata: { annotations?: Record } }; + delete candidate.metadata.annotations; + return YAML.stringify(candidate); +} + describe("declarative llama.cpp server image", () => { const manifestSource = fs.readFileSync(manifestPath, "utf8"); const manifest = YAML.parse(manifestSource) as ImageManifest; @@ -158,7 +173,12 @@ describe("declarative llama.cpp server image", () => { expect(manifest).toMatchObject({ apiVersion: "nemoclaw.nvidia.com/managed-inference/v1", kind: "ServerImageBuild", - metadata: { id: "llama-cpp-server.v1" }, + metadata: { + id: "llama-cpp-server.v1", + annotations: { + "nemoclaw.nvidia.com/request-guard-state": "dormant", + }, + }, spec: { repository: "ghcr.io/nvidia/nemoclaw/llama-cpp-server", build: { backendDirectory: "/opt/llama.cpp/lib" }, @@ -169,6 +189,8 @@ describe("declarative llama.cpp server image", () => { requiredPaths: expect.arrayContaining([ "/opt/llama.cpp/lib/libggml-cuda.so", "/usr/local/bin/llama-server", + "/usr/local/bin/nemoclaw-llama-cpp-request-guard", + "/usr/local/share/licenses/go/copyright", "/usr/local/share/licenses/llama.cpp/LICENSE", ]), writablePaths: ["/tmp"], @@ -363,6 +385,20 @@ describe("declarative llama.cpp server image", () => { "an unexpected top-level field", manifestSource.replace("kind: ServerImageBuild", "kind: ServerImageBuild\nunexpected: true"), ], + ["a missing request guard state annotation", removeManifestAnnotations(manifestSource)], + [ + "a non-dormant request guard state annotation", + configureManifestAnnotations(manifestSource, { + "nemoclaw.nvidia.com/request-guard-state": "active", + }), + ], + [ + "an unexpected image annotation", + configureManifestAnnotations(manifestSource, { + "nemoclaw.nvidia.com/request-guard-state": "dormant", + "nemoclaw.nvidia.com/unexpected": "true", + }), + ], ])("rejects %s before exporting image build inputs (#8231)", (_case, candidate) => { expect(() => loadLlamaCppImageConfig(candidate)).toThrow(); }); @@ -659,6 +695,16 @@ describe("declarative llama.cpp server image", () => { expect(dockerfile).toContain("USER ${RUNTIME_UID}:${RUNTIME_GID}"); expect(dockerfile).toContain('SHELL ["/bin/bash", "-o", "pipefail", "-c"]'); expect(dockerfile).toContain('ENTRYPOINT ["/usr/local/bin/llama-server"]'); + expect(dockerfile).toContain("go test ./..."); + expect(dockerfile).toContain("CGO_ENABLED=0 go build"); + expect(dockerfile).toContain("set -- /usr/share/doc/golang-[0-9]*-go/copyright"); + expect(dockerfile).toContain('cp "$1" /opt/llama.cpp/licenses/go/copyright'); + expect(dockerfile).toContain( + "COPY --from=build --chmod=0555 /opt/llama.cpp/bin/nemoclaw-llama-cpp-request-guard /usr/local/bin/nemoclaw-llama-cpp-request-guard", + ); + expect(dockerfile).toContain( + "COPY --from=build /opt/llama.cpp/licenses/ /usr/local/share/licenses/", + ); expect(dockerfile).toContain("ENV CC=${C_COMPILER}"); expect(dockerfile).toContain("CXX=${CXX_COMPILER}"); expect(dockerfile).toContain("CUDAHOSTCXX=${CUDA_HOST_CXX_COMPILER}");