Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

pkg/recipe/doc.go (1)

31-141: ⚠️ Potential issue | 🟡 Minor

Fix stale godoc: GET query-parameter os list must include amazonlinux and talos.

Criteria.OS comment includes talos, but the GET query parameter section still documents os as ubuntu, cos, rhel, any—missing amazonlinux and talos.

✅ Suggested godoc diff

-//   - os: ubuntu, cos, rhel, any (default: any)
+//   - os: ubuntu, rhel, cos, amazonlinux, talos, any (default: any)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/recipe/doc.go` around lines 31 - 141, The GET query-parameter
documentation in pkg/recipe/doc.go is stale: the "os" list omits amazonlinux and
talos even though Criteria.OS (e.g., constants like CriteriaOSAmazonLinux,
CriteriaOSTalos or related CriteriaOS* types) supports them; update the "os"
line under "Query Parameters (HTTP API - GET)" to include amazonlinux and talos
(so it reads e.g. "ubuntu, cos, rhel, amazonlinux, talos, any") and ensure any
alias/description for the "os" parameter matches the Criteria.OS constants and
examples elsewhere in the file.

docs/user/cli-reference.md (1)

74-94: ⚠️ Potential issue | 🟠 Major

Add missing aicr snapshot --os flag row to the CLI reference.

The docs’ aicr snapshot flags table (lines ~74–94) does not list the new --os flag, even though recipe query mode does. This creates a user-visible mismatch between CLI behavior and documentation.

✅ Suggested doc update (table row insertion)

 | `--toleration` | | string[] | all taints | Tolerations for agent scheduling (key=value:effect, repeatable). **Default: all taints tolerated** (uses `operator: Exists`). Only specify to restrict which taints are tolerated. |
+| `--os` | | string | | Node OS family to use for collectors (e.g. ubuntu, rhel, cos, amazonlinux, talos) |
 | `--timeout` | | duration | 5m | Timeout for agent Job completion |
 | `--no-cleanup` | | bool | false | Skip removal of Job and RBAC resources on completion. **Warning:** leaves a cluster-admin ClusterRoleBinding active. |

Please confirm the exact --os flag description/default from the Go CLI flag definition so the docs mirror the usage text precisely.

#!/bin/bash
# Description: Find the Go definition of the `aicr snapshot --os` flag so we can mirror its usage/default in docs.

set -euo pipefail

# 1) Locate snapshot command source file(s)
SNAPSHOT_FILES=$(fd -a --type f "snapshot.go" pkg/cli pkg || true)
echo "Found snapshot.go files:"
echo "$SNAPSHOT_FILES"

# 2) Search for `os` flag definition within those files
for f in $SNAPSHOT_FILES; do
  echo "---- Searching $f for --os flag definition ----"
  rg -n --fixed-string 'Name:     "os"' "$f" || true
  rg -n --fixed-string 'Name:    "os"' "$f" || true
  rg -n -- '--os' "$f" || true
  rg -n -- "AICR_OS" "$f" || true
  rg -n -- "StringFlag\\{" "$f" | rg -n "os" || true
done

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@docs/user/cli-reference.md` around lines 74 - 94, The docs are missing the
new aicr snapshot --os flag; locate the Go flag definition in the snapshot
command (search for snapshot.go and the StringFlag/Flag with Name: "os" in the
pkg/cli package or related files), copy the exact usage text and default value
for the --os flag from that Go definition, and add a new row to the Flags table
describing `--os` for the `aicr snapshot` command (include
short/Type/Default/Description exactly as in the Go CLI). Ensure the table row
references the `--os` flag name and mirrors the Go usage/default precisely so
docs and CLI output match.

pkg/snapshotter/snapshot.go (1)

85-100: ⚠️ Potential issue | 🟠 Major

Normalize and validate AICR_OS before factory selection.

parseOSEnv() currently forwards raw env input, so casing/whitespace/typos can silently fall back to non-Talos collectors and make failures hard to diagnose.

♻️ Proposed hardening diff

 import (
 	"context"
 	"log/slog"
 	"os"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
@@
 	"github.com/NVIDIA/aicr/pkg/measurement"
+	"github.com/NVIDIA/aicr/pkg/recipe"
 	"github.com/NVIDIA/aicr/pkg/serializer"
 )
@@
 func parseOSEnv() string {
-	return os.Getenv("AICR_OS")
+	val := strings.ToLower(strings.TrimSpace(os.Getenv("AICR_OS")))
+	if val == "" {
+		return ""
+	}
+	if _, err := recipe.ParseCriteriaOSType(val); err != nil {
+		slog.Warn("invalid AICR_OS value, ignoring", slog.String("value", val))
+		return ""
+	}
+	return val
 }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/snapshotter/snapshot.go` around lines 85 - 100, Normalize and validate
