Pickles refused however they come; a string temperature passes again

From review. A mirror set already on the node is looked at again, not
trusted for having been fetched once; Pickled cuts a query or a fragment
from any source, hf: too; .pt2 is a pickle. The flag check reads a flag as
the engine does: underscores as dashes and any beginning of a refused flag
as that flag, since argparse takes --trust-remote for --trust-remote-code;
SGLang refuses --load-format, which would undo the agent's safetensors.

TripoSR's preset says allow_pickle: without it the preset could not be
added and the presets pass failed on it each time. A shared ComfyUI
environment's hidden entry takes its first member's allow_pickle, as it
takes its source.

The GenAI trace attributes made temperature and top_p typed fields of the
request's parse, so "0.7" as a string was a 400; they are read loosely,
as they are read only for the trace.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

CHANGELOG.md

@@ -39,6 +39,17 @@ APIs change between minors without a deprecation period.
 
 ### Fixed
 
+- **Weights that run code got past the refusal in four ways.** A mirror
+  version already on a node was used without being looked at again; a
+  `?query` or `#fragment` on an `hf:` source hid the file's name; `.pt2` was
+  not on the list; and a profile could write `--load_format`,
+  `--trust_remote_code` or an abbreviation such as `--trust-remote`, which
+  vLLM and argparse read as the flag, or give SGLang `--load-format`. All
+  refused now. TripoSR's preset says `allow_pickle`, without which it could
+  not be added, and a shared ComfyUI environment takes its first member's.
+- **A temperature written as a string was refused.** Since the GenAI traces
+  read it, `"temperature": "0.7"` got a 400 before reaching the engine; it
+  is read loosely again, for the trace only.
 - **A reused answer could carry another caller's personal data.** With
   `restore: true` and the answer cache on, two callers' requests that
   redacted to the same text shared one answer — with the first caller's

docs/install.md

