Support extensions and custom controller flags in deploy-to-kube tool

Makefile

@@ -12,17 +12,13 @@ build:
 KIND_CLUSTER=agent-sandbox
 
 .PHONY: deploy-kind
+# `EXTENSIONS=true make deploy-kind` to deploy with Extensions enabled.
+# `CONTROLLER_ARGS="--enable-pprof-debug --zap-log-level=debug" make deploy-kind` to deploy with custom controller flags.
+# `CONTROLLER_ONLY=true make deploy-kind` to build and push only the controller image.
 deploy-kind:
 	./dev/tools/create-kind-cluster --recreate ${KIND_CLUSTER} --kubeconfig bin/KUBECONFIG
-	./dev/tools/push-images --image-prefix=kind.local/ --kind-cluster-name=${KIND_CLUSTER}
-	./dev/tools/deploy-to-kube --image-prefix=kind.local/
-
-	@if [ "$(EXTENSIONS)" = "true" ]; then \
-		echo "🔧 Patching controller to enable extensions..."; \
-		kubectl patch deployment agent-sandbox-controller \
-			-n agent-sandbox-system \
-			-p '{"spec": {"template": {"spec": {"containers": [{"name": "agent-sandbox-controller", "args": ["--extensions=true"]}]}}}}'; \
-	fi
+	./dev/tools/push-images --image-prefix=kind.local/ --kind-cluster-name=${KIND_CLUSTER} $(if $(filter true,$(CONTROLLER_ONLY)),--controller-only)
+	./dev/tools/deploy-to-kube --image-prefix=kind.local/ $(if $(filter true,$(EXTENSIONS)),--extensions) $(if $(CONTROLLER_ARGS),--controller-args="$(CONTROLLER_ARGS)")
 
 .PHONY: deploy-cloud-provider-kind
 deploy-cloud-provider-kind:
@@ -85,7 +81,7 @@ release-manifests:
 	@if [ -z "$(TAG)" ]; then echo "TAG is required (e.g., make release-manifests TAG=vX.Y.Z)"; exit 1; fi
 	go mod tidy
 	go generate ./...
-	./dev/tools/release --tag=${TAG} 
+	./dev/tools/release --tag=${TAG}
 
 # Example usage:
 # make release-python-sdk TAG=v0.1.1rc1 (to release only on TestPyPI, blocked from PyPI in workflow)

dev/ci/shared/runner.py

@@ -53,7 +53,7 @@ def setup_cluster(self, args, extra_push_images_args=None):
         if result.returncode != 0:
             return result
 
-        result = subprocess.run([f"{self.repo_root}/dev/tools/deploy-to-kube", "--image-prefix", args.image_prefix, "--image-tag", image_tag])
+        result = subprocess.run([f"{self.repo_root}/dev/tools/deploy-to-kube", "--image-prefix", args.image_prefix, "--image-tag", image_tag, "--extensions"])
         if result.returncode != 0:
             return result
 

dev/tools/deploy-to-kube

@@ -22,8 +22,18 @@ import yaml
 from shared import utils
 
 
-def find_and_replace_images(doc, service_name, image_prefix, image_tag):
-    """Recursively finds and replaces container images in a Kubernetes manifest."""
+def find_and_replace_images(doc, image_prefix, image_tag):
+    """Recursively finds and replaces container images in a Kubernetes manifest.
+
+    This function specifically targets image fields that either use the 'ko://'
+    prefix or have a ':latest' tag, replacing them with a fully qualified
+    and versioned image name.
+
+    Args:
+        doc (dict): The Kubernetes manifest document to process.
+        image_prefix (str): The image registry prefix (e.g., 'us-docker.pkg.dev/my-project/').
+        image_tag (str): The image tag to apply (e.g., 'v1.2.3').
+    """
     def fn(field_path, v):
         if field_path == ".spec.template.spec.containers[].image":
             # We only replace images that look like :latest
@@ -42,8 +52,34 @@ def find_and_replace_images(doc, service_name, image_prefix, image_tag):
     walk_object(doc, fn)
 
 
