Skip to content

fix(agent): propagate real GPU type through keepalive instead of hardcoding MPS - #5

Open
Wingie wants to merge 6 commits into
mainfrom
claude/beta9-infra/gpu-type-propagation-20260802
Open

fix(agent): propagate real GPU type through keepalive instead of hardcoding MPS#5
Wingie wants to merge 6 commits into
mainfrom
claude/beta9-infra/gpu-type-propagation-20260802

Conversation

@Wingie

@Wingie Wingie commented Aug 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Owner

Summary

  • Re-lands the FlowState-fvks GPU-type propagation work (originally commit e1eceb4 on claude/task/FlowState-fvks) which never made it into main — the branch this repo's HEAD is actually descended from (claude/agentic-rag-browser-use-deps) diverged from it.
  • Gateway keepalive handler (pkg/api/v1/machine.go) was still hardcoding GPUType: "MPS" for every inference-enabled worker. CUDA workers running Ollama were misidentified as MPS nodes, breaking GPU-aware routing.

This PR now also carries PR #4's P0 security fixes into main

PR #4 (security: fix beta9 P0 findings, 95618d0) was merged, but into the
side branch claude/agentic-rag-browser-use-deps — not main. Those
fixes (auth fails-open, unauthenticated control API, SSRF, plaintext creds,
gateway panic) have never actually reached main.

The first version of this PR was a plain cherry-pick of the GPU fix onto
main, bypassing #4 entirely. That was wrong: it would have silently
dropped #4's fix to InferenceRegistry — main still declares
RegisterNode(any) on the local narrow interface in pkg/api/v1/machine.go
while *gateway.ModelRegistry.RegisterNode takes *types.NodeInferenceInfo.
Go requires exact signature matches for interface satisfaction, so on main
today inferenceRegistry.(InferenceRegistry) panics on every gateway
startup that wires in a real *ModelRegistry. #4 already fixed this by
changing the interface to the concrete *types.NodeInferenceInfo signature
plus a non-panicking type-assertion fallback.

