Skip to content

fix(beta9): 5 bugs from Cubic review of upstream PR — status-checking, idempotency, NodePort Service, SDK batch-embeddings - #3

Merged
Wingie merged 5 commits into
mainfrom
fix/cubic-review-findings-20260420
Jul 26, 2026
Merged

fix(beta9): 5 bugs from Cubic review of upstream PR — status-checking, idempotency, NodePort Service, SDK batch-embeddings#3
Wingie merged 5 commits into
mainfrom
fix/cubic-review-findings-20260420

Conversation

@Wingie

@Wingie Wingie commented Apr 20, 2026

Copy link
Copy Markdown
Owner

Summary

Addresses 5 bugs Cubic AI flagged on the (now-closed) upstream PR beam-cloud#1562. All fixes land with tests; full Go + SDK suites pass (62/62 SDK, was 61/62).

Branch off origin/main; one commit per fix for clean review.

Fixes

Commit Finding File
e463ced4 SDK batch-embedding field routing — embedding now receives a single vector, embeddings receives the batch list-of-lists sdk/src/beta9/inference.py
55adb935 UnloadModel status-checks worker response; non-200 logs error with ≤4 KiB body, returns error, leaves LoadState unchanged pkg/gateway/inference_router.go
cd7acd37 /inference/pull status-checks Ollama before decoding; decode errors return 502; includes isAllowedModelName() allowlist (blocks ..) + MaxBytesReader(4096) pkg/agent/control.go
241e932a KeepaliveLoop Start/Stop guarded with sync.Once + atomic started flag; Stop-before-Start no longer hangs; Stop is idempotent; concurrent Stop is safe pkg/agent/keepalive.go
14d47d1f Added Service/gateway-nodeport (type NodePort, ports 31994 + 31993, selector app=gateway) alongside the existing ClusterIP Service manifests/k3d/beta9.yaml

Tests added / updated

  • sdk/tests/test_cubic_fixes.py::test_batch_embeddings — pre-existing, now passes
  • pkg/gateway/inference_router_test.go — new — TestUnloadModelNonOK (200/500/404 table) + TestUnloadModelNodeNotFound
  • pkg/agent/control_test.go — new — TestInferencePullStatus (200/404/500) + TestInferencePullDecodeFailureReturnsBadGateway + TestInferencePullRejectsInvalidModelName + TestInferencePullMethodNotAllowed
  • pkg/agent/keepalive_test.go — new — TestKeepaliveLoopIdempotent + TestKeepaliveLoopConcurrentStop (32 goroutines) + TestKeepaliveLoopStopBeforeStart

Verification

  • go build ./... — clean
  • go test ./... -count=1 -timeout=300s — all packages PASS
  • go vet ./... — clean on touched files (pre-existing warnings in pkg/worker/worker.go and pkg/gateway/services/stub.go unchanged)
  • cd sdk && python -m pytest tests/ -x --tb=short62 passed (was 61 passed / 1 failed before this PR)

Notes

  • The NodePort Service (commit 14d47d1) was validated via Python YAML parse (the build host had no kubectl). Worth re-verifying with kubectl apply --dry-run=client -f manifests/k3d/beta9.yaml before apply.
  • Cubic's line numbers drifted vs current origin/main; fixes are by function identity, not line — the underlying bugs were unchanged.
  • /inference/pull allowlist and body cap are inlined here (not pulled from security-hardening-b9agent branch) since this PR is based on origin/main for a clean diff.
  • Fix context (all 40 Cubic findings + cross-check against our earlier audits): https://github.com/Wingie/flowstate-agents/blob/main/wip-specs/beta9/cubic-review-findings-1562-2026-04-20.md

Co-Authored-By: Claudistrator savetheplanet@agentosaurus.com

Wingie and others added 5 commits April 20, 2026 15:43
Cubic review finding (sdk/src/beta9/inference.py:307):
Batch embed() results were being assigned to the singular `embedding`
field while `embeddings` (the documented plural list-of-lists) stayed
empty. tests/test_cubic_fixes.py::test_batch_embeddings exercises this
path and asserts `len(result.embeddings) == 3`.