+
+def append_controller_flags(doc, flags_str):
+    """Appends arbitrary flags to the controller container arguments."""
+    if not flags_str:
+        return
+    if doc.get("kind") == "Deployment" and doc.get("metadata", {}).get("name") == "agent-sandbox-controller":
+        containers = doc.get("spec", {}).get("template", {}).get("spec", {}).get("containers", [])
+        for container in containers:
+            if container.get("name") == "agent-sandbox-controller":
+                if "args" not in container:
+                    container["args"] = []
+                flags = flags_str.split()
+                for flag in flags:
+                    if flag not in container["args"]:
+                        container["args"].append(flag)
+                        print(f"Appended controller flag: {flag}")
+
+
 def walk_object(obj, fn, field_path=""):
-    """Recursively walks an object, and calls fn with the field_path to allow for value replacement."""
+    """Recursively walks through a nested dictionary or list and applies a function.
+
+    Args:
+        obj (dict | list): The object to walk through.
+        fn (function): A function to apply to each field. It receives the
+            field path (e.g., '.spec.template.spec') and the value.
+        field_path (str, optional): The current path in the object. Used for
+            recursive calls. Defaults to "".
+    """
     if isinstance(obj, dict):
         for k, v in obj.items():
             v2 = fn(f"{field_path}.{k}", v)
@@ -59,7 +95,13 @@ def walk_object(obj, fn, field_path=""):
 
 
 def kubectl_apply_objects(docs, repo_root):
-    """Helper to apply a list of k8s objects."""
+    """Applies a list of Kubernetes documents using 'kubectl apply'.
+
+    Args:
+        docs (list): A list of Kubernetes manifest documents.
+        repo_root (str): The path to the repository root, used as the working
+            directory for the subprocess call.
+    """
     if not docs: return
     for doc in docs:
         kind = doc.get("kind", "Unknown")
@@ -70,61 +112,144 @@ def kubectl_apply_objects(docs, repo_root):
     subprocess.run(["kubectl", "apply", "-f", "-"], cwd=repo_root, check=True, input=string_stream.getvalue(), text=True)
 
 
-def main(args):
-    repo_root = utils.get_repo_root()
+def gather_manifests(manifests_path):
+    """Scans a directory for YAML files and parses them into a structured list.
 
-    manifests_path = os.path.join(repo_root, "k8s")
+    This function walks through the given directory, reads all '.yaml' or '.yml'
+    files, and handles multi-Kubernetes object files by splitting them into individual
+    objects.
+
+    Args:
+        manifests_path (str): The path to the directory containing manifests.
 
+    Returns:
+        list: A flat list of dictionaries. Each dictionary represents a single
+              Kubernetes object and contains the document itself ('doc'), its
+              original 'filename', its 'kind', and the 'root' directory it was found in.
+    """
+    all_docs = []
     if not os.path.isdir(manifests_path):
         print("k8s directory not found")
-        return
-
-    prereq_docs = []
-    other_docs = []
+        return all_docs
 
-    image_prefix = utils.get_image_prefix(args)
-    image_tag = args.image_tag or utils.get_image_tag()
 
     for root, dirs, files in os.walk(manifests_path):
         for filename in files:
             if not (filename.endswith(".yaml") or filename.endswith(".yml")):
                 continue
-
             path = os.path.join(root, filename)
             with open(path, "r") as f:
-                # Use safe_load_all for multi-document YAML files
                 docs = list(yaml.safe_load_all(f))