AICR_OS before using it to select collectors: change parseOSEnv() to return a
trimmed, lowercased value (use strings.TrimSpace + strings.ToLower) and in
NodeSnapshotter.measure() validate that the normalized osVal is one of the
supported OS identifiers (e.g. "talos" or any values accepted by
collector.WithOS); if osVal is empty continue with default behavior, but if it's
present and not in the allowed set return an error (or log a clear warning and
fail fast) instead of silently passing the raw value to collector.WithOS,
referencing parseOSEnv, NodeSnapshotter.measure, and collector.WithOS to locate
the change points.

pkg/collector/talos/service.go (1)

100-104: ⚠️ Potential issue | 🟠 Major

Do not mark containerd.service active for non-containerd or unparseable runtime IDs.

At Line 101, containerd.service becomes active for any non-empty runtime string (e.g., cri-o://... or malformed values), which is an incorrect service-state signal.

🔧 Suggested fix

 	runtimeName, runtimeVersion := splitRuntimeID(info.ContainerRuntimeVersion)
-	b.SetString(keyActiveState, "active").
-		SetString(keyRuntimeName, runtimeName).
-		SetString(keyVersion, runtimeVersion)
+	b.SetString(keyRuntimeName, runtimeName)
+	if runtimeName != "containerd" || runtimeVersion == "" {
+		b.SetString(keyActiveState, "unknown")
+		return b.Build()
+	}
+	b.SetString(keyActiveState, "active").
+		SetString(keyVersion, runtimeVersion)

Please also update pkg/collector/talos/service_test.go expectations for the malformed runtime case accordingly.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/collector/talos/service.go` around lines 100 - 104, The code currently
sets keyActiveState to "active" for any non-empty runtime string; change the
logic around splitRuntimeID(info.ContainerRuntimeVersion) so you only
SetString(keyActiveState, "active") when runtimeName == "containerd" and
runtimeVersion is non-empty/parseable, otherwise do not mark containerd active
(set inactive or omit the active state consistently). Update uses of
keyRuntimeName and keyVersion to still record parsed values only when parse
succeeded, and adjust pkg/collector/talos/service_test.go expectations for the
malformed runtime case to assert containerd is not active (and that
runtime/version are empty or reflect the parsing outcome).

pkg/k8s/agent/job_test.go (1)

119-128: ⚠️ Potential issue | 🟡 Minor

Assert AICR_OS presence/absence, not only value.

At Line 120, gotOSEnv == "" passes both when the variable is absent and when it is present-but-empty, so the empty-OS case is not strictly validated.

🔧 Suggested fix

-			gotOSEnv := ""
+			gotOSEnv := ""
+			foundOSEnv := false
 			for _, e := range spec.Containers[0].Env {
 				if e.Name == "AICR_OS" {
 					gotOSEnv = e.Value
+					foundOSEnv = true
+					break
 				}
 			}
+			if tt.wantAICROSEnv == "" && foundOSEnv {
+				t.Errorf("AICR_OS env should be absent for OS=%q", tt.os)
+			}
+			if tt.wantAICROSEnv != "" && !foundOSEnv {
+				t.Errorf("AICR_OS env should be present for OS=%q", tt.os)
+			}
 			if gotOSEnv != tt.wantAICROSEnv {
 				t.Errorf("AICR_OS env = %q, want %q", gotOSEnv, tt.wantAICROSEnv)
 			}

As per coding guidelines, "Explicitly test error conditions and edge cases in test files".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/k8s/agent/job_test.go` around lines 119 - 128, The test currently only
checks the value of AICR_OS via gotOSEnv which conflates "absent" and
"present-but-empty"; change the loop over spec.Containers[0].Env to record both
a boolean foundOSEnv and the value (e.g., foundOSEnv := false; gotOSEnv := ""
and set foundOSEnv = true when e.Name == "AICR_OS"), then replace the single
value check with two assertions: assert presence/absence (foundOSEnv should
equal (tt.wantAICROSEnv != "")) and, only if foundOSEnv is true, assert gotOSEnv
== tt.wantAICROSEnv so you explicitly test existence and value separately.

pkg/collector/factory.go (1)

26-30: 🧹 Nitpick | 🔵 Trivial

Add a guard test to detect Talos OS constant drift across packages.

The "must stay in sync" comment is good, but enforcement is manual. Add a small invariant test in a package that can import all three constants (recipe, collector, agent) so drift fails fast in CI.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/collector/factory.go` around lines 26 - 30, Add a unit test that asserts
the Talos OS constant is identical across packages to catch drift: create a new
test (e.g., invariants_test) that imports the three symbols OSTalos from
pkg/collector (const OSTalos), CriteriaOSTalos from pkg/recipe, and whatever the
Talos constant is in pkg/agent, and assert they are equal (using testing.Fatal
or require/xtest equivalent) so CI fails on mismatch; name the test clearly
(e.g., TestTalosOSConstantInvariant) and keep it lightweight with no external
dependencies.

pkg/collector/talos/os_test.go (1)

163-164: ⚠️ Potential issue | 🟡 Minor

TestOSCollector_GracefulDegradation bypasses the env-resolution path it prepares.

At Line 163, WithNodeName("missing-node") short-circuits NODE_NAME/KUBERNETES_NODE_NAME/HOSTNAME fallback, so this test does not cover the empty-env path.

Proposed fix

-	c := NewOSCollector(WithClientSet(fake.NewSimpleClientset()), WithNodeName("missing-node"))
+	c := NewOSCollector(WithClientSet(fake.NewSimpleClientset()))

As per coding guidelines, "Explicitly test error conditions and edge cases in test files."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/collector/talos/os_test.go` around lines 163 - 164, The test is bypassing
the environment-resolution path by passing WithNodeName("missing-node") so it
never exercises NODE_NAME/KUBERNETES_NODE_NAME/HOSTNAME fallbacks; remove the
WithNodeName(...) option when constructing NewOSCollector (leave
NewOSCollector(WithClientSet(fake.NewSimpleClientset()))), ensure the
environment variables NODE_NAME, KUBERNETES_NODE_NAME and HOSTNAME are unset or
cleared in the test setup, then call c.Collect(ctx) to verify graceful
degradation; reference NewOSCollector, WithClientSet, WithNodeName and Collect
when making the change.

pkg/collector/talos/talos.go (1)

101-101: ⚠️ Potential issue | 🟠 Major

Enforce a 10s timeout at the fetchNode I/O boundary.

fetchNode issues Kubernetes API I/O directly on the incoming ctx at Line 101; this helper should guarantee its own timeout so it cannot block indefinitely when reused without a deadline.

Proposed fix

 import (
 	"context"
 	"log/slog"
 
 	"github.com/NVIDIA/aicr/pkg/collector/k8s"
+	"github.com/NVIDIA/aicr/pkg/defaults"
 	"github.com/NVIDIA/aicr/pkg/errors"
 	k8sclient "github.com/NVIDIA/aicr/pkg/k8s/client"
 	corev1 "k8s.io/api/core/v1"
@@
-	node, err := cs.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+	getCtx, cancel := context.WithTimeout(ctx, defaults.CollectorTimeout)
+	defer cancel()
+	node, err := cs.CoreV1().Nodes().Get(getCtx, nodeName, metav1.GetOptions{})

As per coding guidelines, "All I/O operations and HTTP handlers must use context with explicit timeouts via context.WithTimeout(). Collectors: 10s timeout."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/collector/talos/talos.go` at line 101, The fetchNode helper currently
calls cs.CoreV1().Nodes().Get using the incoming ctx and must enforce a 10s I/O
timeout; wrap the call in a context.WithTimeout(ctx, 10*time.Second), defer
cancel(), and use that timed context when calling cs.CoreV1().Nodes().Get (the
node retrieval in fetchNode) so the function cannot block indefinitely when
invoked without a deadline; ensure you import time if not present and always
call cancel to avoid leaks.

pkg/collector/talos/os.go (1)

134-156: ⚠️ Potential issue | 🟡 Minor

parseOSImage emits misleading ID/VERSION_ID for non‑Talos OSImage formats.

At Line 146+ this parser converts multi-word distro names into invalid structured fields (e.g., RHEL-like strings become ID=red, VERSION_ID=Plow). Prefer Talos-only normalization and leave ID/VERSION_ID unset when format is not confidently mapped.

Proposed fix

 func parseOSImage(image string) (id, version string) {
 	if image == "" {
 		return "", ""
 	}
 	parts := strings.SplitN(image, " (", 2)
-	id = strings.ToLower(strings.TrimSpace(parts[0]))
-	// Strip secondary words from the ID for distros that omit the
-	// parenthesis form (e.g., "Ubuntu 22.04.5 LTS" → "ubuntu").
-	if spaceIdx := strings.Index(id, " "); spaceIdx >= 0 {
-		id = id[:spaceIdx]
-	}
-	if len(parts) == 2 {
+	base := strings.TrimSpace(parts[0])
+	if !strings.EqualFold(base, "talos") {
+		return "", ""
+	}
+	id = "talos"
+	if len(parts) == 2 {
 		v := strings.TrimSuffix(strings.TrimSpace(parts[1]), ")")
 		v = strings.TrimPrefix(v, "v")
 		version = v
 	}
 	return id, version
 }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@pkg/collector/talos/os.go` around lines 134 - 156, parseOSImage currently
force-parses multi-word distro strings into invalid ID/VERSION_ID; change it to
only yield values for confident Talos-style patterns: in parseOSImage check for
a parenthesized version suffix (parts := strings.SplitN(image, " (", 2)) and
only proceed if len(parts) == 2 AND the left-hand part contains no spaces OR the
left-hand part equals "Talos" (case-insensitive) — otherwise return empty
strings; when proceeding, keep the existing lowercasing/trim and version
stripping logic (v prefix trim) so only true Talos-style inputs produce ID and
VERSION_ID.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes