fix(kaito): detect KAITO installed via AKS AI-toolchain add-on (#322)

Signed-off-by: Suraj Deshmukh <suraj.deshmukh@microsoft.com>
Co-authored-by: Robbie Cronin <robert.owen.cronin@gmail.com>

Author: surajssd

.gitignore

@@ -63,6 +63,7 @@ providers/*/bin/
 /test-results/
 /playwright-report/
 /playwright/.cache/
+.playwright-mcp/
 frontend/test-results/
 frontend/playwright-report/
 frontend/blob-report/

backend/src/services/kubernetes-runtime-status.test.ts

@@ -378,6 +378,33 @@ describe('KubernetesService - Runtime Status', () => {
     expect(status.message).toBe('KAITO workspace CRD found and KAITO operator pods are ready');
   });
 
+  test('reports KAITO as installed when the AKS AI-toolchain-operator add-on pod is running in kube-system', async () => {
+    restores.push(
+      mockServiceMethod(kubernetesService, 'checkCRDExists', async (crdName: string) => crdName === 'workspaces.kaito.sh'),
+    );
+    // The AKS add-on runs the KAITO operator in kube-system, labeled
+    // app=ai-toolchain-operator rather than the upstream Helm chart labels, so
+    // it only surfaces through the cross-namespace fallback search.
+    mockOperatorPods('kaito-workspace', kaitoOperatorSelector, [], [
+      {
+        metadata: { namespace: 'kube-system', name: 'kaito-workspace-557dbc5ffb-smczp', labels: { app: 'ai-toolchain-operator' } },
+        status: {
+          phase: 'Running',
+          containerStatuses: [
+            { ready: true, restartCount: 0 },
+          ],
+        },
+      },
+    ]);
+
+    const status = await kubernetesService.checkKaitoInstallationStatus();
+
+    expect(status.installed).toBe(true);
+    expect(status.crdFound).toBe(true);
+    expect(status.operatorRunning).toBe(true);
+    expect(status.message).toBe('KAITO workspace CRD found and KAITO operator pods are ready in kube-system');
+  });
+
   test('reports Dynamo as installed when a ready operator pod is found', async () => {
     restores.push(
       mockServiceMethod(kubernetesService, 'checkCRDExists', async (crdName: string) => crdName === 'dynamographdeployments.nvidia.com'),

backend/src/services/kubernetes.ts

@@ -31,6 +31,17 @@ const INFERENCE_EXTENSION_VERSION_ANNOTATIONS = [
 const KAITO_WORKSPACE_CRD = 'workspaces.kaito.sh';
 const KAITO_NAMESPACE = 'kaito-workspace';
 const KAITO_OPERATOR_POD_SELECTOR = 'app.kubernetes.io/name=workspace,app.kubernetes.io/instance=kaito-workspace';
+// The AKS AI-toolchain-operator add-on installs KAITO in kube-system. Verified
+// against a live `--enable-ai-toolchain-operator` cluster, the add-on operator
+// POD carries ONLY the bare `app=ai-toolchain-operator` label — it does NOT
+// carry `app.kubernetes.io/name` (that key is present on the Deployment but not
+// propagated to the pod template). So this pod probe must match on `app`; using
+// `app.kubernetes.io/name=ai-toolchain-operator` here would match nothing.
+// NOTE: the Go provider shim probes the Deployment instead and intentionally
+// uses `app.kubernetes.io/name=ai-toolchain-operator` — see
+// providers/kaito/upstream_health.go (listWorkspaceController). The two paths
+// key off different labels on purpose because they inspect different objects.
+const KAITO_AKS_ADDON_POD_SELECTOR = 'app=ai-toolchain-operator';
 const DYNAMO_CRD = 'dynamographdeployments.nvidia.com';
 const DYNAMO_NAMESPACE = 'dynamo-system';
 const DYNAMO_OPERATOR_POD_SELECTOR = 'control-plane=controller-manager,app.kubernetes.io/name=dynamo-operator,app.kubernetes.io/instance=dynamo-platform';
@@ -156,6 +167,14 @@ const RUNTIME_INSTALLATION_PROBES: Record<string, RuntimeInstallationProbe> = {
     operatorNamespace: KAITO_NAMESPACE,
     operatorPodSelectors: [KAITO_OPERATOR_POD_SELECTOR],
     fallbackPodSelectors: ['app.kubernetes.io/name=workspace'],
+    // The AKS add-on pod only ever lives in kube-system, never in
+    // kaito-workspace, so it is matched exclusively in the cross-namespace
+    // pass. Listing it explicitly here (rather than relying on the implicit
+    // `crossNamespaceFallbackPodSelectors = fallbackPodSelectors` default)
+    // keeps add-on detection working even if KAITO later gains other
+    // same-namespace fallbacks, and avoids a guaranteed-empty query for
+    // `app=ai-toolchain-operator` against kaito-workspace on every probe.
+    crossNamespaceFallbackPodSelectors: ['app.kubernetes.io/name=workspace', KAITO_AKS_ADDON_POD_SELECTOR],
   },
   dynamo: {
     providerName: 'Dynamo',

providers/kaito/upstream_health.go

@@ -17,7 +17,9 @@ import (
 
 	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/api/meta"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/selection"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -43,9 +45,23 @@ const (
 
 // Well-known resource selectors the probe looks up.
 const (
-	kaitoDeploymentSelectorKey    = "app.kubernetes.io/name"
-	kaitoDeploymentSelectorValue  = "workspace"
-	controllerMissingUserMessage  = "The KAITO workspace controller is not running. Install KAITO with `helm install kaito-workspace kaito/workspace`. If this cluster was provisioned with `--enable-ai-toolchain-operator`, disable the AKS extension first (`az aks update --disable-ai-toolchain-operator ...`)."
+	kaitoDeploymentSelectorKey   = "app.kubernetes.io/name"
+	kaitoDeploymentSelectorValue = "workspace"
+	// aksAddonSelectorValue matches the KAITO controller Deployment installed by
+	// the AKS AI-toolchain-operator add-on. Verified against a live
+	// `--enable-ai-toolchain-operator` cluster, the add-on Deployment carries
+	// BOTH app.kubernetes.io/name=ai-toolchain-operator AND app=ai-toolchain-operator
+	// (in kube-system), so probing the dotted key here is correct.
+	// NOTE: the add-on POD only carries the bare `app` label, so the TypeScript
+	// pod probe in backend/src/services/kubernetes.ts intentionally matches
+	// `app=ai-toolchain-operator` instead. The two paths use different label
+	// keys on purpose because they inspect different objects (Deployment here,
+	// Pod there).
+	aksAddonSelectorValue = "ai-toolchain-operator"
+	// controllerMissingUserMessage covers both the "never installed" case and the
+	// "add-on enabled but unhealthy" case, pointing at the namespace to inspect
+	// for each install path.
+	controllerMissingUserMessage  = "The KAITO workspace controller is not running. Install it with `helm install kaito-workspace kaito/workspace` (check the kaito-workspace namespace), or via the AKS AI toolchain operator add-on `az aks update --enable-ai-toolchain-operator ...` (check the kube-system namespace)."
 	controllerNotReadyUserMessage = "The KAITO workspace controller Deployment %s/%s exists but has no ready replicas."
 	crdMissingUserMessage         = "KAITO Workspace CRD not found. Install KAITO."
 )
@@ -132,16 +148,35 @@ func isNoKindMatch(err error) bool {
 // workspace controller label selector. It also returns a second return value
 // indicating whether any Deployment with the selector was found (so callers
 // can distinguish "missing" from "not ready").
+//
+// The selector matches both the upstream Helm chart
+// (app.kubernetes.io/name=workspace) and the AKS AI-toolchain-operator add-on
+// (app.kubernetes.io/name=ai-toolchain-operator). The List is cluster-wide so
+// the controller is found regardless of which namespace it runs in
+// (kaito-workspace for the chart, kube-system for the add-on).
 func listWorkspaceController(ctx context.Context, direct client.Client) (*appsv1.Deployment, bool, error) {
+	req, err := labels.NewRequirement(
+		kaitoDeploymentSelectorKey,
+		selection.In,
+		[]string{kaitoDeploymentSelectorValue, aksAddonSelectorValue},
+	)
+	if err != nil {
+		return nil, false, fmt.Errorf("build controller selector: %w", err)
+	}
+	selector := labels.NewSelector().Add(*req)
+
 	list := &appsv1.DeploymentList{}
-	if err := direct.List(ctx, list, client.MatchingLabels{kaitoDeploymentSelectorKey: kaitoDeploymentSelectorValue}); err != nil {
+	if err := direct.List(ctx, list, client.MatchingLabelsSelector{Selector: selector}); err != nil {
 		return nil, false, fmt.Errorf("list deployments: %w", err)
 	}
 	if len(list.Items) == 0 {
 		return nil, false, nil
 	}
 	// Prefer a ready one; otherwise return the first item so the caller can
-	// reference the namespace/name in the message.
+	// reference the namespace/name in the message. When both the Helm chart and
+	// the AKS add-on are present, the In selector returns both Deployments and
+	// this loop reports the first ready one — installed/healthy is what matters,
+	// not which install path wins the tiebreak.
 	for i := range list.Items {
 		d := &list.Items[i]
 		if d.Status.ReadyReplicas > 0 {

providers/kaito/upstream_health_test.go

@@ -123,6 +123,26 @@ func newKaitoDeployment(namespace, name string, readyReplicas int32) *appsv1.Dep
 	}
 }
 
+// newAKSAddonDeployment builds a Deployment that mimics the KAITO controller
+// installed by the AKS AI-toolchain-operator add-on. The label set mirrors what
+// a live `--enable-ai-toolchain-operator` cluster emits: the Deployment carries
+// BOTH app.kubernetes.io/name=ai-toolchain-operator and app=ai-toolchain-operator
+// (typically in kube-system), unlike the upstream Helm chart's
+// app.kubernetes.io/name=workspace. The probe matches on the dotted key.
+func newAKSAddonDeployment(namespace, name string, readyReplicas int32) *appsv1.Deployment {
+	return &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+			Labels: map[string]string{
+				kaitoDeploymentSelectorKey: aksAddonSelectorValue,
+				"app":                      aksAddonSelectorValue,
+			},
+		},
+		Status: appsv1.DeploymentStatus{ReadyReplicas: readyReplicas},
+	}
+}
+
 func stringContains(s, substr string) bool {
 	for i := 0; i+len(substr) <= len(s); i++ {
 		if s[i:i+len(substr)] == substr {
@@ -161,6 +181,47 @@ func TestProbe_ControllerReady(t *testing.T) {
 	}
 }
 
+func TestProbe_ControllerReady_AKSAddon(t *testing.T) {
+	// KAITO installed via the AKS AI-toolchain-operator add-on: the controller
+	// runs in kube-system with app.kubernetes.io/name=ai-toolchain-operator.
+	// The probe must recognise it as a healthy upstream controller.
+	d := newAKSAddonDeployment("kube-system", "kaito-workspace", 1)
+	c := probeClientBuilderWithWorkspace(t).
+		WithObjects(d).
+		Build()
+
+	got := probeUpstreamController(context.Background(), c)
+
+	if !got.Healthy {
+		t.Errorf("expected Healthy=true, got %+v", got)
+	}
+	if got.Reason != ReasonUpstreamHealthy {
+		t.Errorf("expected Reason=%s, got %s", ReasonUpstreamHealthy, got.Reason)
+	}
+}
+
+func TestProbe_ControllerNotReady_AKSAddon(t *testing.T) {
+	// The add-on Deployment exists in kube-system but has no ready replicas:
+	// the probe must report NotReady (not Missing), referencing its location.
+	d := newAKSAddonDeployment("kube-system", "kaito-workspace", 0)
+	c := probeClientBuilderWithWorkspace(t).
+		WithObjects(d).
+		Build()
+
+	got := probeUpstreamController(context.Background(), c)
+
+	if got.Healthy {
+		t.Error("expected Healthy=false")
+	}
+	if got.Reason != ReasonUpstreamControllerNotReady {
+		t.Errorf("expected Reason=%s, got %s", ReasonUpstreamControllerNotReady, got.Reason)
+	}
+	want := "kube-system/kaito-workspace"
+	if !stringContains(got.Message, want) {
+		t.Errorf("expected Message to contain %q, got %q", want, got.Message)
+	}
+}
+
 func TestProbe_ControllerNotReady(t *testing.T) {
 	d := newKaitoDeployment("kaito-workspace", "kaito-workspace", 0)
 	c := probeClientBuilderWithWorkspace(t).