+                for doc in docs:
+                    if doc:
+                        all_docs.append({
+                            "doc": doc,
+                            "filename": filename,
+                            "kind": doc.get("kind"),
+                        })
+    return all_docs
+
+
+def process_manifests(all_docs, args):
+    """Filters, modifies, and sorts a list of raw manifest objects.
+
+    This function applies the core deployment logic based on the script's
+    arguments. It filters out unwanted documents, replaces image tags, and
+    sorts the documents into three batches for ordered application:
+    prerequisites, standard documents, and extensions documents.
+
+    Args:
+        all_docs (list): The raw list of manifest objects from gather_manifests.
+        args (argparse.Namespace): The parsed command-line arguments.
+
+    Returns:
+        tuple: A tuple containing three lists of documents:
+               (prereq_docs, other_docs, extensions_docs).
+    """
+    prereq_docs = []
+    extensions_docs = []
+    other_docs = []
+
+    image_prefix = utils.get_image_prefix(args)
+    image_tag = utils.get_image_tag()
+
+    for item in all_docs:
+        doc = item["doc"]
+        filename = item["filename"]
+        kind = item["kind"]
+
+        is_extension_file = filename.startswith("extensions")
+        if is_extension_file and not args.extensions:
+            continue
+
+        # When extensions is enabled, we skip the core controller manifest because
+        # it will be replaced by the one in the extensions directory.
+        is_core_controller = (
+            doc.get("metadata", {}).get("name") == "agent-sandbox-controller" and
+            kind in ["StatefulSet", "Deployment"] and
+            not is_extension_file
+        )
+        if args.extensions and is_core_controller:
+            continue
+
+        find_and_replace_images(doc, image_prefix, image_tag)
+        append_controller_flags(doc, getattr(args, "controller_args", ""))
+
+        if kind in ["CustomResourceDefinition", "Namespace"]:
+            prereq_docs.append(doc)
+        elif is_extension_file:
+            extensions_docs.append(doc)
+        else:
+            other_docs.append(doc)
+
+    return prereq_docs, other_docs, extensions_docs
+
+
+def apply_manifests(prereq_docs, other_docs, extensions_docs, repo_root, apply_extensions):
+    """Applies the processed manifests to the cluster in the correct order.
 
-            service_name = os.path.basename(root)
-
-            # Process each document in the file
-            for doc in docs:
-                if not doc: continue
-
-                # Rule 1: CRDs need to be applied first
-                is_crd = doc.get("kind") == "CustomResourceDefinition"
-                if is_crd:
-                      prereq_docs.append(doc)
-                      continue
-
-                # Rule 2: Namespaces also need to be applied first
-                if doc.get("kind") == "Namespace":
-                      prereq_docs.append(doc)
-                      continue
-
-                # Rule 3: Everything else (Deployments, StatefulSets, etc.)
-                # If we got here, it's a resource we want to apply.
-                find_and_replace_images(doc, service_name, image_prefix, image_tag)
-                other_docs.append(doc)
+    The order of application is:
+    1. Prerequisites (Namespaces, CRDs)
+    2. Standard resources
+    3. Extension resources (if enabled)
 
+    Args:
+        prereq_docs (list): A list of prerequisite manifest documents.
+        other_docs (list): A list of standard manifest documents.
+        extension_docs (list): A list of extension manifest documents.
+        repo_root (str): The path to the repository root.
+        apply_extensions (bool): Whether to apply the extension manifests.
+    """
     print(f"Applying {len(prereq_docs)} prerequisite manifests (Namespace, CRDs)...")
     kubectl_apply_objects(prereq_docs, repo_root)
 
     print(f"Applying {len(other_docs)} other manifests...")
     kubectl_apply_objects(other_docs, repo_root)
+
+    if apply_extensions:
+        print(f"Applying {len(extensions_docs)} extensions manifests...")
+        kubectl_apply_objects(extensions_docs, repo_root)
+
     print("Deployment complete.")
 
 
+def main(args):
+    """Orchestrates the Kubernetes deployment process.
+
+    This function calls the gather, process, and apply functions in sequence
+    to carry out the full deployment workflow.
+
+    Args:
+        args (argparse.Namespace): The parsed command-line arguments.
+    """
+    repo_root = utils.get_repo_root()
+    manifests_path = os.path.join(repo_root, "k8s")
+
+    all_docs = gather_manifests(manifests_path)
+    prereq_docs, other_docs, extension_docs = process_manifests(all_docs, args)
+    apply_manifests(prereq_docs, other_docs, extension_docs, repo_root, args.extensions)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Deploy manifests under k8s/ to Kubernetes cluster")
     parser.add_argument("--image-prefix",
@@ -137,5 +262,15 @@ if __name__ == "__main__":
                         help="Image tag for this build. If not set, a default tag is generated based on date and git commit",
                         type=str,
                         default=os.getenv("IMAGE_TAG"))
+    parser.add_argument("--extensions",
+                        action="store_true",
+                        help="Apply the extensions controller manifest. This flag must be set to deploy extensions.")
+    parser.add_argument("--controller-args",
+                        dest="controller_args",
+                        help="Space-separated flags to pass to the controller container. Example " \
+                        "usage: --controller-args \"--enable-pprof-debug --zap-log-level=debug\". " \
+                        "Do not use this flag to enable extensions, use --extensions instead.",
+                        type=str,
+                        default="")
     args = parser.parse_args()
     main(args)