expose Hami MCP#1925
Conversation
Signed-off-by: Tim <tim.wang03@sap.com>
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: haitwang-cloud The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive execution plan (hami-mcp-server.md) to add a Model Context Protocol (MCP) server to HAMi for exposing read-only GPU and scheduler states. The review feedback provides valuable improvements to the plan, including correcting Kubernetes RBAC constraints for ConfigMaps, expanding secret redaction to cover init and ephemeral containers, reusing existing HAMi client initialization utilities, refining static check regexes to avoid false positives in tests, and adjusting the E2E test client constructor to be more idiomatic for Ginkgo suites.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # ServiceAccount: hami-mcp-server | ||
| # ClusterRole verbs (HARD LIMIT): | ||
| # - get/list/watch on: nodes, pods, namespaces, configmaps (hami-scheduler-config only) | ||
| # - NO write verbs anywhere | ||
| # - NO get on secrets |
There was a problem hiding this comment.
In Kubernetes RBAC, you cannot restrict list or watch requests by resourceNames. The API server does not support filtering lists by name at the authorization level, so any attempt to list or watch ConfigMaps with this ClusterRole will result in a 403 Forbidden error. Since the MCP server likely only needs to read the specific hami-scheduler-config ConfigMap, you should only grant the get verb for the specific ConfigMap name under resourceNames or allow listing all ConfigMaps.
| # ServiceAccount: hami-mcp-server | |
| # ClusterRole verbs (HARD LIMIT): | |
| # - get/list/watch on: nodes, pods, namespaces, configmaps (hami-scheduler-config only) | |
| # - NO write verbs anywhere | |
| # - NO get on secrets | |
| # ServiceAccount: hami-mcp-server | |
| # ClusterRole verbs (HARD LIMIT): | |
| # - get/list/watch on: nodes, pods, namespaces | |
| # - get on: configmaps (restricted to hami-scheduler-config via resourceNames) | |
| # - NO write verbs anywhere | |
| # - NO get on secrets |
|
|
||
| ### 3.B Output redaction | ||
| `pkg/mcp/redact/redact.go` strips: | ||
| - All container `env` entries whose name matches `(?i)(token|secret|password|key|cred)` |
There was a problem hiding this comment.
The redaction plan currently only mentions stripping secrets from "container env entries". Kubernetes pods can also run initContainers and ephemeralContainers, which can similarly contain sensitive environment variables. To prevent accidental leaks, the redaction logic in pkg/mcp/redact/redact.go should explicitly process all container types (spec.containers, spec.initContainers, and spec.ephemeralContainers).
| - All container `env` entries whose name matches `(?i)(token|secret|password|key|cred)` | |
| - All container, initContainer, and ephemeralContainer `env` entries whose name matches `(?i)(token|secret|password|key|cred)` |
| pkg/mcp/ | ||
| server.go # Server constructor, tool/resource registration | ||
| client/ | ||
| k8s.go # Wraps client-go: list nodes, pods, configmaps |
There was a problem hiding this comment.
HAMi already has a robust, centralized Kubernetes client initialization package in pkg/util/client/client.go which automatically handles loading kubeconfig from the environment (KUBECONFIG) or falling back to in-cluster config. Instead of creating a new custom client wrapper in pkg/mcp/client/k8s.go, you should reuse pkg/util/client to avoid code duplication and ensure consistent client configuration across the project.
| k8s.go # Wraps client-go: list nodes, pods, configmaps | |
| k8s.go # Uses pkg/util/client to list nodes, pods, configmaps |
| # Forbid write verbs in mcp packages | ||
| if grep -rE '\.(Update|Create|Delete|Patch|Apply)\(' pkg/mcp/ cmd/mcp-server/ 2>/dev/null; then | ||
| echo "ERROR: mcp packages must not call mutating K8s APIs" | ||
| exit 1 | ||
| fi |
There was a problem hiding this comment.
The proposed static check uses a simple grep pattern \.(Update|Create|Delete|Patch|Apply)\( to forbid write operations. This regex is highly prone to false positives. For example, calling a completely unrelated local method or builder pattern like mcpServer.Create(...) or tools.Create(...) will trigger a build failure. Consider refining the regex to target K8s client interfaces specifically, or using a more robust AST-based linter/static analysis tool to enforce read-only API usage.
| # Forbid write verbs in mcp packages | |
| if grep -rE '\.(Update|Create|Delete|Patch|Apply)\(' pkg/mcp/ cmd/mcp-server/ 2>/dev/null; then | |
| echo "ERROR: mcp packages must not call mutating K8s APIs" | |
| exit 1 | |
| fi | |
| # Forbid write verbs in mcp packages (excluding test files to avoid false positives on mocks/fixtures) | |
| if grep -rE --exclude='*_test.go' '\.(Update|Create|Delete|Patch|Apply)\(' pkg/mcp/ cmd/mcp-server/ 2>/dev/null; then | |
| echo "ERROR: mcp packages must not call mutating K8s APIs" | |
| exit 1 | |
| fi |
Add a Model Context Protocol (MCP) server that exposes HAMi GPU scheduling state to AI assistants (Claude Desktop, Claude Code, Cursor, etc.). ## New Components - cmd/mcp-server: Standalone MCP server binary with CLI flags - pkg/mcp/server: MCP server constructor with tool/resource registration - pkg/mcp/client: K8s and Prometheus client wrappers (read-only) - pkg/mcp/tools: 5 MCP tools for GPU queries - pkg/mcp/redact: Secret redaction for sensitive data - pkg/mcp/resources: HAMi config resource ## MCP Tools - list_gpu_nodes: List GPU nodes with vendor/count/memory info - list_gpu_pods: List pods with GPU resource requests - get_quota_usage: Get GPU quota usage per namespace - get_gpu_metrics: Query Prometheus GPU metrics - describe_node: Detailed node info with GPU devices ## Build & Deployment - Makefile: Add mcp-server and docker-mcp targets - docker/Dockerfile.mcp: Multi-stage distroless image - charts/hami/templates/mcp-server/: Helm templates - docs/mcp-server.md: User documentation ## E2E Tests - Full MCP protocol stack tests with in-memory transports - Fake K8s clientset and stubbed Prometheus server - Coverage for all 5 tools, resources, and error handling ## Security - Read-only K8s API access (no write verbs) - RBAC with minimal permissions (no secrets access) - Automatic redaction of sensitive data in responses Signed-off-by: Tim <tim@example.com> Signed-off-by: Tim <tim.wang03@sap.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
📝 WalkthroughWalkthroughThis PR adds a new standalone HAMi MCP (Model Context Protocol) server exposing read-only GPU/Kubernetes/Prometheus data through five tools and a config resource. It includes a Kubernetes client, Prometheus client, JSON redaction utility, server wiring, CLI, Dockerfile, Helm chart, documentation, and extensive unit/e2e tests. ChangesHAMi MCP Server
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant MCPServer
participant K8sClient
participant PrometheusClient
participant RedactUtil
Client->>MCPServer: CallTool(list_gpu_pods, namespace)
MCPServer->>K8sClient: ListGPUPods(ctx, namespace, phase)
K8sClient-->>MCPServer: []Pod
MCPServer->>MCPServer: extractGPUPodInfo(pods)
MCPServer->>RedactUtil: Redact(jsonPayload)
RedactUtil-->>MCPServer: redactedJSON
MCPServer-->>Client: CallToolResult(redactedJSON)
Client->>MCPServer: CallTool(get_gpu_metrics, metric, node)
MCPServer->>PrometheusClient: QueryCustomMetric(metric, node)
PrometheusClient->>PrometheusClient: validateMetricName/validateNodeLabel
PrometheusClient-->>MCPServer: []MetricValue
MCPServer->>RedactUtil: Redact(jsonPayload)
RedactUtil-->>MCPServer: redactedJSON
MCPServer-->>Client: CallToolResult(redactedJSON)
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 15
🧹 Nitpick comments (13)
pkg/mcp/client/prometheus.go (2)
52-55: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueModernize: use
anyinstead ofinterface{}.Static analysis flags line 54 (
Value []interface{}); Go 1.18+ prefers theanyalias.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/client/prometheus.go` around lines 52 - 55, The PrometheusResult struct is using the older interface{} type for Value, which modern Go style prefers to express with any. Update the PrometheusResult definition to replace the Value field’s []interface{} with []any while keeping the JSON tag and the rest of the struct unchanged.Source: Linters/SAST tools
40-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick wingofmt: struct tag alignment.
Static analysis flags this block as not gofmt-formatted (misaligned field tags between
ResultType/Result). Rungofmt -w/goimportson the file.🔧 Proposed fix
Data struct { - ResultType string `json:"resultType"` - Result []PrometheusResult `json:"result"` + ResultType string `json:"resultType"` + Result []PrometheusResult `json:"result"` } `json:"data"`🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/client/prometheus.go` around lines 40 - 49, The PrometheusResponse struct in PrometheusResponse needs gofmt-style alignment for the nested Data fields, since the struct tags/spacing on ResultType and Result are misaligned. Update the field layout in the PrometheusResponse definition so the nested struct formatting matches standard gofmt output, then run gofmt (or goimports) on the file to normalize the entire block.Source: Linters/SAST tools
pkg/mcp/client/k8s.go (2)
113-133: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate GPU resource-name list.
gpuResourceNamesis defined identically in bothhasGPUResources(Lines 114-124) andhasGPURequests(Lines 191-201). Extract to a shared package-level slice to avoid drift if a new vendor is added.♻️ Proposed refactor
+var gpuResourceNames = []string{ + "nvidia.com/gpu", + "nvidia.com/gpumem", + "nvidia.com/gpucores", + "cambricon.com/vmlu", + "hygon.com/dcunum", + "metax-tech.com/sgpu", + "enflame.com/drs-gcu", + "kunlunxin.com/xpu", + "vastaitech.com/va", +} + func hasGPUResources(node *corev1.Node) bool { - gpuResourceNames := []string{ ... } for _, resourceName := range gpuResourceNames { ...Also applies to: 189-214
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/client/k8s.go` around lines 113 - 133, The GPU resource-name list is duplicated across hasGPUResources and hasGPURequests, so extract the shared vendor list into a package-level slice and reuse it from both helpers. Update the logic in k8s.go so hasGPUResources and hasGPURequests reference the same shared symbol, keeping the resource-name set consistent and easier to extend when new vendors are added.
114-124: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExtract the GPU resource-name list into a shared helper
The same list is repeated in both GPU node and pod checks; centralizing it will keep the two paths in sync when a vendor name changes or a new resource is added.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/client/k8s.go` around lines 114 - 124, The GPU resource-name slice is duplicated in the node and pod GPU checks, so changes can drift between the two paths. Extract the list from the k8s client logic into a shared helper or package-level function and have both the node and pod check code use that single source of truth. Use the existing GPU-check code in k8s.go as the reference point so both callers reuse the same helper whenever vendor resource names change.pkg/mcp/redact/redact.go (2)
45-45: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueModernize: use
anyinstead ofinterface{}.Flagged by static analysis at Lines 45 and 65.
Also applies to: 65-65
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/redact/redact.go` at line 45, Replace the legacy empty interface usage in redact.go with the Go 1.18+ alias any. Update the data variable declaration in the relevant redact logic and any other matching interface{} usage flagged in the same flow, keeping the existing behavior unchanged. Use the local symbols around the redact function and any helper that assigns or inspects data to find both occurrences.Source: Linters/SAST tools
42-42: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuegodot: comment missing trailing period.
✏️ Fix
-// - Pod spec.volumes[*].secret +// - Pod spec.volumes[*].secret.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/redact/redact.go` at line 42, The comment in redact.go is missing the required trailing period; update the affected spec.volumes[*].secret comment in the redaction list to end with a period so it matches godot style. Look for the comment near the Pod spec.volumes section and adjust only the comment text.Source: Linters/SAST tools
pkg/mcp/redact/redact_test.go (1)
136-170: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd camelCase regression cases to
TestIsSensitiveKey.Given the regex gap noted in
redact.go(Lines 25-29), add cases likeauthToken,clientSecret,apiTokento lock in the desired behavior once fixed.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/redact/redact_test.go` around lines 136 - 170, Add camelCase regression cases to TestIsSensitiveKey so the sensitivity regex behavior stays locked in after the fix. Update the table in TestIsSensitiveKey to include camelCase inputs such as authToken, clientSecret, and apiToken with expected true, alongside the existing isSensitiveKey cases, so the test covers the gap in redact.go’s matching logic.charts/hami/values.yaml (1)
529-545: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winMisleading default
prometheusUrl.
http://localhost:9090will not resolve to a real Prometheus instance since themcp-servercontainer (seedeployment.yaml) has no Prometheus sidecar. Users who enable the chart without overriding this value will silently get failingget_gpu_metricscalls.Consider defaulting to a cluster-local Prometheus service address (if HAMi ships one) or adding an inline comment flagging that this must be overridden.
💡 Suggested comment
enabled: false image: repository: "hami-mcp" tag: "" pullPolicy: IfNotPresent - prometheusUrl: "http://localhost:9090" + # NOTE: Override with your in-cluster Prometheus service address, + # e.g. http://prometheus-server.monitoring.svc.cluster.local:9090 + prometheusUrl: "http://localhost:9090"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@charts/hami/values.yaml` around lines 529 - 545, The default prometheusUrl in mcpServer is misleading because localhost:9090 will not reach a Prometheus instance from the mcp-server container. Update the mcpServer configuration in values.yaml to use a realistic cluster-local default if available, or add a clear inline note that this value must be overridden; keep the change aligned with the mcp-server deployment setup so users don’t get broken get_gpu_metrics calls by default.charts/hami/templates/mcp-server/deployment.yaml (1)
21-48: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAdd pod/container
securityContexthardening.Given the PR emphasizes minimal privileges and read-only access, the container should also run with a hardened security context (non-root, read-only root filesystem, dropped capabilities, no privilege escalation).
🔒 Suggested hardening
spec: serviceAccountName: {{ include "hami-vgpu.mcp-server" . }} + securityContext: + runAsNonRoot: true + seccompProfile: + type: RuntimeDefault containers: - name: mcp-server image: "{{ .Values.global.imageRegistry }}{{ .Values.mcpServer.image.repository }}:{{ .Values.mcpServer.image.tag | default .Values.global.imageTag }}" imagePullPolicy: {{ .Values.mcpServer.image.pullPolicy }} + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL args:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@charts/hami/templates/mcp-server/deployment.yaml` around lines 21 - 48, The mcp-server deployment is missing container/pod security hardening, so update the deployment template to add a hardened securityContext for the pod and the mcp-server container. In the deployment spec, configure the container to run as non-root, use a read-only root filesystem, drop all capabilities, and disallow privilege escalation; add any needed pod-level settings to support this safely. Use the existing mcp-server container block and the surrounding spec/template structure to place the new securityContext fields consistently.pkg/mcp/tools/list_gpu_nodes.go (2)
157-167: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
errorResultis a shared cross-tool helper tucked into this specific tool file.It's also used by
get_quota_usage.go. Consider moving it to a shared file (e.g.,tools/common.go) for discoverability, since other tool files implicitly depend on this one for it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/tools/list_gpu_nodes.go` around lines 157 - 167, The shared errorResult helper is defined in list_gpu_nodes.go but is also depended on by get_quota_usage.go, making the dependency hidden and hard to discover. Move errorResult into a shared tools-level file such as common.go (or another clearly shared location), keep its signature unchanged, and update both list_gpu_nodes and get_quota_usage to use the shared helper so cross-tool usage is explicit and easy to find.
120-136: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNon-deterministic vendor selection on multi-vendor annotations.
Iterating a Go
maphas randomized order; if a node ever reports capacity for more than one vendor resource name, the detected vendor/count could vary between calls.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/tools/list_gpu_nodes.go` around lines 120 - 136, The GPU vendor detection in list_gpu_nodes.go is non-deterministic because the range over gpuResourceNames uses a Go map, so nodes with multiple vendor capacities may return different results between calls. Update the vendor selection logic around gpuResourceNames and the loop in the GPU capacity check to use a fixed priority order (for example, an ordered slice of resource/vendor pairs) and keep the existing first-match behavior deterministic.pkg/mcp/tools/list_gpu_pods.go (1)
131-149: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHardcoded vendor GPU resource-name list risks drifting from the canonical list.
This 7-entry
gpuResourceNamesslice duplicates resource-name knowledge that's likely already centralized (e.g. inpkg/mcp/client/k8s.go'shasGPURequests, referenced by the upstreamListGPUPodscontract, or inpkg/device/pkg/util). If the two lists diverge,ListGPUPodscould return pods that this function then reports withRequestedGPU: 0, silently confusing MCP clients.Consider extracting a single shared list/constant that both the filter and the extractor consume.
#!/bin/bash # Compare resource-name lists used by hasGPURequests vs extractGPUPodInfo. rg -n --type=go -A15 'func hasGPURequests' pkg/mcp/client/k8s.go🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/tools/list_gpu_pods.go` around lines 131 - 149, The hardcoded gpuResourceNames slice in extractGPUPodInfo is duplicating GPU resource-name knowledge and can drift from the canonical list used by hasGPURequests/ListGPUPods. Move the vendor resource names into a shared constant or helper in a common package (for example alongside hasGPURequests in pkg/mcp/client/k8s.go or a shared util/package), then make both the pod filter and extractGPUPodInfo consume that same source so RequestedGPU stays consistent with what ListGPUPods returns.pkg/mcp/tools/describe_node.go (1)
173-173: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueUse
strings.SplitSeqhere. Go 1.26.2 already supports it, so this loop can avoid allocating the split slice:for part := range strings.SplitSeq(deviceStr, ":") {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/mcp/tools/describe_node.go` at line 173, The loop in describe_node.go still uses strings.Split, which allocates a slice unnecessarily; update the deviceStr iteration in the relevant parsing block to use strings.SplitSeq instead. Locate the logic that iterates over deviceStr parts and switch the range form to strings.SplitSeq(deviceStr, ":") so it streams parts without the extra allocation.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pkg/mcp/redact/redact_test.go`:
- Around line 105-116: The test in redact_test.go uses unchecked type assertions
on elements from parsed["imagePullSecrets"], which can panic and violates
forcetypeassert. Update the assertions in the redaction test to use comma-ok
checks for each arr element before treating it as map[string]interface{}, and
fail the test with a clear message if the shape is unexpected. Keep the fix
localized to the parsed array handling around arr, first, and second in the
test.
In `@pkg/mcp/redact/redact.go`:
- Around line 25-29: The sensitive key detection in isSensitiveKey is missing
camelCase names, so redaction can be bypassed by keys like authToken,
clientSecret, apiToken, and accessToken. Update the matching logic around
sensitivePattern in redact.go to normalize or split camelCase before applying
the existing sensitive-word check, and keep the current separator-based matching
so both snake_case and camelCase secrets are caught.
In `@pkg/mcp/resources/config.go`:
- Around line 64-77: The config serialization in `configMap.Data` only redacts
after `json.Marshal`, so secrets inside raw YAML/text blob values can still
leak. Update the logic around `configData`, `json.Marshal`, and `redact.Redact`
in this config path to detect string/blob values before marshaling and redact
sensitive content within those raw fields as well, not just JSON objects/arrays.
Keep the existing `redact.Redact` call for structured data, but ensure any raw
config blobs are sanitized first or separately before being included in
`redactedData`.
- Around line 58-62: The ConfigResource logic is hardcoding the HAMi ConfigMap
namespace and name instead of using the chart values, so update the
ConfigResource path to derive both from the rendered chart configuration. In
config.go, thread the namespace and ConfigMap name into the resource lookup used
by the method that calls r.k8sClient.GetConfigMap, rather than always using
hami-system and hami-scheduler-config. Make the ConfigMap reference come from
the same values that render the chart so namespaceOverride and chart-generated
names are handled correctly.
In `@pkg/mcp/server_test.go`:
- Around line 95-96: The test setup in server_test.go is using an explicit
context.WithCancel on context.Background(), which triggers the modernize lint
warning. Update the test code in the relevant test function to use t.Context()
instead of creating a manual cancelable context, and remove the now-unneeded
cancel handling while keeping the existing context variable usage aligned with
the test name and MCP server test helpers.
In `@pkg/mcp/server.go`:
- Around line 41-46: The Server struct field alignment is inconsistent and is
triggering the gofmt lint failure. Update the field formatting in Server so the
struct fields are aligned consistently with standard Go formatting, keeping the
existing field names and types in place. Use the Server type definition as the
target for the gofmt cleanup.
- Around line 169-184: The HTTP server in RunHTTP only sets ReadHeaderTimeout,
so add the missing server timeouts to the http.Server configuration. Update the
server setup in Server.RunHTTP to include ReadTimeout, WriteTimeout, and
IdleTimeout alongside ReadHeaderTimeout, using sensible duration values
consistent with the existing timeout style in this codebase. Keep the change
localized to the srv initialization so the /mcp and /healthz handlers continue
to work unchanged.
In `@pkg/mcp/tools/describe_node.go`:
- Around line 156-165: The extractGPUDevices method is only reading the
fixture-specific hami.io/node-devices-to-register annotation, so gpuDevices
stays empty for real HAMi nodes. Update DescribeNodeTool.extractGPUDevices to
check the actual device-plugin registration annotations used in production, such
as hami.io/node-nvidia-register, hami.io/node-dcu-register,
hami.io/node-register-xpu, and hami.io/node-register-Ascend910B, or delegate to
the device-specific decoders before falling back to nil. Keep the parsing logic
compatible with parseDeviceString and preserve the current empty/nil behavior
when no supported registration annotation is present.
In `@pkg/mcp/tools/get_gpu_metrics.go`:
- Around line 96-103: The GPUMetric timestamp serialization in
get_gpu_metrics.go is formatting a local time with a hardcoded UTC suffix.
Update the GPUMetric construction in the results loop to convert r.Time to UTC
before formatting, and use a timezone-aware format such as time.RFC3339 so the
serialized timestamp in GPUMetric.Time matches its actual timezone.
In `@pkg/mcp/tools/get_quota_usage.go`:
- Around line 63-71: The quota fields in QuotaUsage are currently serialized as
0, which can be mistaken for a real quota value instead of “unknown.” Update the
get_quota_usage tool response model so GPUMemoryQuotaGiB and GPUCoreQuota are
nullable or omitted when not available, and adjust the code that builds this
response to leave them unset until ResourceQuota/namespace annotation lookup
exists. If you keep the current shape for now, update the tool Description in
get_quota_usage.go to clearly state that quota limits are not yet tracked and 0
does not mean a real limit.
In `@pkg/mcp/tools/list_gpu_nodes.go`:
- Around line 62-70: The GPUNodeInfo.AllocatedMemoryGB field is never populated,
so extractGPUNodeInfo and the related response path always return misleading
zeroes. Either populate AllocatedMemoryGB in extractGPUNodeInfo by computing
node GPU memory usage from the pods/requests already being inspected, or remove
AllocatedMemoryGB from GPUNodeInfo and any JSON output until that calculation
exists. Keep the fix aligned with GPUNodeInfo and extractGPUNodeInfo so the MCP
response only exposes supported values.
- Around line 138-152: The GPU memory and core values are being read from the
wrong node metadata source in the list GPU nodes flow. Update the extraction
logic in the helper that builds the node info to read `nvidia.com/gpu.memory`
and `nvidia.com/gpu.cores` from `node.Labels` instead of `node.Annotations`, and
make sure the test fixture reflects labels so the `ListGpuNodes`-related parsing
continues to work on real GPU nodes.
In `@pkg/mcp/tools/list_gpu_pods.go`:
- Around line 152-156: The allocation parsing in list_gpu_pods.go is using the
wrong annotations and payload format, so AllocatedDeviceUUIDs never reflects
real HAMi pods. Update the logic in the pod annotation handling to read
hami.io/vgpu-devices-allocated in the list GPU pods flow, and decode the HAMi
UUID,Type,mem,cores:;... entries to extract only the UUID values. Remove the
fallback to nvidia.com/gpu-devices-to-use and hami.io/gpu-devices-to-use in this
path, and keep the fix localized around the annotation parsing branch that
populates info.AllocatedDeviceUUIDs.
In `@test/e2e/mcp/mcp_e2e_test.go`:
- Around line 1-374: The MCP end-to-end suite is written with plain testing.T
assertions, but the project guideline requires Ginkgo/Gomega for e2e tests.
Refactor the tests in mcp_e2e_test.go, including setupFixture, callTool, and the
TestMCPServer_* cases, into Ginkgo specs using Describe/Context/It and Gomega
expectations, while keeping the same MCP client/server behavior and assertions.
- Line 147: The test setup is using context.WithCancel(context.Background())
instead of the newer t.Context() pattern, matching the modernization issue seen
in the MCP tests. Update the context creation in the relevant e2e test code to
derive from t.Context() and keep the cancel handling aligned with the
surrounding test lifecycle; use the existing test context in the mcp_e2e_test
helper or test body where ctx is created.
---
Nitpick comments:
In `@charts/hami/templates/mcp-server/deployment.yaml`:
- Around line 21-48: The mcp-server deployment is missing container/pod security
hardening, so update the deployment template to add a hardened securityContext
for the pod and the mcp-server container. In the deployment spec, configure the
container to run as non-root, use a read-only root filesystem, drop all
capabilities, and disallow privilege escalation; add any needed pod-level
settings to support this safely. Use the existing mcp-server container block and
the surrounding spec/template structure to place the new securityContext fields
consistently.
In `@charts/hami/values.yaml`:
- Around line 529-545: The default prometheusUrl in mcpServer is misleading
because localhost:9090 will not reach a Prometheus instance from the mcp-server
container. Update the mcpServer configuration in values.yaml to use a realistic
cluster-local default if available, or add a clear inline note that this value
must be overridden; keep the change aligned with the mcp-server deployment setup
so users don’t get broken get_gpu_metrics calls by default.
In `@pkg/mcp/client/k8s.go`:
- Around line 113-133: The GPU resource-name list is duplicated across
hasGPUResources and hasGPURequests, so extract the shared vendor list into a
package-level slice and reuse it from both helpers. Update the logic in k8s.go
so hasGPUResources and hasGPURequests reference the same shared symbol, keeping
the resource-name set consistent and easier to extend when new vendors are
added.
- Around line 114-124: The GPU resource-name slice is duplicated in the node and
pod GPU checks, so changes can drift between the two paths. Extract the list
from the k8s client logic into a shared helper or package-level function and
have both the node and pod check code use that single source of truth. Use the
existing GPU-check code in k8s.go as the reference point so both callers reuse
the same helper whenever vendor resource names change.
In `@pkg/mcp/client/prometheus.go`:
- Around line 52-55: The PrometheusResult struct is using the older interface{}
type for Value, which modern Go style prefers to express with any. Update the
PrometheusResult definition to replace the Value field’s []interface{} with
[]any while keeping the JSON tag and the rest of the struct unchanged.
- Around line 40-49: The PrometheusResponse struct in PrometheusResponse needs
gofmt-style alignment for the nested Data fields, since the struct tags/spacing
on ResultType and Result are misaligned. Update the field layout in the
PrometheusResponse definition so the nested struct formatting matches standard
gofmt output, then run gofmt (or goimports) on the file to normalize the entire
block.
In `@pkg/mcp/redact/redact_test.go`:
- Around line 136-170: Add camelCase regression cases to TestIsSensitiveKey so
the sensitivity regex behavior stays locked in after the fix. Update the table
in TestIsSensitiveKey to include camelCase inputs such as authToken,
clientSecret, and apiToken with expected true, alongside the existing
isSensitiveKey cases, so the test covers the gap in redact.go’s matching logic.
In `@pkg/mcp/redact/redact.go`:
- Line 45: Replace the legacy empty interface usage in redact.go with the Go
1.18+ alias any. Update the data variable declaration in the relevant redact
logic and any other matching interface{} usage flagged in the same flow, keeping
the existing behavior unchanged. Use the local symbols around the redact
function and any helper that assigns or inspects data to find both occurrences.
- Line 42: The comment in redact.go is missing the required trailing period;
update the affected spec.volumes[*].secret comment in the redaction list to end
with a period so it matches godot style. Look for the comment near the Pod
spec.volumes section and adjust only the comment text.
In `@pkg/mcp/tools/describe_node.go`:
- Line 173: The loop in describe_node.go still uses strings.Split, which
allocates a slice unnecessarily; update the deviceStr iteration in the relevant
parsing block to use strings.SplitSeq instead. Locate the logic that iterates
over deviceStr parts and switch the range form to strings.SplitSeq(deviceStr,
":") so it streams parts without the extra allocation.
In `@pkg/mcp/tools/list_gpu_nodes.go`:
- Around line 157-167: The shared errorResult helper is defined in
list_gpu_nodes.go but is also depended on by get_quota_usage.go, making the
dependency hidden and hard to discover. Move errorResult into a shared
tools-level file such as common.go (or another clearly shared location), keep
its signature unchanged, and update both list_gpu_nodes and get_quota_usage to
use the shared helper so cross-tool usage is explicit and easy to find.
- Around line 120-136: The GPU vendor detection in list_gpu_nodes.go is
non-deterministic because the range over gpuResourceNames uses a Go map, so
nodes with multiple vendor capacities may return different results between
calls. Update the vendor selection logic around gpuResourceNames and the loop in
the GPU capacity check to use a fixed priority order (for example, an ordered
slice of resource/vendor pairs) and keep the existing first-match behavior
deterministic.
In `@pkg/mcp/tools/list_gpu_pods.go`:
- Around line 131-149: The hardcoded gpuResourceNames slice in extractGPUPodInfo
is duplicating GPU resource-name knowledge and can drift from the canonical list
used by hasGPURequests/ListGPUPods. Move the vendor resource names into a shared
constant or helper in a common package (for example alongside hasGPURequests in
pkg/mcp/client/k8s.go or a shared util/package), then make both the pod filter
and extractGPUPodInfo consume that same source so RequestedGPU stays consistent
with what ListGPUPods returns.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6a2c1e4b-0f4a-4c69-bee1-04edd80ff597
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (35)
Makefilecharts/hami/templates/_helpers.tplcharts/hami/templates/mcp-server/clusterrole.yamlcharts/hami/templates/mcp-server/clusterrolebinding.yamlcharts/hami/templates/mcp-server/deployment.yamlcharts/hami/templates/mcp-server/service.yamlcharts/hami/templates/mcp-server/serviceaccount.yamlcharts/hami/values.yamlcmd/mcp-server/flags.gocmd/mcp-server/main.godocker/Dockerfile.mcpdocs/mcp-server.mdgo.modhami-mcp-server.mdpkg/mcp/client/k8s.gopkg/mcp/client/k8s_test.gopkg/mcp/client/prometheus.gopkg/mcp/client/prometheus_test.gopkg/mcp/redact/redact.gopkg/mcp/redact/redact_test.gopkg/mcp/resources/config.gopkg/mcp/resources/config_test.gopkg/mcp/server.gopkg/mcp/server_test.gopkg/mcp/tools/describe_node.gopkg/mcp/tools/describe_node_test.gopkg/mcp/tools/get_gpu_metrics.gopkg/mcp/tools/get_gpu_metrics_test.gopkg/mcp/tools/get_quota_usage.gopkg/mcp/tools/get_quota_usage_test.gopkg/mcp/tools/list_gpu_nodes.gopkg/mcp/tools/list_gpu_nodes_test.gopkg/mcp/tools/list_gpu_pods.gopkg/mcp/tools/list_gpu_pods_test.gotest/e2e/mcp/mcp_e2e_test.go
| arr, ok := parsed["imagePullSecrets"].([]interface{}) | ||
| if !ok || len(arr) != 2 { | ||
| t.Fatalf("expected array of length 2, got %T %v", parsed["imagePullSecrets"], parsed["imagePullSecrets"]) | ||
| } | ||
| first := arr[0].(map[string]interface{}) | ||
| if first["name"] != "[REDACTED]" { | ||
| t.Errorf("expected name to be redacted, got %v", first["name"]) | ||
| } | ||
| second := arr[1].(map[string]interface{}) | ||
| if second["name"] != "[REDACTED]" { | ||
| t.Errorf("expected second name to be redacted, got %v", second["name"]) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Unchecked type assertions failing CI lint (forcetypeassert).
Pipeline failure and static analysis both flag Lines 109 and 113: arr[0].(map[string]interface{}) / arr[1].(map[string]interface{}) are unchecked. Use the comma-ok form to satisfy golangci-lint and avoid a test panic with an unhelpful message if the shape changes.
✅ Proposed fix
- first := arr[0].(map[string]interface{})
+ first, ok := arr[0].(map[string]interface{})
+ if !ok {
+ t.Fatalf("expected map, got %T", arr[0])
+ }
if first["name"] != "[REDACTED]" {
t.Errorf("expected name to be redacted, got %v", first["name"])
}
- second := arr[1].(map[string]interface{})
+ second, ok := arr[1].(map[string]interface{})
+ if !ok {
+ t.Fatalf("expected map, got %T", arr[1])
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| arr, ok := parsed["imagePullSecrets"].([]interface{}) | |
| if !ok || len(arr) != 2 { | |
| t.Fatalf("expected array of length 2, got %T %v", parsed["imagePullSecrets"], parsed["imagePullSecrets"]) | |
| } | |
| first := arr[0].(map[string]interface{}) | |
| if first["name"] != "[REDACTED]" { | |
| t.Errorf("expected name to be redacted, got %v", first["name"]) | |
| } | |
| second := arr[1].(map[string]interface{}) | |
| if second["name"] != "[REDACTED]" { | |
| t.Errorf("expected second name to be redacted, got %v", second["name"]) | |
| } | |
| arr, ok := parsed["imagePullSecrets"].([]interface{}) | |
| if !ok || len(arr) != 2 { | |
| t.Fatalf("expected array of length 2, got %T %v", parsed["imagePullSecrets"], parsed["imagePullSecrets"]) | |
| } | |
| first, ok := arr[0].(map[string]interface{}) | |
| if !ok { | |
| t.Fatalf("expected map, got %T", arr[0]) | |
| } | |
| if first["name"] != "[REDACTED]" { | |
| t.Errorf("expected name to be redacted, got %v", first["name"]) | |
| } | |
| second, ok := arr[1].(map[string]interface{}) | |
| if !ok { | |
| t.Fatalf("expected map, got %T", arr[1]) | |
| } | |
| if second["name"] != "[REDACTED]" { | |
| t.Errorf("expected second name to be redacted, got %v", second["name"]) | |
| } |
🧰 Tools
🪛 GitHub Actions: CI Workflow / 5_lint.txt
[error] 109-109: golangci-lint: type assertion must be checked (forcetypeassert)
🪛 GitHub Check: lint
[failure] 113-113:
type assertion must be checked (forcetypeassert)
[failure] 109-109:
type assertion must be checked (forcetypeassert)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/redact/redact_test.go` around lines 105 - 116, The test in
redact_test.go uses unchecked type assertions on elements from
parsed["imagePullSecrets"], which can panic and violates forcetypeassert. Update
the assertions in the redaction test to use comma-ok checks for each arr element
before treating it as map[string]interface{}, and fail the test with a clear
message if the shape is unexpected. Keep the fix localized to the parsed array
handling around arr, first, and second in the test.
Sources: Linters/SAST tools, Pipeline failures
| // sensitivePattern matches environment variable names and annotation keys that | ||
| // likely carry secrets. We use word-boundary anchors to avoid false positives | ||
| // like "apiVersion" / "keyName" / "credentialPolicy" / "authority" being | ||
| // matched purely on substring. | ||
| var sensitivePattern = regexp.MustCompile(`(?i)(?:^|[._\-/])(token|secret|password|passwd|credential|cred|auth|api[_\-]?key|key)(?:$|[._\-/])`) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Go regexp RE2 lookahead support for word boundary matching
💡 Result:
The Go regexp package, which is based on the RE2 syntax, does not support lookahead or lookbehind assertions [1][2][3]. As a design principle, RE2 and Go's regexp package omit features that require backtracking, such as lookaround and backreferences, to ensure linear-time performance guarantees relative to the input size [2][4]. Regarding word boundary matching: - Go's regexp package does support word boundary matching using the \b and \B tokens [5][6]. - These tokens are strictly limited to ASCII word boundaries [7][8]. They define a word boundary as the position between a word character (\w) and a non-word character (\W), the beginning of the text (\A), or the end of the text (\z) [5]. - The word character class \w is defined strictly as [0-9A-Za-z_] [5]. It does not support Unicode word characters [7][8]. If your application requires lookahead or advanced Unicode word boundary support, it is recommended to handle that logic at the application level—for example, by performing multiple passes with regex or using standard string manipulation—rather than relying on complex regex assertions [4].
Citations:
- 1: https://pkg.go.dev/regexp
- 2: https://github.com/google/re2/
- 3: https://docs.cloud.google.com/load-balancing/docs/re2-support
- 4: Support for negative lookahead google/re2#365
- 5: https://go.dev/src/regexp/syntax/doc.go?m=text
- 6: https://go.dev/src/regexp/syntax/regexp.go
- 7: regexp/syntax: document that \b and \B are ASCII-only golang/go#5896
- 8: [regexp feature request] Support 'Ş' character on word boundary search, like in .NET/Python3/Java regex engine. golang/go#58863
🏁 Script executed:
#!/bin/bash
set -euo pipefail
git ls-files 'pkg/mcp/redact/*' 'pkg/mcp/**/*test*' | sed -n '1,200p'
printf '\n--- redact.go ---\n'
cat -n pkg/mcp/redact/redact.go | sed -n '1,220p'
printf '\n--- references ---\n'
rg -n "isSensitiveKey|sensitivePattern|splitCamelCase|camelBoundaryRE|nonSensitiveKeys" pkg/mcp -SRepository: Project-HAMi/HAMi
Length of output: 8946
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- isSensitiveKey implementation ---\n'
sed -n '232,255p' pkg/mcp/redact/redact.go | cat -n
printf '\n--- tests around isSensitiveKey ---\n'
sed -n '140,210p' pkg/mcp/redact/redact_test.go | cat -n
printf '\n--- behavioral probe ---\n'
python3 - <<'PY'
import re
pattern = re.compile(r'(?i)(?:^|[._\-/])(token|secret|password|passwd|credential|cred|auth|api[_\-]?key|key)(?:$|[._\-/])')
samples = [
"authToken",
"clientSecret",
"apiToken",
"accessToken",
"token",
"my.token",
"my-secret",
"apiKey",
"api-key",
"credentialPolicy",
"authority",
]
for s in samples:
print(f"{s:18} -> {bool(pattern.search(s.lower()))}")
PYRepository: Project-HAMi/HAMi
Length of output: 2111
Redaction misses camelCase secret keys isSensitiveKey only matches separators around sensitive words, so keys like authToken, clientSecret, apiToken, and accessToken bypass redaction and can leak through MCP output. Normalize camelCase before matching.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/redact/redact.go` around lines 25 - 29, The sensitive key detection
in isSensitiveKey is missing camelCase names, so redaction can be bypassed by
keys like authToken, clientSecret, apiToken, and accessToken. Update the
matching logic around sensitivePattern in redact.go to normalize or split
camelCase before applying the existing sensitive-word check, and keep the
current separator-based matching so both snake_case and camelCase secrets are
caught.
| // Get the HAMi scheduler config ConfigMap | ||
| configMap, err := r.k8sClient.GetConfigMap(ctx, "hami-system", "hami-scheduler-config") | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to get HAMi config: %w", err) | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== config.go around target lines ==\n'
sed -n '1,180p' pkg/mcp/resources/config.go
printf '\n== search for hardcoded configmap refs ==\n'
rg -n 'hami-system|hami-scheduler-config|GetConfigMap\(' .
printf '\n== search for config/flags plumbing around scheduler config ==\n'
rg -n 'scheduler-config|namespace|config map|ConfigMap' pkg cmd internal api .Repository: Project-HAMi/HAMi
Length of output: 50373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== MCP server deployment ==\n'
sed -n '1,220p' charts/hami/templates/mcp-server/deployment.yaml
printf '\n== scheduler ConfigMap template ==\n'
sed -n '1,120p' charts/hami/templates/scheduler/configmap.yaml
printf '\n== ConfigResource constructor usages ==\n'
rg -n 'NewConfigResource\(' .Repository: Project-HAMi/HAMi
Length of output: 4876
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== charts/hami/templates/mcp-server/deployment.yaml =='
sed -n '1,220p' charts/hami/templates/mcp-server/deployment.yaml
echo '== charts/hami/templates/scheduler/configmap.yaml =='
sed -n '1,120p' charts/hami/templates/scheduler/configmap.yaml
echo '== NewConfigResource usages =='
rg -n 'NewConfigResource\(' .Repository: Project-HAMi/HAMi
Length of output: 4909
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== helpers for scheduler naming =='
sed -n '1,120p' charts/hami/templates/_helpers.tpl
echo '== scheduler-related values =='
sed -n '320,390p' charts/hami/values.yaml
echo '== search for configmap name overrides =='
rg -n 'configmap.*name|scheduler.*name|hami-scheduler-config|existing ConfigMap|namespaceOverride' charts/hami/values.yaml charts/hami/templates -g '!**/dist/**'Repository: Project-HAMi/HAMi
Length of output: 8779
Derive the ConfigMap reference from chart values.
pkg/mcp/resources/config.go always reads hami-system/hami-scheduler-config, but the chart renders the scheduler ConfigMap into hami-vgpu.namespace with a chart-generated name. That makes this resource miss the ConfigMap in normal installs and any namespaceOverride setup. Thread both namespace and name through ConfigResource instead of hardcoding them.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/resources/config.go` around lines 58 - 62, The ConfigResource logic
is hardcoding the HAMi ConfigMap namespace and name instead of using the chart
values, so update the ConfigResource path to derive both from the rendered chart
configuration. In config.go, thread the namespace and ConfigMap name into the
resource lookup used by the method that calls r.k8sClient.GetConfigMap, rather
than always using hami-system and hami-scheduler-config. Make the ConfigMap
reference come from the same values that render the chart so namespaceOverride
and chart-generated names are handled correctly.
| // Extract config data | ||
| configData := make(map[string]interface{}) | ||
| for key, value := range configMap.Data { | ||
| configData[key] = value | ||
| } | ||
|
|
||
| // Marshal to JSON | ||
| data, err := json.Marshal(configData) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to marshal config: %w", err) | ||
| } | ||
|
|
||
| // Apply redaction | ||
| redactedData := redact.Redact(string(data)) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect redact.Redact implementation to confirm whether it scans opaque string values for secret patterns.
fd -a redact.go
cat -n pkg/mcp/redact/redact.go 2>/dev/null | sed -n '1,200p'Repository: Project-HAMi/HAMi
Length of output: 6905
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== pkg/mcp/resources/config.go =="
cat -n pkg/mcp/resources/config.go | sed -n '1,180p'
echo
echo "== redact-related call sites =="
rg -n "redact\.Redact|Redact\(" pkg/mcp -g '!**/*_test.go'Repository: Project-HAMi/HAMi
Length of output: 3932
Redact raw config blobs too. configMap.Data values are emitted here as JSON string fields, but redact.Redact() only traverses JSON objects/arrays. Secrets embedded inside YAML/text blobs will remain visible.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/resources/config.go` around lines 64 - 77, The config serialization
in `configMap.Data` only redacts after `json.Marshal`, so secrets inside raw
YAML/text blob values can still leak. Update the logic around `configData`,
`json.Marshal`, and `redact.Redact` in this config path to detect string/blob
values before marshaling and redact sensitive content within those raw fields as
well, not just JSON objects/arrays. Keep the existing `redact.Redact` call for
structured data, but ensure any raw config blobs are sanitized first or
separately before being included in `redactedData`.
| ctx, cancel := context.WithCancel(context.Background()) | ||
| defer cancel() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use t.Context() (modernize lint failure).
🔧 Fix
- ctx, cancel := context.WithCancel(context.Background())
+ ctx, cancel := context.WithCancel(t.Context())
defer cancel()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| ctx, cancel := context.WithCancel(context.Background()) | |
| defer cancel() | |
| ctx, cancel := context.WithCancel(t.Context()) | |
| defer cancel() |
🧰 Tools
🪛 GitHub Check: lint
[failure] 95-95:
testingcontext: context.WithCancel can be modernized using t.Context (modernize)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/server_test.go` around lines 95 - 96, The test setup in
server_test.go is using an explicit context.WithCancel on context.Background(),
which triggers the modernize lint warning. Update the test code in the relevant
test function to use t.Context() instead of creating a manual cancelable
context, and remove the now-unneeded cancel handling while keeping the existing
context variable usage aligned with the test name and MCP server test helpers.
Source: Linters/SAST tools
| // GPUNodeInfo represents GPU information for a node. | ||
| type GPUNodeInfo struct { | ||
| Name string `json:"name"` | ||
| GPUVendor string `json:"gpuVendor"` | ||
| GPUCount int `json:"gpuCount"` | ||
| AllocatedMemoryGB float64 `json:"allocatedMemoryGB"` | ||
| TotalMemoryGB float64 `json:"totalMemoryGB"` | ||
| AllocatedCoresPct float64 `json:"allocatedCoresPct"` | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
AllocatedMemoryGB is declared but never populated.
The field is always 0 in every response regardless of node state, since nothing in extractGPUNodeInfo sets it. Either implement the calculation (e.g., sum GPU memory requests of pods on the node) or drop the field until implemented, to avoid presenting misleading zero values to MCP consumers.
🔧 Minimal fix: remove unpopulated field until implemented
type GPUNodeInfo struct {
Name string `json:"name"`
GPUVendor string `json:"gpuVendor"`
GPUCount int `json:"gpuCount"`
- AllocatedMemoryGB float64 `json:"allocatedMemoryGB"`
TotalMemoryGB float64 `json:"totalMemoryGB"`
AllocatedCoresPct float64 `json:"allocatedCoresPct"`
}Also applies to: 114-155
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/tools/list_gpu_nodes.go` around lines 62 - 70, The
GPUNodeInfo.AllocatedMemoryGB field is never populated, so extractGPUNodeInfo
and the related response path always return misleading zeroes. Either populate
AllocatedMemoryGB in extractGPUNodeInfo by computing node GPU memory usage from
the pods/requests already being inspected, or remove AllocatedMemoryGB from
GPUNodeInfo and any JSON output until that calculation exists. Keep the fix
aligned with GPUNodeInfo and extractGPUNodeInfo so the MCP response only exposes
supported values.
| // Extract memory information from annotations | ||
| if memStr, ok := node.Annotations["nvidia.com/gpu.memory"]; ok { | ||
| var totalMem float64 | ||
| if _, err := fmt.Sscanf(memStr, "%f", &totalMem); err == nil { | ||
| info.TotalMemoryGB = totalMem / 1024 // Convert MB to GB | ||
| } | ||
| } | ||
|
|
||
| // Extract core information from annotations | ||
| if coreStr, ok := node.Annotations["nvidia.com/gpu.cores"]; ok { | ||
| var totalCores float64 | ||
| if _, err := fmt.Sscanf(coreStr, "%f", &totalCores); err == nil { | ||
| info.AllocatedCoresPct = totalCores | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
HAMi util.DecodeNodeDevices EncodeNodeDevices format hami.io/node-nvidia-register
💡 Result:
In the HAMi project, hami.io/node-nvidia-register is a Kubernetes node annotation used by the device plugin to report detailed GPU specifications—such as UUID, VRAM, compute limit, and health status—to the scheduler [1][2]. Because standard Kubernetes device plugins only support integer counts, HAMi uses this annotation to provide the extended metadata required for its scheduling and virtualization logic [2]. The EncodeNodeDevices and DecodeNodeDevices utility functions manage the serialization and deserialization of these device records [3][4]. Data Format The annotation stores a colon-separated string where each segment represents a device [1][5]. The internal format of a device record is comma-separated [1][2]: {Device UUID},{device split count},{device memory limit},{device core limit},{device type},{device numa},{healthy},{index},{mode} [6][3] Key fields include: - Device UUID: Unique identifier for the GPU [2]. - Count: The inflated logical device count [2]. - Devmem: VRAM limit in MiB [2]. - Devcore: Compute limit as a percentage [2]. - Type: The GPU model string [2]. - Numa: The NUMA node number [2]. - Health: A boolean status [2]. - Index: The device index [3]. - Mode: The operation mode (e.g., hami-core, mig, mps) [7][3]. Implementation These utilities are located within the HAMi package, typically under pkg/device/ [6][3]. - EncodeNodeDevices: Constructs the string by iterating through a slice of DeviceInfo structs and joining fields with commas, using colons to separate multiple devices [3]. - DecodeNodeDevices: Parses the annotation string back into a slice of DeviceInfo objects, validating the input format and propagating errors if parsing fails [8][9]. The scheduler watches these annotations on all GPU-enabled nodes to maintain a global view of cluster resources [2]. If a node fails to update this annotation within 5 minutes, the scheduler marks the node as unavailable [1][5].
Citations:
- 1: https://project-hami.io/docs/developers/protocol
- 2: https://project-hami.io/docs/next/core-concepts/gpu-virtualization
- 3: https://github.com/Project-HAMi/HAMi/blob/f0dd4393/pkg/device/devices.go
- 4: https://deepwiki.com/Project-HAMi/HAMi/9-api-reference
- 5: https://project-hami.io/docs/next/developers/protocol
- 6: https://github.com/Project-HAMi/HAMi/blob/f0dd4393/pkg/device/devices_test.go
- 7: https://github.com/Project-HAMi/HAMi/blob/v2.6.1/pkg/device/nvidia/device.go
- 8: https://github.com/Project-HAMi/HAMi/blob/f0dd4393/pkg/device/nvidia/device.go
- 9: 18b8fd4
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and nearby symbols
ast-grep outline pkg/mcp/tools/list_gpu_nodes.go --view expanded || true
# Find all references to the HAMi node registration annotation and node device decoding
rg -n 'hami\.io/node-nvidia-register|nvidia\.com/gpu\.memory|nvidia\.com/gpu\.cores|DecodeNodeDevices|EncodeNodeDevices' pkg -S
# Read the target region with line numbers
sed -n '1,240p' pkg/mcp/tools/list_gpu_nodes.go | cat -n
# Read the device encoding/decoding implementation
sed -n '1,260p' pkg/device/devices.go | cat -nRepository: Project-HAMi/HAMi
Length of output: 17343
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the writer for hami.io/node-nvidia-register
sed -n '1,430p' pkg/device/nvidia/device.go | cat -n
# Show the decoder/encoder around the record format
sed -n '240,380p' pkg/device/devices.go | cat -n
# Check whether list_gpu_nodes parses HAMi annotations anywhere else
rg -n 'node-nvidia-register|gpu\.memory|gpu\.cores|DecodeNodeDevices\(' pkg/mcp pkg/device -SRepository: Project-HAMi/HAMi
Length of output: 23439
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the list_gpu_nodes tests
sed -n '1,220p' pkg/mcp/tools/list_gpu_nodes_test.go | cat -n
# Search for any writer of the NVIDIA memory/core annotations
rg -n 'nvidia\.com/gpu\.memory|nvidia\.com/gpu\.cores|gpu\.memory|gpu\.cores' pkg -S
# Search for any code that reads hami.io/node-nvidia-register in MCP tool path
rg -n 'node-nvidia-register|UnMarshalNodeDevices|DecodeNodeDevices|EncodeNodeDevices' pkg/mcp pkg/device -SRepository: Project-HAMi/HAMi
Length of output: 7613
🌐 Web query:
"nvidia.com/gpu.memory" "nvidia.com/gpu.cores" Kubernetes annotation NVIDIA GPU Feature Discovery
💡 Result:
In the context of NVIDIA GPU Feature Discovery (GFD) for Kubernetes, nvidia.com/gpu.memory and nvidia.com/gpu.cores are node labels (not Kubernetes annotations) automatically generated to describe GPU hardware properties [1][2]. GFD is a component that inspects GPU hardware via the NVIDIA Management Library (NVML) and labels Kubernetes nodes accordingly to facilitate GPU-aware scheduling [1][2]. Key details regarding these labels include: nvidia.com/gpu.memory: This label reports the memory capacity of the GPU [3]. Historically, this value was documented as being in mebibytes (MiB) [3][4]. However, discrepancies have been noted where the reported unit could be in bytes depending on the GFD version and configuration, though fixes were introduced in later versions to standardize reporting [5]. When Multi-Instance GPU (MIG) is enabled, this label can be overridden to represent the memory of each specific MIG device [3][4]. nvidia.com/gpu.cores: While the user query mentions this, official GFD documentation and label lists primarily reference other attributes such as nvidia.com/gpu.multiprocessors (representing the number of streaming multiprocessors) or specific engine counts (e.g., encoders, decoders, JPEG engines) for MIG devices [3][4]. Users should verify their specific GFD version's documentation, as labels may evolve or differ based on the deployment configuration [3][6]. These labels are used in Kubernetes node selectors or node affinity rules, allowing administrators to target specific GPU configurations (e.g., selecting nodes with a minimum amount of memory or specific architectural features) [2]. They are not typically used as annotations on pods or other objects [3][2].
Citations:
- 1: https://github.com/NVIDIA/k8s-device-plugin/blob/main/docs/gpu-feature-discovery/README.md
- 2: https://kubernetes.recipes/recipes/ai/nvidia-gpu-feature-discovery-kubernetes/
- 3: https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/main/docs/gpu-feature-discovery/README.md
- 4: https://github.com/NVIDIA/k8s-device-plugin/tree/main/docs/gpu-feature-discovery
- 5: gpu.memory label for non mig describes size in bytes NVIDIA/gpu-feature-discovery#26
- 6: https://github.com/NVIDIA/gpu-feature-discovery/tree/v0.2.0
Read these values from node labels, not annotations. nvidia.com/gpu.memory and nvidia.com/gpu.cores are exposed by NVIDIA GFD as node labels, so node.Annotations[...] will usually miss them and this tool will return 0 for memory/core on real GPU nodes. Switch to node.Labels[...] (or the actual source used by your cluster) and update the test fixture accordingly.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/tools/list_gpu_nodes.go` around lines 138 - 152, The GPU memory and
core values are being read from the wrong node metadata source in the list GPU
nodes flow. Update the extraction logic in the helper that builds the node info
to read `nvidia.com/gpu.memory` and `nvidia.com/gpu.cores` from `node.Labels`
instead of `node.Annotations`, and make sure the test fixture reflects labels so
the `ListGpuNodes`-related parsing continues to work on real GPU nodes.
Source: Coding guidelines
| if uuids, ok := pod.Annotations["nvidia.com/gpu-devices-to-use"]; ok && uuids != "" { | ||
| info.AllocatedDeviceUUIDs = splitAndClean(uuids, ",") | ||
| } else if uuids, ok := pod.Annotations["hami.io/gpu-devices-to-use"]; ok && uuids != "" { | ||
| info.AllocatedDeviceUUIDs = splitAndClean(uuids, ",") | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find the actual annotation keys used for pod GPU device allocation.
rg -n --type=go -C2 'vgpu-devices-allocated|vgpu-devices-to-allocate|gpu-devices-to-use' -g '!pkg/mcp/**'Repository: Project-HAMi/HAMi
Length of output: 155
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== annotation-related symbols ==\n'
rg -n --type=go 'vgpu-devices|gpu-devices|Annotation|annotations' pkg scheduler util api pkg/mcp tools -g '!pkg/mcp/**' || true
printf '\n== target file around relevant lines ==\n'
sed -n '1,240p' pkg/mcp/tools/list_gpu_pods.go
printf '\n== files mentioning hami.io or nvidia.com keys ==\n'
rg -n --type=go 'hami\.io/|nvidia\.com/' .Repository: Project-HAMi/HAMi
Length of output: 50374
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== target file relevant slice ==\n'
sed -n '130,190p' pkg/mcp/tools/list_gpu_pods.go
printf '\n== encode/decode definitions ==\n'
rg -n --type=go 'func EncodePodSingleDevice|func DecodePodDevices|type PodSingleDevice|SupportDevices|InRequestDevices' pkg/device pkg -g '!pkg/mcp/**'
printf '\n== nvidia device annotation wiring ==\n'
sed -n '150,220p' pkg/device/nvidia/device.go
printf '\n== device package decode/encode file if present ==\n'
fd -a 'device.go' pkg/device | sed -n '1,40p'Repository: Project-HAMi/HAMi
Length of output: 38399
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== pkg/mcp/tools/list_gpu_pods.go =='
nl -ba pkg/mcp/tools/list_gpu_pods.go | sed -n '140,180p'
echo
echo '== device encode/decode and annotation keys =='
rg -n --type=go 'func EncodePodSingleDevice|func DecodePodDevices|type PodSingleDevice|SupportDevices|InRequestDevices' pkg/device pkg -g '!pkg/mcp/**'
echo
echo '== nvidia annotation wiring =='
nl -ba pkg/device/nvidia/device.go | sed -n '160,190p'Repository: Project-HAMi/HAMi
Length of output: 232
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '380,470p' pkg/device/devices.goRepository: Project-HAMi/HAMi
Length of output: 3194
Read HAMi’s allocation annotation instead of gpu-devices-to-use.
pkg/mcp/tools/list_gpu_pods.go should read hami.io/vgpu-devices-allocated and decode HAMi’s UUID,Type,mem,cores:;... payload, not look for nvidia.com/gpu-devices-to-use / hami.io/gpu-devices-to-use and split on commas. As written, AllocatedDeviceUUIDs will stay empty or pick up non-UUID fields on real pods.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/mcp/tools/list_gpu_pods.go` around lines 152 - 156, The allocation
parsing in list_gpu_pods.go is using the wrong annotations and payload format,
so AllocatedDeviceUUIDs never reflects real HAMi pods. Update the logic in the
pod annotation handling to read hami.io/vgpu-devices-allocated in the list GPU
pods flow, and decode the HAMi UUID,Type,mem,cores:;... entries to extract only
the UUID values. Remove the fallback to nvidia.com/gpu-devices-to-use and
hami.io/gpu-devices-to-use in this path, and keep the fix localized around the
annotation parsing branch that populates info.AllocatedDeviceUUIDs.
| /* | ||
| Copyright 2024 The HAMi Authors. | ||
|
|
||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
|
|
||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
|
|
||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. | ||
| */ | ||
|
|
||
| // Package mcp_e2e exercises the full HAMi MCP server stack end-to-end via the | ||
| // in-memory MCP transport. It uses a fake Kubernetes clientset and a stubbed | ||
| // Prometheus HTTP server so the suite runs hermetically (no real cluster | ||
| // required), while still going through the real MCP protocol and JSON-RPC | ||
| // machinery the production binary uses. | ||
| package mcp_e2e | ||
|
|
||
| import ( | ||
| "context" | ||
| "encoding/json" | ||
| "net/http" | ||
| "net/http/httptest" | ||
| "strings" | ||
| "testing" | ||
|
|
||
| mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" | ||
| corev1 "k8s.io/api/core/v1" | ||
| "k8s.io/apimachinery/pkg/api/resource" | ||
| metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" | ||
| "k8s.io/client-go/kubernetes/fake" | ||
|
|
||
| hamimcp "github.com/Project-HAMi/HAMi/pkg/mcp" | ||
| "github.com/Project-HAMi/HAMi/pkg/mcp/client" | ||
| "github.com/Project-HAMi/HAMi/pkg/mcp/tools" | ||
| ) | ||
|
|
||
| // e2eFixture builds a fully-wired MCP server connected to an MCP client via | ||
| // in-memory transports. It returns the client session and a cleanup func. | ||
| type e2eFixture struct { | ||
| clientSession *mcpsdk.ClientSession | ||
| cleanup func() | ||
| } | ||
|
|
||
| func setupFixture(t *testing.T) *e2eFixture { | ||
| t.Helper() | ||
|
|
||
| gpuNode := &corev1.Node{ | ||
| ObjectMeta: metav1.ObjectMeta{ | ||
| Name: "gpu-node-1", | ||
| Labels: map[string]string{ | ||
| "gpu": "on", | ||
| }, | ||
| Annotations: map[string]string{ | ||
| "nvidia.com/gpu.memory": "16384", | ||
| "nvidia.com/gpu.cores": "100", | ||
| "hami.io/node-devices-to-register": "GPU-aaa,0,16384,100,A100,0,true,1,exclusive:", | ||
| // Sensitive — should be redacted out of describe_node output | ||
| "my-api-token": "super-secret", | ||
| }, | ||
| }, | ||
| Status: corev1.NodeStatus{ | ||
| Capacity: corev1.ResourceList{ | ||
| corev1.ResourceName("nvidia.com/gpu"): resource.MustParse("4"), | ||
| corev1.ResourceCPU: resource.MustParse("32"), | ||
| }, | ||
| Allocatable: corev1.ResourceList{ | ||
| corev1.ResourceName("nvidia.com/gpu"): resource.MustParse("4"), | ||
| }, | ||
| }, | ||
| } | ||
| cpuNode := &corev1.Node{ | ||
| ObjectMeta: metav1.ObjectMeta{Name: "cpu-only"}, | ||
| Status: corev1.NodeStatus{ | ||
| Capacity: corev1.ResourceList{ | ||
| corev1.ResourceCPU: resource.MustParse("8"), | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
| ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "ai-team"}} | ||
| gpuPod := &corev1.Pod{ | ||
| ObjectMeta: metav1.ObjectMeta{ | ||
| Name: "training-job", | ||
| Namespace: "ai-team", | ||
| Annotations: map[string]string{ | ||
| "hami.io/gpu-devices-to-use": "GPU-aaa", | ||
| }, | ||
| }, | ||
| Spec: corev1.PodSpec{ | ||
| NodeName: "gpu-node-1", | ||
| Containers: []corev1.Container{{ | ||
| Name: "trainer", | ||
| Resources: corev1.ResourceRequirements{ | ||
| Requests: corev1.ResourceList{ | ||
| corev1.ResourceName("nvidia.com/gpu"): resource.MustParse("1"), | ||
| corev1.ResourceName("nvidia.com/gpumem"): resource.MustParse("4096"), | ||
| corev1.ResourceName("nvidia.com/gpucores"): resource.MustParse("80"), | ||
| }, | ||
| }, | ||
| }}, | ||
| }, | ||
| Status: corev1.PodStatus{Phase: corev1.PodRunning}, | ||
| } | ||
|
|
||
| configMap := &corev1.ConfigMap{ | ||
| ObjectMeta: metav1.ObjectMeta{ | ||
| Name: "hami-scheduler-config", | ||
| Namespace: "hami-system", | ||
| }, | ||
| Data: map[string]string{ | ||
| "scheduler-config.yaml": "policy: binpack", | ||
| "api-token": "should-be-redacted-by-mcp", | ||
| }, | ||
| } | ||
|
|
||
| cs := fake.NewClientset(gpuNode, cpuNode, ns, gpuPod, configMap) | ||
|
|
||
| prom := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| w.Header().Set("Content-Type", "application/json") | ||
| w.Write([]byte(`{ | ||
| "status":"success", | ||
| "data":{"resultType":"vector","result":[ | ||
| {"metric":{"node":"gpu-node-1","gpu":"0"},"value":[1700000000,"42"]} | ||
| ]} | ||
| }`)) | ||
| })) | ||
|
|
||
| pc, err := client.NewPrometheusClient(prom.URL) | ||
| if err != nil { | ||
| t.Fatalf("prom client: %v", err) | ||
| } | ||
| k8s := client.NewK8sClientFromInterface(cs) | ||
|
|
||
| srv, err := hamimcp.NewServerWithClients(&hamimcp.ServerConfig{ | ||
| PrometheusURL: prom.URL, | ||
| }, k8s, pc) | ||
| if err != nil { | ||
| t.Fatalf("server init: %v", err) | ||
| } | ||
|
|
||
| ctx, cancel := context.WithCancel(context.Background()) | ||
| clientTr, serverTr := mcpsdk.NewInMemoryTransports() | ||
|
|
||
| if _, err := srv.Connect(ctx, serverTr); err != nil { | ||
| cancel() | ||
| prom.Close() | ||
| t.Fatalf("server connect: %v", err) | ||
| } | ||
|
|
||
| mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "hami-e2e-client", Version: "v0"}, nil) | ||
| clientSession, err := mcpClient.Connect(ctx, clientTr, nil) | ||
| if err != nil { | ||
| cancel() | ||
| prom.Close() | ||
| t.Fatalf("client connect: %v", err) | ||
| } | ||
|
|
||
| return &e2eFixture{ | ||
| clientSession: clientSession, | ||
| cleanup: func() { | ||
| _ = clientSession.Close() | ||
| cancel() | ||
| prom.Close() | ||
| }, | ||
| } | ||
| } | ||
|
|
||
| func extractToolText(t *testing.T, res *mcpsdk.CallToolResult) string { | ||
| t.Helper() | ||
| if res == nil || len(res.Content) == 0 { | ||
| t.Fatalf("empty tool result") | ||
| } | ||
| tc, ok := res.Content[0].(*mcpsdk.TextContent) | ||
| if !ok { | ||
| t.Fatalf("expected text content, got %T", res.Content[0]) | ||
| } | ||
| return tc.Text | ||
| } | ||
|
|
||
| func callTool(t *testing.T, cs *mcpsdk.ClientSession, name string, args map[string]any) *mcpsdk.CallToolResult { | ||
| t.Helper() | ||
| res, err := cs.CallTool(context.Background(), &mcpsdk.CallToolParams{ | ||
| Name: name, | ||
| Arguments: args, | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("CallTool(%s): %v", name, err) | ||
| } | ||
| return res | ||
| } | ||
|
|
||
| func TestMCPServer_ListTools(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| resp, err := f.clientSession.ListTools(context.Background(), nil) | ||
| if err != nil { | ||
| t.Fatalf("ListTools: %v", err) | ||
| } | ||
| if len(resp.Tools) != 5 { | ||
| t.Fatalf("expected 5 tools, got %d", len(resp.Tools)) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_ListGPUNodes(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "list_gpu_nodes", nil) | ||
| if res.IsError { | ||
| t.Fatalf("tool returned error: %s", extractToolText(t, res)) | ||
| } | ||
|
|
||
| var nodes []tools.GPUNodeInfo | ||
| if err := json.Unmarshal([]byte(extractToolText(t, res)), &nodes); err != nil { | ||
| t.Fatalf("unmarshal: %v", err) | ||
| } | ||
| if len(nodes) != 1 || nodes[0].Name != "gpu-node-1" { | ||
| t.Fatalf("expected 1 node 'gpu-node-1', got %+v", nodes) | ||
| } | ||
| if nodes[0].GPUVendor != "NVIDIA" || nodes[0].GPUCount != 4 { | ||
| t.Errorf("unexpected GPU info: %+v", nodes[0]) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_ListGPUPods(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "list_gpu_pods", map[string]any{"namespace": "ai-team"}) | ||
| if res.IsError { | ||
| t.Fatalf("tool returned error: %s", extractToolText(t, res)) | ||
| } | ||
|
|
||
| var pods []tools.GPUPodInfo | ||
| if err := json.Unmarshal([]byte(extractToolText(t, res)), &pods); err != nil { | ||
| t.Fatalf("unmarshal: %v", err) | ||
| } | ||
| if len(pods) != 1 { | ||
| t.Fatalf("expected 1 pod, got %d", len(pods)) | ||
| } | ||
| if pods[0].Name != "training-job" || pods[0].Node != "gpu-node-1" { | ||
| t.Errorf("unexpected pod info: %+v", pods[0]) | ||
| } | ||
| if pods[0].RequestedGPU != 1 { | ||
| t.Errorf("expected RequestedGPU=1, got %d", pods[0].RequestedGPU) | ||
| } | ||
| if len(pods[0].AllocatedDeviceUUIDs) != 1 || pods[0].AllocatedDeviceUUIDs[0] != "GPU-aaa" { | ||
| t.Errorf("expected allocated UUID GPU-aaa, got %v", pods[0].AllocatedDeviceUUIDs) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_DescribeNode(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "describe_node", map[string]any{"node": "gpu-node-1"}) | ||
| if res.IsError { | ||
| t.Fatalf("tool returned error: %s", extractToolText(t, res)) | ||
| } | ||
|
|
||
| body := extractToolText(t, res) | ||
| var desc tools.NodeDescription | ||
| if err := json.Unmarshal([]byte(body), &desc); err != nil { | ||
| t.Fatalf("unmarshal: %v", err) | ||
| } | ||
| if desc.Name != "gpu-node-1" { | ||
| t.Errorf("expected node name gpu-node-1, got %s", desc.Name) | ||
| } | ||
| if len(desc.GPUDevices) != 1 || desc.GPUDevices[0].Type != "A100" { | ||
| t.Errorf("expected one A100 device, got %+v", desc.GPUDevices) | ||
| } | ||
|
|
||
| // The node has annotation "my-api-token"; redact pass should mask its value. | ||
| if !strings.Contains(body, "REDACTED") { | ||
| t.Errorf("expected sensitive annotation to be redacted, body=%s", body) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_DescribeNode_Missing(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "describe_node", map[string]any{"node": "does-not-exist"}) | ||
| if !res.IsError { | ||
| t.Fatalf("expected error result for missing node, got %s", extractToolText(t, res)) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_GetQuotaUsage(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "get_quota_usage", map[string]any{"namespace": "ai-team"}) | ||
| if res.IsError { | ||
| t.Fatalf("tool returned error: %s", extractToolText(t, res)) | ||
| } | ||
|
|
||
| var usage tools.QuotaUsage | ||
| if err := json.Unmarshal([]byte(extractToolText(t, res)), &usage); err != nil { | ||
| t.Fatalf("unmarshal: %v", err) | ||
| } | ||
| if usage.Namespace != "ai-team" { | ||
| t.Errorf("expected namespace ai-team, got %s", usage.Namespace) | ||
| } | ||
| if usage.GPUMemoryUsedGiB != 4 { | ||
| t.Errorf("expected 4 GiB used (4096 MiB), got %v", usage.GPUMemoryUsedGiB) | ||
| } | ||
| if usage.GPUCoreUsed != 80 { | ||
| t.Errorf("expected 80 cores, got %v", usage.GPUCoreUsed) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_GetGPUMetrics(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "get_gpu_metrics", map[string]any{ | ||
| "metric": "hami_gpu_device_count", | ||
| "node": "gpu-node-1", | ||
| }) | ||
| if res.IsError { | ||
| t.Fatalf("tool returned error: %s", extractToolText(t, res)) | ||
| } | ||
|
|
||
| var metrics []tools.GPUMetric | ||
| if err := json.Unmarshal([]byte(extractToolText(t, res)), &metrics); err != nil { | ||
| t.Fatalf("unmarshal: %v", err) | ||
| } | ||
| if len(metrics) != 1 { | ||
| t.Fatalf("expected 1 metric value, got %d", len(metrics)) | ||
| } | ||
| if metrics[0].Value != 42 { | ||
| t.Errorf("expected value 42, got %v", metrics[0].Value) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_ReadConfigResource(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res, err := f.clientSession.ReadResource(context.Background(), &mcpsdk.ReadResourceParams{ | ||
| URI: "hami://config/scheduler", | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("ReadResource: %v", err) | ||
| } | ||
| if len(res.Contents) != 1 { | ||
| t.Fatalf("expected 1 content, got %d", len(res.Contents)) | ||
| } | ||
| body := res.Contents[0].Text | ||
| if !strings.Contains(body, "binpack") { | ||
| t.Errorf("expected scheduler config to contain 'binpack', got %s", body) | ||
| } | ||
| if !strings.Contains(body, "REDACTED") { | ||
| t.Errorf("expected api-token to be redacted, got %s", body) | ||
| } | ||
| } | ||
|
|
||
| func TestMCPServer_GetGPUMetrics_RequiresMetric(t *testing.T) { | ||
| f := setupFixture(t) | ||
| defer f.cleanup() | ||
|
|
||
| res := callTool(t, f.clientSession, "get_gpu_metrics", map[string]any{}) | ||
| if !res.IsError { | ||
| t.Fatalf("expected error for missing metric, got %s", extractToolText(t, res)) | ||
| } | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
E2E suite doesn't use Ginkgo/Gomega.
This file lives under test/e2e/mcp/ and uses plain testing.T assertions throughout instead of Ginkgo/Gomega (Describe/It/Expect).
As per coding guidelines, "E2E tests must use Ginkgo/Gomega."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/e2e/mcp/mcp_e2e_test.go` around lines 1 - 374, The MCP end-to-end suite
is written with plain testing.T assertions, but the project guideline requires
Ginkgo/Gomega for e2e tests. Refactor the tests in mcp_e2e_test.go, including
setupFixture, callTool, and the TestMCPServer_* cases, into Ginkgo specs using
Describe/Context/It and Gomega expectations, while keeping the same MCP
client/server behavior and assertions.
Source: Coding guidelines
| t.Fatalf("server init: %v", err) | ||
| } | ||
|
|
||
| ctx, cancel := context.WithCancel(context.Background()) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Likely same modernize-lint failure as server_test.go.
context.WithCancel(context.Background()) here mirrors the pattern flagged for t.Context() modernization in pkg/mcp/server_test.go:95.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/e2e/mcp/mcp_e2e_test.go` at line 147, The test setup is using
context.WithCancel(context.Background()) instead of the newer t.Context()
pattern, matching the modernization issue seen in the MCP tests. Update the
context creation in the relevant e2e test code to derive from t.Context() and
keep the cancel handling aligned with the surrounding test lifecycle; use the
existing test context in the mcp_e2e_test helper or test body where ctx is
created.
Source: Linters/SAST tools
2 participants
What type of PR is this?
/kind feature
What this PR does / why we need it:
Which issue(s) this PR fixes:
Fixes #
Special notes for your reviewer:
Does this PR introduce a user-facing change?:
Summary by CodeRabbit