Fix: track whether caller passed a string (single) or a list (batch),
and populate:
  - `embedding` = flat vector for single-input calls (back-compat)
  - `embeddings` = full list-of-lists for batch calls

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
Cubic review finding (pkg/gateway/inference_router.go:463):
UnloadModel previously cleared the registry's LoadState to Idle no
matter what the worker returned. Any 4xx/5xx response would leave the
registry inconsistent with the worker's actual in-memory model set —
a cache-coherency bug.

Fix:
- On non-200 response: log error with status + up to 4 KiB of body,
  return an error, and leave the model's LoadState unchanged.
- Only on 200: drain body and transition to LoadStateIdle.

Added table-driven TestUnloadModelNonOK + TestUnloadModelNodeNotFound
in pkg/gateway/inference_router_test.go covering 200/500/404 paths
and the missing-node early return.

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
…broke

Cubic review finding (pkg/agent/control.go:227):
handleInferencePull forwarded to Ollama's /api/pull, looped over the
NDJSON body with a decoder, and returned StatusOK unconditionally. A
non-200 from Ollama or a malformed stream both produced a false
success response.

Fix:
- Check resp.StatusCode != 200 before streaming; non-OK now returns
  502 with a 512-byte body preview from Ollama for diagnosis.
- Track decoder errors; if decode fails (or the stream is empty) the
  handler returns 502 instead of pretending success.
- Success case now returns 201 Created to distinguish from the pre-fix
  false-success 200.

Inline P0-A hardening (required by spec, not yet on origin/main):
- MaxBytesReader caps request body at 4 KiB.
- isAllowedModelName() allowlist [A-Za-z0-9._/:-] + length <= 128 and
  blocks ".." path-traversal sequences before they reach Ollama.

Added TestInferencePullStatus (200/404/500 upstream),
TestInferencePullDecodeFailureReturnsBadGateway,
TestInferencePullRejectsInvalidModelName, TestInferencePullMethodNotAllowed
in pkg/agent/control_test.go. ControlServer gained an optional
ollamaBaseURL field so httptest servers can stand in for the daemon.

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
Cubic review finding (pkg/agent/keepalive.go, Start ~L113 / Stop):
Calling Stop() twice — or Start() after Stop() — closed an already
closed channel and produced a "close of closed channel" panic. The
bug is reachable during shutdown races and during retries after a
transient failure.

Fix:
- Added `startOnce` + `stopOnce` sync.Once guards plus an atomic
  `started` flag. Start is a no-op after the first call; subsequent
  Stops are no-ops too.
- Stop-before-Start is also safe: when started == 0 we return without
  waiting on doneCh (which would otherwise never close).

Tests in pkg/agent/keepalive_test.go:
- TestKeepaliveLoopIdempotent — Start, Stop, Stop, Start, Stop
- TestKeepaliveLoopConcurrentStop — 32 goroutines stopping in parallel
- TestKeepaliveLoopStopBeforeStart — Stop never blocks pre-Start

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
Cubic review finding (manifests/k3d/beta9.yaml around L361):
The manifest's comment declared "using NodePort instead" but the only
gateway Service was ClusterIP on 1994/1993/9090. Tailscale peers and
external workers hitting :31994/:31993 on the OCI node saw connection
refused because no NodePort Service ever materialised — matches
empty output of `ss -tlnp | grep 31994` on the live host.

Fix: add `Service/gateway-nodeport` (type: NodePort) alongside the
existing ClusterIP Service, with:
  - nodePort: 31994 -> targetPort 1994 (HTTP)
  - nodePort: 31993 -> targetPort 1993 (gRPC)
  - selector: app=gateway

Kept the ClusterIP Service unchanged for in-cluster traffic.
Validated with `python3 -c "import yaml; yaml.safe_load_all(...)"`;
kubectl dry-run not available on the build host but the Service
matches the Kubernetes v1 Service schema used elsewhere in the repo.

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
@Wingie

Wingie commented Jul 26, 2026

Copy link
Copy Markdown
Owner Author

Manual review (no @claude workflow on this repo, so no bot review is coming)