Rebuilt this branch as: merge main (PR #3, cubic-review-findings) into #4's
branch, resolve the two overlapping fixes (see below), then reapply the
GPU-type change on top using types.NodeInferenceInfo (matching #4's fix,
not main's still-broken local type). Net effect: merging this PR both
lands GPU-type propagation and resurrects #4's P0 fixes onto main.

Overlapping-fix resolution (both branches independently fixed the same 2 bugs)

Changes (GPU-type propagation)

  • pkg/agent/keepalive.go — GPUType field on InferenceStatus payload
  • pkg/agent/state.go — InferenceGPUType on AgentState/AgentStateSnapshot; UpdateInferenceWithGPU() (keeps UpdateInference() for compat)
  • pkg/agent/agent.go — StartInference() reads GPU type from OllamaManager.GetStatus()
  • pkg/api/v1/machine.go — keepalive handler uses request.Inference.GPUType, falls back to "MPS" only for old agents that omit the field

Test plan

  • No Go toolchain in the environment this PR was authored in — CI must build/vet/test this before merge, especially the merge-conflict resolutions above
  • lint_and_test_go_pkg has been skipped (not run) on every CI run so far on this PR and on main — verify it actually runs and passes before merge, not just shows green from being skipped
  • lint_and_test_python_sdk fails on main itself (874 pre-existing ruff violations across the whole SDK, unrelated to this change) — expect it to keep failing; not this PR's responsibility to fix

Wingie and others added 2 commits May 18, 2026 21:27
@Wingie Claudistrator
deps(sdk): bump rich<15 + websockets<17 for browser-use compat
84ea11c 
Parent repo needs browser-use>=0.11 for shop kit deep-research gap-fill
(agentic_kit_deep_research_task). browser-use 0.11.x requires rich==14.3.1
and websockets>=15.0.1; existing beta9 sdk pins were rich<14 and
websockets<15. Bumped upper bounds to allow rich 14 and websockets 16
— both are minor-version compatible with beta9's console/table API
and websocket-client usage.

No code changes; pin relaxation only.

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
@Wingie Claudistrator
security: fix beta9 P0 findings (auth fails-open, unauth control API,…
95618d0 
… SSRF, plaintext creds, gateway panic) (#4)

* security: fix beta9 P0 findings (auth fails-open, unauth control API, SSRF, plaintext creds, gateway panic)

Applies the P0 security hardening from the FlowState audit
(wip-specs/beta9/2026-07-13-p0-security-fixes.md). Built against fork HEAD
84ea11c; cross-compiles clean for linux/amd64.

1. Inference API auth was fails-open — AuthMiddleware only populated the auth
   context, it never rejected anonymous callers, so /chat/completions,
   /embeddings, /models, /nodes/register, /nodes/:id/heartbeat ran regardless
   of auth. Node/model management now requires TokenTypeClusterAdmin;
   chat/embeddings/listing require any authenticated token.
2. SSRF via unvalidated tailscale_ip — handleRegisterNode and MachineKeepalive
   stored a caller-supplied IP verbatim and used it to build outbound requests.
   Both now reject anything outside the allowed private range.
3. Unauthenticated agent control API — pkg/agent/control.go endpoints hardened.
4. Plaintext credentials in manifests/k3d/beta9.yaml — removed.
5. Gateway panic path fixed.

Files: manifests/k3d/beta9.yaml, pkg/agent/control.go, pkg/api/v1/machine.go,
pkg/api/v1/workspace.go, pkg/gateway/inference_handlers.go

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

* fix(sdk): populate EmbeddingResult.embeddings for batch input (found via CI gate)

The lint_and_test_python_sdk check on PR #4 flagged ruff F401 (unused
ChatMessage/EmbeddingResult imports in test_cubic_fixes.py). Rather than just
delete the imports, running the suite behind that gate surfaced a real bug:

- inference.embed() overloaded the singular `embedding` field with the whole
  batch list and never populated the plural `embeddings` field, so callers of
  result.embeddings got an empty list for batch input. test_batch_embeddings
  was already failing (0 != 3) — hidden because the lint step failed first.

Fixes:
- embed() now sets `embedding` = first vector (legacy accessor) and
  `embeddings` = the full batch, per EmbeddingResult's documented field intent.
- test_batch_embeddings now asserts the return type (uses EmbeddingResult).
- Added test_chat_forwards_chatmessage covering the previously-untested
  ChatMessage -> dict conversion in chat() (uses ChatMessage).

Full SDK suite: 63 passed; ruff clean.

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

---------

Co-authored-by: Claudistrator <savetheplanet@agentosaurus.com>
@Wingie
Wingie force-pushed the claude/beta9-infra/gpu-type-propagation-20260802 branch from baaca51 to 6848153 Compare August 3, 2026 06:11
Wingie and others added 2 commits August 3, 2026 08:15
@Wingie
Merge origin/main into P0 security fixes (#4)
2dbf684 
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.
@Wingie Claudistrator
feat(agent): propagate GPU type through keepalive to gateway
907c90b 
- 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. Applied on top of the merge of
origin/main into PR #4 (95618d0) rather than directly on origin/main, so
this branch keeps PR #4's InferenceRegistry signature fix — main's info :=
&NodeInferenceInfo{} (locally-scoped, plain type) doesn't satisfy
InferenceRegistry.RegisterNode(*types.NodeInferenceInfo), so main still
carries the P0 panic-on-startup bug PR #4 fixed. Uses types.NodeInferenceInfo
here to match.

BD: FlowState-fvks

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
@Wingie
Wingie force-pushed the claude/beta9-infra/gpu-type-propagation-20260802 branch from 6848153 to 907c90b Compare August 3, 2026 06:16
Wingie and others added 2 commits August 8, 2026 10:58
@Wingie Claudistrator
ci: unchain the three test jobs so Go tests actually run
fc887e4 
lint_and_test_go_pkg and verify_protobufs were chained behind
lint_and_test_python_sdk with `needs:`. The python job fails on a pre-existing
~874-violation ruff backlog unrelated to any Go change, so GitHub Actions
skipped both downstream jobs on every push — a silent SKIPPED that reads as a
pass in the checks UI.

Consequence: the Go code PR #5 changes (pkg/agent/keepalive.go, state.go,
agent.go, pkg/api/v1/machine.go) has never once been CI-tested, on that PR or
any prior one.

Removing both `needs:` lines lets all three jobs run independently and report
their real status. The python job will keep failing on the ruff backlog until
that is addressed separately — but it no longer masks the Go and protobuf jobs.

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
@Wingie Claudistrator
fix(sdk): resync generated protobuf clients with genproto verifier
4380035 
verify_protobufs started running for real once lint_and_test_go_pkg
was unchained from the skipped python job (fc887e4); it then failed
on pre-existing drift between the committed generated stubs and the
current betterproto codegen output (one stray blank line before
'if TYPE_CHECKING:' in 14 files) that predates this PR's Go changes
and was never caught because CI never actually ran this check.

Co-Authored-By: Claudistrator <savetheplanet@agentosaurus.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Wingie