@@ -891,8 +891,9 @@ files were checked, can say so:
   allow_pickle: true
 ```
 
-A file put on the machine by hand, as a path, is the operator's own and is not
-looked at.
+A local path is judged by its name as any source is; a local directory is
+not looked into, nor a model an engine downloads itself other than by its
+Hugging Face id (an adapter by id, TensorRT-LLM's checkpoints).
 
 Put them where the node can read them. `scp`, a shared filesystem, or the
 engine's own downloader are all valid, and which you want depends on how many
@@ -2326,6 +2327,7 @@ triposr`. What that adds is this entry, and its workflow is built in:
 ```yaml
 - name: triposr
   source: hf:stabilityai/TripoSR/model.ckpt
+  allow_pickle: true               # published as this checkpoint only
   comfyui:
     workflow: triposr              # built in; a model of your own names a JSON file on every control plane
     output: meshes

internal/agent/engines/pickle_test.go

@@ -20,3 +20,25 @@ func TestAModelDownloadedByTheEngineIsReadAsSafetensorsOnly(t *testing.T) {
 		t.Error("allow_pickle still forced safetensors")
 	}
 }
+
+func TestAForbiddenFlagIsKnownHoweverItIsSpelled(t *testing.T) {
+	for engine, args := range map[string][]string{
+		"vllm":   {"--trust_remote_code"},
+		"sglang": {"--load-format", "pt"},
+	} {
+		if UnsafeArgs(engine, args) == nil {
+			t.Errorf("%s %v was let through", engine, args)
+		}
+	}
+	for _, a := range [][]string{{"--load_format", "pt"}, {"--trust-remote"}, {"--load-form=pt"}} {
+		if UnsafeArgs("vllm", a) == nil {
+			t.Errorf("vllm %v was let through", a)
+		}
+	}
+	// What is not one of them, nor the start of one, still passes.
+	for _, a := range [][]string{{"--max-model-len", "8192"}, {"--tokenizer-mode", "mistral"}, {"--gpu-memory-utilization=0.9"}} {
+		if err := UnsafeArgs("vllm", a); err != nil {
+			t.Errorf("vllm %v was refused: %v", a, err)
+		}
+	}
+}

internal/agent/engines/unsafeenv.go

@@ -119,6 +119,7 @@ var flagsNever = map[string]map[string]string{
 		"--file-storage-path": "writes files", "--log-requests": "logs prompts",
 		"--model-path": "is the model's source, set by the agent", "--served-model-name": "is set by the agent",
 		"--lora-paths": "loads files", "--download-dir": "writes files",
+		"--load-format": "changes how weights are loaded, some ways running code",
 	},
 	// ComfyUI's main.py, behind the gate that sets where it listens and
 	// where its files are.
@@ -164,9 +165,23 @@ var flagsNever = map[string]map[string]string{
 // unsafeFlag is why an engine may not be given a flag, or "".
 func unsafeFlag(engine, flag, _ string) string {
 	name, _, _ := strings.Cut(flag, "=")
+	if !strings.HasPrefix(name, "--") {
+		return ""
+	}
+	// Written as the engine would read it: vLLM takes --load_format for
+	// --load-format, TensorRT-LLM the other way round, and argparse takes
+	// any unambiguous beginning of a flag for the flag — --trust-remote is
+	// --trust-remote-code.
+	dashed := "--" + strings.ReplaceAll(name[2:], "_", "-")
 	for _, set := range []string{"", engine} {
-		if why, ok := flagsNever[set][name]; ok && why != "" {
-			return why
+		for never, why := range flagsNever[set] {
+			if why == "" {
+				continue
+			}
+			n := "--" + strings.ReplaceAll(never[2:], "_", "-")
+			if n == dashed || (len(dashed) > 4 && strings.HasPrefix(n, dashed)) {
+				return why
+			}
 		}
 	}
 	return ""

internal/agent/weights/pickle.go

@@ -16,11 +16,13 @@ import (
 
 // pickleExts are the extensions of PyTorch's pickle formats, and of pickle
 // itself. A .bin beside safetensors is the same model in pickle form.
-var pickleExts = []string{".bin", ".pt", ".pth", ".ckpt", ".pkl", ".pickle"}
+var pickleExts = []string{".bin", ".pt", ".pt2", ".pth", ".ckpt", ".pkl", ".pickle"}
 
 // Pickled reports whether a file name or URL is in a pickle format.
 func Pickled(name string) bool {
-	if i := strings.IndexAny(name, "?#"); i >= 0 && strings.Contains(name, "://") {
+	// A query or a fragment is not the file's name, on a URL or an hf:
+	// source alike: model.bin?download=true is still model.bin.
+	if i := strings.IndexAny(name, "?#"); i >= 0 {
 		name = name[:i]
 	}
 	ext := strings.ToLower(path.Ext(name))

internal/agent/weights/pickle_test.go

@@ -10,6 +10,7 @@ func TestPickleFormatsAreKnownByTheirNames(t *testing.T) {
 	for name, want := range map[string]bool{
 		"pytorch_model.bin": true, "model.ckpt": true, "x.PT": true, "w.pth": true, "a.pkl": true,
 		"https://h/r/resolve/main/model.ckpt?download=true": true,
+		"hf:o/r/model.bin?download=true":                    true, "hf:o/r/model.ckpt#x": true, "m.pt2": true,
 		"model.safetensors": false, "m-Q4_K_M.gguf": false, "Qwen/Qwen2.5-7B-Instruct": false,
 		"config.json": false, "": false,
 	} {
@@ -51,3 +52,15 @@ func TestASetWithAPickleFileIsRefusedWhole(t *testing.T) {
 		t.Errorf("allowed, the set was still refused: %v", err)
 	}
 }
+
+func TestASetAlreadyHereIsLookedAtAgain(t *testing.T) {
+	// Fetched once for a model that allowed pickles, or before they were
+	// refused: another asking for it without leave is refused all the same.
+	f, srv, _ := mirrorOf(t, map[string]string{"pytorch_model.bin": "pickle", "config.json": "{}"})
+	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", true); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); !errors.Is(err, ErrPickle) {
+		t.Errorf("a set already on the node was not looked at: %v", err)
+	}
+}

internal/agent/weights/set.go

@@ -56,6 +56,16 @@ func (f *Fetcher) ensureSet(ctx context.Context, source string, allowPickle bool
 	// exactly the failure the marker exists to prevent.
 	done := filepath.Join(dir, ".complete")
 	if _, err := os.Stat(done); err == nil {
+		// Already here, maybe fetched for a model that allowed pickles, or
+		// before they were refused: looked at again, not trusted.
+		if !allowPickle {
+			entries, _ := os.ReadDir(dir)
+			for _, e := range entries {
+				if Pickled(e.Name()) {
+					return "", refusePickle(source + e.Name())
+				}
+			}
+		}
 		return dir, nil
 	}
 

internal/control/catalog/pickle_test.go

@@ -28,3 +28,13 @@ func TestAModelInAPickleFormatIsRefusedUnlessItSaysSo(t *testing.T) {
 		t.Errorf("a GGUF was refused: %v", err)
 	}
 }
+
+func TestTheTripoSRPresetSaysItsCheckpointIsAllowed(t *testing.T) {
+	p, ok := PresetNamed("triposr")
+	if !ok {
+		t.Fatal("no triposr preset")
+	}
+	if err := p.Model.CheckPickle(); err != nil {
+		t.Errorf("the preset llmgw ships is refused by its own check: %v", err)
+	}
+}

internal/control/catalog/presets.go

@@ -52,6 +52,9 @@ func Presets() []Preset {
 				Name:        "triposr",
 				Description: "TripoSR: a 3D model (GLB) from a picture of one object on a plain background, as a job",
 				Source:      "hf:stabilityai/TripoSR/model.ckpt",
+				// TripoSR is published as this checkpoint only; it is
+				// Stability AI's, pinned, and the one pickle llmgw ships.
+				AllowPickle: true,
 				ComfyUI: &ComfyUI{
 					Workflow: "triposr",
 					Output:   OutputMeshes,

internal/control/catalog/shared.go

@@ -99,6 +99,7 @@ func envModel(env Environment, members []Model) Model {
 		Role:        RoleEnvironment,
 		Description: "the shared ComfyUI replica of environment " + env.Name,
 		Source:      first.Source,
+		AllowPickle: first.AllowPickle,
 		ComfyUI:     &ComfyUI{Environment: env.Name, Nodes: env.Nodes, Output: OutputEnvironment},
 		Profiles:    profs,
 		Scaling:     Scaling{Min: minN, Max: maxN, IdleTimeout: idle},

internal/router/handlers.go

@@ -516,8 +516,10 @@ type inferenceBody struct {
 	MaxTokens           int             `json:"max_tokens"`
 	MaxCompletionTokens int             `json:"max_completion_tokens"`
 	// Read only to say them in the request's trace; sent on untouched.
-	Temperature   *float64 `json:"temperature"`
-	TopP          *float64 `json:"top_p"`
+	// Temperature and TopP are read loosely, for the trace only: a client
+	// that sends "0.7" is the engine's to accept or refuse, as before.
+	Temperature   looseFloat `json:"temperature"`
+	TopP          looseFloat `json:"top_p"`
 	StreamOptions *struct {
 		IncludeUsage bool `json:"include_usage"`
 	} `json:"stream_options"`
@@ -907,8 +909,8 @@ func (r *Router) parse(req *http.Request, path string) (*parsedRequest, error) {
 		Band:            req.Header.Get(r.cfg.Admission.BandHeader),
 		ProfileOverride: req.Header.Get("x-llmgw-profile"),
 		MaxTokens:       max(b.MaxTokens, b.MaxCompletionTokens),
-		Temperature:     b.Temperature,
-		TopP:            b.TopP,
+		Temperature:     b.Temperature.v,
+		TopP:            b.TopP.v,
 		PromptPrefix:    promptPrefix(&b),
 		ThreadID:        req.Header.Get(threadHeader),
 	}
@@ -1536,3 +1538,15 @@ func withAdapter(body []byte, model string, index int, scale float64) []byte {
 	}
 	return out
 }
+
+// looseFloat is a number, or a string holding one, or anything else read as
+// none: a field a client may write in more ways than one, read only to say.
+type looseFloat struct{ v *float64 }
+
+func (f *looseFloat) UnmarshalJSON(b []byte) error {
+	s := strings.Trim(string(b), `"`)
+	if v, err := strconv.ParseFloat(s, 64); err == nil {
+		f.v = &v
+	}
+	return nil
+}

internal/router/trace_test.go

@@ -141,3 +141,14 @@ func TestEveryInferencePathHasItsGenAIOperation(t *testing.T) {
 		}
 	}
 }
+
+func TestATemperatureWrittenAsAStringIsTheEnginesToJudge(t *testing.T) {
+	// Read for the trace only: before it was, "0.7" reached the engine.
+	e := newFakeEngine(t, "m")
+	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))
+	rec := post(t, rt.Handler(), map[string]any{"model": "m", "temperature": "0.7",
+		"messages": []any{map[string]any{"role": "user", "content": "hi"}}}, nil)
+	if rec.Code != 200 {
+		t.Errorf("status %d: %s", rec.Code, rec.Body.String())
+	}
+}