Verdict: APPROVE. Five fixes, one commit each, 9 new Go tests, suites green (62/62 SDK, was 61/62). This is the shape a bugfix PR should have.

isAllowedModelName() is built correctly — an allowlist, not a blocklist: length-bounded (1–128), explicit .. rejection, explicit character whitelist with default: return false. URL-encoded traversal (%2e%2e) can't slip through because % isn't permitted. Rejecting on strings.Contains(name, "..") is the right call here — it over-rejects rather than under-rejects, which is the correct failure direction for a security guard.

One gap worth a follow-up: / is allowed and there's no leading-slash check, so /etc/passwd passes the allowlist. .. blocks upward traversal, but absolute paths aren't rejected. Whether that's exploitable depends on how Ollama consumes the name (a leading / probably just 404s upstream), so not blocking — but if strings.HasPrefix(name, "/") { return false } would close it for the cost of one line, and defense-in-depth on a path-ish input is cheap.

The rest read well: status-checking before decode (the old code swallowed decode errors and returned success unconditionally), MaxBytesReader(4096) and the ≤4 KiB error-body cap bounding untrusted input, and the sync.Once + atomic guard fixing a real Stop-before-Start hang and making Stop idempotent under concurrency. Leaving LoadState unchanged on a non-200 UnloadModel is the right conservative choice.

Merging with --merge rather than squash, consistent with how the other submodules are handled here, so the head SHA stays reachable if a parent gitlink ever pins it.

@Wingie
Wingie merged commit 689961d into main Jul 26, 2026
2 of 3 checks passed
Wingie added a commit that referenced this pull request Aug 3, 2026
- Add GPUType field to InferenceStatus in keepalive payload
- Add InferenceGPUType to AgentState and AgentStateSnapshot
- Add UpdateInferenceWithGPU() method; UpdateInference() stays for compat
- StartInference() reads gpu_type from OllamaManager.GetStatus() after start
- Gateway keepalive handler uses reported GPUType instead of hardcoded "MPS"
- Falls back to "MPS" for old agents that omit the field (zero-value compat)

Fixes the TODO at pkg/api/v1/machine.go:289.

Rebased onto origin/main (689961d, PR #3 cubic-review-findings) — the
prior branch tip was stacked on unmerged PR #4 (security fixes) which had
since diverged from main, producing conflicts. Re-resolved cleanly: kept
main's current locally-scoped NodeInferenceInfo type in machine.go (main
never migrated it to types.NodeInferenceInfo, unlike what the original
fvks-branch commit assumed) and layered the gpuType variable/fallback on
top of it.

BD: FlowState-fvks

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
Wingie added a commit that referenced this pull request Aug 3, 2026
Both branches independently fixed the same two bugs; kept the stronger
side of each and dropped the now-redundant duplicate:

- pkg/agent/control.go: two overlapping model-name validators
  (isValidOllamaModelName from this branch, isAllowedModelName from
  origin/main PR #3). Kept isValidOllamaModelName — it's a superset
  (also rejects fully-qualified registry paths, not just shell
  metacharacters/traversal) — and removed the now-unused duplicate.
  Existing TestInferencePullRejectsInvalidModelName payloads (shell
  metachars, path traversal, empty) are all still rejected.

- sdk/src/beta9/inference.py: two fixes for the same embed() batch-
  routing bug. Kept this branch's simpler always-populate-both-fields
  version (matches the EmbeddingResult dataclass's own docstring:
  embedding="Legacy: First embedding if batch", embeddings="New: All
  embeddings") over origin/main's is_batch-conditional version, which
  zeroes out .embedding for batches of 2+ results — an untested edge
  case that contradicts the documented contract. test_cubic_fixes.py's
  test_batch_embeddings only asserts .embeddings for the batch case, so
  both variants passed it; picked by contract match, not by the test.

This merge exists so the P0 security fixes here (InferenceRegistry
signature fix — see the added comment in pkg/api/v1/machine.go history —
auth fails-open, unauth control API, SSRF, plaintext creds, gateway
panic) don't get silently dropped from the beta9-infra submodule pin
by rebasing straight onto origin/main, which was the first (wrong)
approach tried for PR #5's GPU-type-propagation conflict.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant