From 029aabc4cbbb391f9ad1bcfe8cf12eef0b651dbf Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 30 Oct 2025 06:26:30 +0000
Subject: [PATCH 1/6] Initial plan


From 880ffb54db3561062d426a44212d50445c064025 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 30 Oct 2025 06:43:29 +0000
Subject: [PATCH 2/6] Add pipeline step secrets support to SDK and CLI

- Implement step_version_secrets parsing in pipeline_builder.py
- Add helper methods to Pipeline client (get_pipeline_version, create_pipeline_version, list_step_secrets)
- Include step secrets in lockfile generation
- Update pipeline config template with example step secrets
- Add comprehensive test coverage for step secrets functionality

Co-authored-by: ydixit-clarifai <219138670+ydixit-clarifai@users.noreply.github.com>
---
 clarifai/cli/templates/pipeline_templates.py  |   9 +
 clarifai/client/pipeline.py                   | 159 ++++++++++
 .../runners/pipelines/pipeline_builder.py     |  46 +++
 tests/test_pipeline_client.py                 | 181 +++++++++++
 tests/test_pipeline_step_secrets.py           | 300 ++++++++++++++++++
 5 files changed, 695 insertions(+)
 create mode 100644 tests/test_pipeline_step_secrets.py

diff --git a/clarifai/cli/templates/pipeline_templates.py b/clarifai/cli/templates/pipeline_templates.py
index 854685b0..513df4fc 100644
--- a/clarifai/cli/templates/pipeline_templates.py
+++ b/clarifai/cli/templates/pipeline_templates.py
@@ -45,6 +45,15 @@ def get_pipeline_config_template(
         - name: sequence
           steps:
 {steps_yaml}
+    # Optional: Define secrets for pipeline steps
+    # step_version_secrets:
+    #   step-0:
+    #     secrets:
+    #       API_KEY: users/{user_id}/secrets/my-api-key
+    #       DB_PASSWORD: users/{user_id}/secrets/db-secret
+    #   step-1:
+    #     secrets:
+    #       EMAIL_TOKEN: users/{user_id}/secrets/email-token
 """
 
 
diff --git a/clarifai/client/pipeline.py b/clarifai/client/pipeline.py
index c4fd75b3..3729de3a 100644
--- a/clarifai/client/pipeline.py
+++ b/clarifai/client/pipeline.py
@@ -330,3 +330,162 @@ def _display_new_logs(self, run_id: str, seen_logs: set, current_page: int = 1)
             logger.debug(f"Error fetching logs: {e}")
             # Return current page on error to retry the same page next fetch
             return current_page
+
+    def get_pipeline_version(self, pipeline_version_id: Optional[str] = None) -> Dict:
+        """Get pipeline version details including step secrets.
+
+        Args:
+            pipeline_version_id (Optional[str]): The pipeline version ID. If None, uses self.pipeline_version_id.
+
+        Returns:
+            Dict: Pipeline version information including step_version_secrets if configured.
+        """
+        version_id = pipeline_version_id or self.pipeline_version_id
+        if not version_id:
+            raise UserError("pipeline_version_id is required")
+
+        request = service_pb2.GetPipelineVersionRequest()
+        request.user_app_id.CopyFrom(self.user_app_id)
+        request.pipeline_id = self.pipeline_id
+        request.pipeline_version_id = version_id
+
+        response = self.STUB.GetPipelineVersion(request, metadata=self.auth_helper.metadata)
+
+        if response.status.code != status_code_pb2.StatusCode.SUCCESS:
+            raise UserError(
+                f"Failed to get pipeline version: {response.status.description}. "
+                f"Details: {response.status.details}"
+            )
+
+        return json_format.MessageToDict(
+            response.pipeline_version, preserving_proto_field_name=True
+        )
+
+    def create_pipeline_version(
+        self,
+        orchestration_spec: Dict,
+        step_version_secrets: Optional[Dict[str, Dict[str, str]]] = None,
+        description: Optional[str] = None,
+    ) -> str:
+        """Create a new pipeline version with optional step secrets.
+
+        Note: This creates a new version by patching the pipeline with a new version.
+
+        Args:
+            orchestration_spec (Dict): The orchestration specification for the pipeline.
+            step_version_secrets (Optional[Dict[str, Dict[str, str]]]): Map of step references to their secrets.
+                Format: {step_ref: {secret_name: secret_path}}
+                Example: {"step-0": {"API_KEY": "users/user123/secrets/my-api-key"}}
+            description (Optional[str]): Description for the pipeline version.
+
+        Returns:
+            str: The created pipeline version ID.
+        """
+        pipeline_version = resources_pb2.PipelineVersion()
+        if description:
+            pipeline_version.description = description
+
+        # Set orchestration spec
+        if "argo_orchestration_spec" in orchestration_spec:
+            argo_spec_str = orchestration_spec["argo_orchestration_spec"]
+            import yaml
+
+            argo_spec = yaml.safe_load(argo_spec_str)
+            api_version = argo_spec.get("apiVersion", "argoproj.io/v1alpha1")
+
+            orchestration_spec_proto = resources_pb2.OrchestrationSpec()
+            argo_orchestration_spec_proto = resources_pb2.ArgoOrchestrationSpec()
+            argo_orchestration_spec_proto.api_version = api_version
+            import json
+
+            argo_orchestration_spec_proto.spec_json = json.dumps(argo_spec)
+
+            orchestration_spec_proto.argo_orchestration_spec.CopyFrom(
+                argo_orchestration_spec_proto
+            )
+            pipeline_version.orchestration_spec.CopyFrom(orchestration_spec_proto)
+
+        # Add step_version_secrets if provided
+        if step_version_secrets:
+            for step_ref, secrets in step_version_secrets.items():
+                if not secrets:
+                    continue
+                step_secret_config = resources_pb2.StepSecretConfig()
+                for secret_name, secret_path in secrets.items():
+                    step_secret_config.secrets[secret_name] = secret_path
+                pipeline_version.config.step_version_secrets[step_ref].CopyFrom(step_secret_config)
+
+        # Make the API call using PatchPipelineVersions
+        # This creates a new version for an existing pipeline
+        request = service_pb2.PatchPipelineVersionsRequest()
+        request.user_app_id.CopyFrom(self.user_app_id)
+        request.pipeline_id = self.pipeline_id
+        request.pipeline_versions.append(pipeline_version)
+        request.action = "overwrite"  # Create a new version
+
+        response = self.STUB.PatchPipelineVersions(request, metadata=self.auth_helper.metadata)
+
+        if response.status.code != status_code_pb2.StatusCode.SUCCESS:
+            raise UserError(
+                f"Failed to create pipeline version: {response.status.description}. "
+                f"Details: {response.status.details}"
+            )
+
+        if not response.pipeline_versions:
+            raise UserError("No pipeline version was created")
+
+        created_version = response.pipeline_versions[0]
+        logger.info(f"Created pipeline version: {created_version.id}")
+        return created_version.id
+
+    def add_step_secret(
+        self,
+        step_ref: str,
+        secret_name: str,
+        secret_ref: str,
+        pipeline_version_id: Optional[str] = None,
+    ) -> None:
+        """Add a secret to a specific pipeline step.
+
+        Args:
+            step_ref (str): The step reference (e.g., "step-0", "step-1").
+            secret_name (str): The name of the secret environment variable.
+            secret_ref (str): The secret reference path (e.g., "users/user123/secrets/my-api-key").
+            pipeline_version_id (Optional[str]): The pipeline version ID. If None, uses self.pipeline_version_id.
+
+        Note:
+            This is a convenience method. For production use, manage secrets via the config.yaml
+            orchestration spec and use the pipeline upload command.
+        """
+        raise NotImplementedError(
+            "Adding secrets to existing pipeline versions is not supported. "
+            "Please define step secrets in your config.yaml orchestration spec "
+            "and use 'clarifai pipeline upload' to create a new pipeline version."
+        )
+
+    def list_step_secrets(
+        self, step_ref: Optional[str] = None, pipeline_version_id: Optional[str] = None
+    ) -> Dict[str, Dict[str, str]]:
+        """List secrets configured for pipeline steps.
+
+        Args:
+            step_ref (Optional[str]): If provided, only return secrets for this step.
+            pipeline_version_id (Optional[str]): The pipeline version ID. If None, uses self.pipeline_version_id.
+
+        Returns:
+            Dict[str, Dict[str, str]]: Map of step references to their secrets.
+                Format: {step_ref: {secret_name: secret_path}}
+        """
+        version_data = self.get_pipeline_version(pipeline_version_id)
+        config = version_data.get("config", {})
+        step_version_secrets = config.get("step_version_secrets", {})
+
+        if step_ref:
+            # Return only the specified step's secrets
+            return {step_ref: step_version_secrets.get(step_ref, {}).get("secrets", {})}
+
+        # Return all step secrets
+        result = {}
+        for step, config in step_version_secrets.items():
+            result[step] = config.get("secrets", {})
+        return result
diff --git a/clarifai/runners/pipelines/pipeline_builder.py b/clarifai/runners/pipelines/pipeline_builder.py
index e22c08a6..055bed0c 100644
--- a/clarifai/runners/pipelines/pipeline_builder.py
+++ b/clarifai/runners/pipelines/pipeline_builder.py
@@ -202,6 +202,13 @@ def prepare_lockfile_with_step_versions(self) -> Dict[str, Any]:
             }
         }
 
+        # Include step_version_secrets if present in config
+        step_version_secrets = orchestration_spec.get("step_version_secrets", {})
+        if step_version_secrets:
+            lockfile_data["pipeline"]["orchestration_spec"]["step_version_secrets"] = (
+                step_version_secrets
+            )
+
         return lockfile_data
 
     def update_lockfile_with_pipeline_info(
@@ -246,6 +253,13 @@ def generate_lockfile_data(
             }
         }
 
+        # Include step_version_secrets if present in config
+        step_version_secrets = orchestration_spec.get("step_version_secrets", {})
+        if step_version_secrets:
+            lockfile_data["pipeline"]["orchestration_spec"]["step_version_secrets"] = (
+                step_version_secrets
+            )
+
         return lockfile_data
 
     def save_lockfile(self, lockfile_data: Dict[str, Any], lockfile_path: str = None) -> None:
@@ -362,6 +376,33 @@ def _get_version_from_config_lock(self, step_name: str) -> str:
 
         return None
 
+    def _add_step_version_secrets(
+        self, pipeline_version: resources_pb2.PipelineVersion, step_version_secrets: Dict[str, Any]
+    ) -> None:
+        """Add step_version_secrets to the pipeline version config.
+
+        Args:
+            pipeline_version: The PipelineVersion proto to update
+            step_version_secrets: Dictionary mapping step references to their secret configs
+                                 Format: {step_ref: {secrets: {secret_name: secret_path}}}
+        """
+        logger.info(f"Adding step version secrets for {len(step_version_secrets)} steps")
+
+        for step_ref, step_config in step_version_secrets.items():
+            secrets = step_config.get("secrets", {})
+            if not secrets:
+                logger.warning(f"No secrets found for step {step_ref}, skipping")
+                continue
+
+            # Create StepSecretConfig proto
+            step_secret_config = resources_pb2.StepSecretConfig()
+            for secret_name, secret_ref in secrets.items():
+                step_secret_config.secrets[secret_name] = secret_ref
+                logger.info(f"Added secret {secret_name} for step {step_ref}")
+
+            # Add to pipeline version config
+            pipeline_version.config.step_version_secrets[step_ref].CopyFrom(step_secret_config)
+
     def create_pipeline(self) -> tuple[bool, str]:
         """Create the pipeline using PostPipelines RPC.
 
@@ -404,6 +445,11 @@ def create_pipeline(self) -> tuple[bool, str]:
             )
             pipeline_version.orchestration_spec.CopyFrom(orchestration_spec_proto)
 
+            # Add step_version_secrets if present in config
+            step_version_secrets = orchestration_spec.get("step_version_secrets", {})
+            if step_version_secrets:
+                self._add_step_version_secrets(pipeline_version, step_version_secrets)
+
             pipeline.pipeline_version.CopyFrom(pipeline_version)
 
             # Make the RPC call
diff --git a/tests/test_pipeline_client.py b/tests/test_pipeline_client.py
index 8dab4fbe..3fa2f080 100644
--- a/tests/test_pipeline_client.py
+++ b/tests/test_pipeline_client.py
@@ -329,3 +329,184 @@ def test_display_new_logs_pagination(self, mock_init):
         next_page = pipeline._display_new_logs('test-run-123', seen_logs_error, current_page=3)
         assert next_page == 3
         assert len(seen_logs_error) == 0
+
+    @patch('clarifai.client.pipeline.BaseClient.__init__')
+    def test_get_pipeline_version_with_step_secrets(self, mock_init):
+        """Test getting pipeline version with step secrets."""
+        mock_init.return_value = None
+
+        pipeline = Pipeline(
+            pipeline_id='test-pipeline',
+            pipeline_version_id='test-version-123',
+            user_id='test-user',
+            app_id='test-app',
+            pat='test-pat',
+        )
+
+        # Mock the required attributes
+        pipeline.user_app_id = resources_pb2.UserAppIDSet(user_id="test-user", app_id="test-app")
+        pipeline.STUB = Mock()
+        pipeline.auth_helper = Mock()
+        pipeline.auth_helper.metadata = []
+
+        # Create mock response with step secrets
+        mock_response = Mock()
+        mock_response.status.code = status_code_pb2.StatusCode.SUCCESS
+
+        pipeline_version = resources_pb2.PipelineVersion()
+        pipeline_version.id = 'test-version-123'
+
+        # Add step secrets
+        step_secret_config = resources_pb2.StepSecretConfig()
+        step_secret_config.secrets['API_KEY'] = 'users/test-user/secrets/my-api-key'
+        step_secret_config.secrets['DB_PASSWORD'] = 'users/test-user/secrets/db-secret'
+        pipeline_version.config.step_version_secrets['step-0'].CopyFrom(step_secret_config)
+
+        mock_response.pipeline_version = pipeline_version
+        pipeline.STUB.GetPipelineVersion.return_value = mock_response
+
+        # Execute
+        result = pipeline.get_pipeline_version()
+
+        # Verify
+        assert result['id'] == 'test-version-123'
+        assert 'config' in result
+        assert 'step_version_secrets' in result['config']
+        assert 'step-0' in result['config']['step_version_secrets']
+        secrets = result['config']['step_version_secrets']['step-0']['secrets']
+        assert secrets['API_KEY'] == 'users/test-user/secrets/my-api-key'
+        assert secrets['DB_PASSWORD'] == 'users/test-user/secrets/db-secret'
+
+    @patch('clarifai.client.pipeline.BaseClient.__init__')
+    def test_create_pipeline_version_with_step_secrets(self, mock_init):
+        """Test creating pipeline version with step secrets."""
+        mock_init.return_value = None
+
+        pipeline = Pipeline(
+            pipeline_id='test-pipeline',
+            user_id='test-user',
+            app_id='test-app',
+            pat='test-pat',
+        )
+
+        # Mock the required attributes
+        pipeline.user_app_id = resources_pb2.UserAppIDSet(user_id="test-user", app_id="test-app")
+        pipeline.STUB = Mock()
+        pipeline.auth_helper = Mock()
+        pipeline.auth_helper.metadata = []
+
+        # Create mock response
+        mock_response = Mock()
+        mock_response.status.code = status_code_pb2.StatusCode.SUCCESS
+        created_version = resources_pb2.PipelineVersion()
+        created_version.id = 'new-version-456'
+        mock_response.pipeline_versions = [created_version]
+        pipeline.STUB.PatchPipelineVersions.return_value = mock_response
+
+        # Execute
+        orchestration_spec = {
+            "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates: []
+"""
+        }
+        step_version_secrets = {
+            "step-0": {"API_KEY": "users/test-user/secrets/my-api-key"},
+            "step-1": {"EMAIL_TOKEN": "users/test-user/secrets/email-token"},
+        }
+
+        version_id = pipeline.create_pipeline_version(
+            orchestration_spec=orchestration_spec,
+            step_version_secrets=step_version_secrets,
+            description="Test pipeline version",
+        )
+
+        # Verify
+        assert version_id == 'new-version-456'
+        pipeline.STUB.PatchPipelineVersions.assert_called_once()
+
+        # Verify the request has step secrets
+        call_args = pipeline.STUB.PatchPipelineVersions.call_args
+        request = call_args[0][0]
+        assert len(request.pipeline_versions) == 1
+        pv = request.pipeline_versions[0]
+        assert 'step-0' in pv.config.step_version_secrets
+        assert 'step-1' in pv.config.step_version_secrets
+        assert (
+            pv.config.step_version_secrets['step-0'].secrets['API_KEY']
+            == 'users/test-user/secrets/my-api-key'
+        )
+
+    @patch('clarifai.client.pipeline.BaseClient.__init__')
+    def test_list_step_secrets(self, mock_init):
+        """Test listing step secrets for a pipeline version."""
+        mock_init.return_value = None
+
+        pipeline = Pipeline(
+            pipeline_id='test-pipeline',
+            pipeline_version_id='test-version-123',
+            user_id='test-user',
+            app_id='test-app',
+            pat='test-pat',
+        )
+
+        # Mock the required attributes
+        pipeline.user_app_id = resources_pb2.UserAppIDSet(user_id="test-user", app_id="test-app")
+        pipeline.STUB = Mock()
+        pipeline.auth_helper = Mock()
+        pipeline.auth_helper.metadata = []
+
+        # Create mock response with step secrets
+        mock_response = Mock()
+        mock_response.status.code = status_code_pb2.StatusCode.SUCCESS
+
+        pipeline_version = resources_pb2.PipelineVersion()
+        pipeline_version.id = 'test-version-123'
+
+        # Add step secrets for multiple steps
+        step_secret_config_0 = resources_pb2.StepSecretConfig()
+        step_secret_config_0.secrets['API_KEY'] = 'users/test-user/secrets/my-api-key'
+        pipeline_version.config.step_version_secrets['step-0'].CopyFrom(step_secret_config_0)
+
+        step_secret_config_1 = resources_pb2.StepSecretConfig()
+        step_secret_config_1.secrets['EMAIL_TOKEN'] = 'users/test-user/secrets/email-token'
+        pipeline_version.config.step_version_secrets['step-1'].CopyFrom(step_secret_config_1)
+
+        mock_response.pipeline_version = pipeline_version
+        pipeline.STUB.GetPipelineVersion.return_value = mock_response
+
+        # Test listing all secrets
+        all_secrets = pipeline.list_step_secrets()
+        assert 'step-0' in all_secrets
+        assert 'step-1' in all_secrets
+        assert all_secrets['step-0']['API_KEY'] == 'users/test-user/secrets/my-api-key'
+        assert all_secrets['step-1']['EMAIL_TOKEN'] == 'users/test-user/secrets/email-token'
+
+        # Test listing secrets for specific step
+        step0_secrets = pipeline.list_step_secrets(step_ref='step-0')
+        assert 'step-0' in step0_secrets
+        assert 'step-1' not in step0_secrets
+        assert step0_secrets['step-0']['API_KEY'] == 'users/test-user/secrets/my-api-key'
+
+    @patch('clarifai.client.pipeline.BaseClient.__init__')
+    def test_add_step_secret_not_implemented(self, mock_init):
+        """Test that add_step_secret raises NotImplementedError."""
+        mock_init.return_value = None
+
+        pipeline = Pipeline(
+            pipeline_id='test-pipeline',
+            pipeline_version_id='test-version-123',
+            user_id='test-user',
+            app_id='test-app',
+            pat='test-pat',
+        )
+
+        # Verify it raises NotImplementedError
+        with pytest.raises(
+            NotImplementedError,
+            match="Adding secrets to existing pipeline versions is not supported",
+        ):
+            pipeline.add_step_secret('step-0', 'API_KEY', 'users/test-user/secrets/my-api-key')
diff --git a/tests/test_pipeline_step_secrets.py b/tests/test_pipeline_step_secrets.py
new file mode 100644
index 00000000..1a4881f4
--- /dev/null
+++ b/tests/test_pipeline_step_secrets.py
@@ -0,0 +1,300 @@
+"""Tests for pipeline step secrets functionality."""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+import yaml
+from clarifai_grpc.grpc.api import resources_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2
+
+from clarifai.runners.pipelines.pipeline_builder import PipelineBuilder
+
+
+class TestPipelineStepSecrets:
+    """Test cases for pipeline step secrets feature."""
+
+    def test_pipeline_builder_with_step_secrets(self):
+        """Test that PipelineBuilder correctly handles step_version_secrets from config."""
+        # Create a temporary config file with step secrets
+        config = {
+            "pipeline": {
+                "id": "test-pipeline",
+                "user_id": "test-user",
+                "app_id": "test-app",
+                "step_directories": ["step1", "step2"],
+                "orchestration_spec": {
+                    "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates:
+  - name: sequence
+    steps:
+      - - name: step-0
+          templateRef:
+            name: users/test-user/apps/test-app/pipeline_steps/step1
+            template: users/test-user/apps/test-app/pipeline_steps/step1
+""",
+                    "step_version_secrets": {
+                        "step-0": {
+                            "secrets": {
+                                "API_KEY": "users/test-user/secrets/my-api-key",
+                                "DB_PASSWORD": "users/test-user/secrets/db-secret",
+                            }
+                        },
+                        "step-1": {
+                            "secrets": {"EMAIL_TOKEN": "users/test-user/secrets/email-token"}
+                        },
+                    },
+                },
+            }
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+            yaml.dump(config, f)
+            config_path = f.name
+
+        try:
+            # Initialize builder
+            with patch('clarifai.runners.pipelines.pipeline_builder.BaseClient'):
+                builder = PipelineBuilder(config_path)
+
+            # Verify config was loaded with step secrets
+            assert "step_version_secrets" in builder.config["pipeline"]["orchestration_spec"]
+            step_secrets = builder.config["pipeline"]["orchestration_spec"]["step_version_secrets"]
+            assert "step-0" in step_secrets
+            assert "step-1" in step_secrets
+            assert (
+                step_secrets["step-0"]["secrets"]["API_KEY"]
+                == "users/test-user/secrets/my-api-key"
+            )
+        finally:
+            Path(config_path).unlink()
+
+    def test_add_step_version_secrets_to_pipeline_version(self):
+        """Test the _add_step_version_secrets helper method."""
+        config = {
+            "pipeline": {
+                "id": "test-pipeline",
+                "user_id": "test-user",
+                "app_id": "test-app",
+                "step_directories": [],
+                "orchestration_spec": {
+                    "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates: []
+"""
+                },
+            }
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+            yaml.dump(config, f)
+            config_path = f.name
+
+        try:
+            with patch('clarifai.runners.pipelines.pipeline_builder.BaseClient'):
+                builder = PipelineBuilder(config_path)
+
+            # Create pipeline version proto
+            pipeline_version = resources_pb2.PipelineVersion()
+
+            # Define step secrets
+            step_version_secrets = {
+                "step-0": {
+                    "secrets": {
+                        "API_KEY": "users/test-user/secrets/my-api-key",
+                        "DB_PASSWORD": "users/test-user/secrets/db-secret",
+                    }
+                },
+                "step-1": {"secrets": {"EMAIL_TOKEN": "users/test-user/secrets/email-token"}},
+            }
+
+            # Call the helper method
+            builder._add_step_version_secrets(pipeline_version, step_version_secrets)
+
+            # Verify secrets were added
+            assert "step-0" in pipeline_version.config.step_version_secrets
+            assert "step-1" in pipeline_version.config.step_version_secrets
+
+            step0_secrets = pipeline_version.config.step_version_secrets["step-0"].secrets
+            assert step0_secrets["API_KEY"] == "users/test-user/secrets/my-api-key"
+            assert step0_secrets["DB_PASSWORD"] == "users/test-user/secrets/db-secret"
+
+            step1_secrets = pipeline_version.config.step_version_secrets["step-1"].secrets
+            assert step1_secrets["EMAIL_TOKEN"] == "users/test-user/secrets/email-token"
+        finally:
+            Path(config_path).unlink()
+
+    def test_lockfile_includes_step_secrets(self):
+        """Test that lockfile generation includes step_version_secrets."""
+        config = {
+            "pipeline": {
+                "id": "test-pipeline",
+                "user_id": "test-user",
+                "app_id": "test-app",
+                "step_directories": [],
+                "orchestration_spec": {
+                    "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates: []
+""",
+                    "step_version_secrets": {
+                        "step-0": {"secrets": {"API_KEY": "users/test-user/secrets/my-api-key"}}
+                    },
+                },
+            }
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+            yaml.dump(config, f)
+            config_path = f.name
+
+        try:
+            with patch('clarifai.runners.pipelines.pipeline_builder.BaseClient'):
+                builder = PipelineBuilder(config_path)
+
+            # Generate lockfile data
+            lockfile_data = builder.prepare_lockfile_with_step_versions()
+
+            # Verify step secrets are in lockfile
+            assert "orchestration_spec" in lockfile_data["pipeline"]
+            assert "step_version_secrets" in lockfile_data["pipeline"]["orchestration_spec"]
+            step_secrets = lockfile_data["pipeline"]["orchestration_spec"]["step_version_secrets"]
+            assert "step-0" in step_secrets
+            assert (
+                step_secrets["step-0"]["secrets"]["API_KEY"]
+                == "users/test-user/secrets/my-api-key"
+            )
+        finally:
+            Path(config_path).unlink()
+
+    def test_create_pipeline_with_step_secrets(self):
+        """Test full pipeline creation with step secrets."""
+        config = {
+            "pipeline": {
+                "id": "test-pipeline",
+                "user_id": "test-user",
+                "app_id": "test-app",
+                "step_directories": [],
+                "orchestration_spec": {
+                    "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates:
+  - name: sequence
+    steps:
+      - - name: step-0
+          templateRef:
+            name: users/test-user/apps/test-app/pipeline_steps/step1/versions/v1
+            template: users/test-user/apps/test-app/pipeline_steps/step1/versions/v1
+""",
+                    "step_version_secrets": {
+                        "step-0": {
+                            "secrets": {
+                                "API_KEY": "users/test-user/secrets/my-api-key",
+                                "DB_PASSWORD": "users/test-user/secrets/db-secret",
+                            }
+                        }
+                    },
+                },
+            }
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+            yaml.dump(config, f)
+            config_path = f.name
+
+        try:
+            with patch('clarifai.runners.pipelines.pipeline_builder.BaseClient') as mock_client:
+                builder = PipelineBuilder(config_path)
+
+                # Mock the STUB and response
+                mock_stub = Mock()
+                mock_response = Mock()
+                mock_response.status.code = status_code_pb2.SUCCESS
+
+                mock_pipeline = resources_pb2.Pipeline()
+                mock_pipeline.id = "test-pipeline"
+                mock_pipeline_version = resources_pb2.PipelineVersion()
+                mock_pipeline_version.id = "test-version-123"
+                mock_pipeline.pipeline_version.CopyFrom(mock_pipeline_version)
+
+                mock_response.pipelines = [mock_pipeline]
+                mock_stub.PostPipelines.return_value = mock_response
+
+                builder.client.STUB = mock_stub
+                builder.client.user_app_id = resources_pb2.UserAppIDSet(
+                    user_id="test-user", app_id="test-app"
+                )
+
+                # Create the pipeline
+                success, version_id = builder.create_pipeline()
+
+                # Verify success
+                assert success is True
+                assert version_id == "test-version-123"
+
+                # Verify PostPipelines was called
+                mock_stub.PostPipelines.assert_called_once()
+
+                # Verify the request includes step secrets
+                call_args = mock_stub.PostPipelines.call_args
+                request = call_args[0][0]
+                assert len(request.pipelines) == 1
+                pipeline = request.pipelines[0]
+                assert pipeline.pipeline_version.HasField("config")
+                assert "step-0" in pipeline.pipeline_version.config.step_version_secrets
+                step_secrets = pipeline.pipeline_version.config.step_version_secrets[
+                    "step-0"
+                ].secrets
+                assert step_secrets["API_KEY"] == "users/test-user/secrets/my-api-key"
+                assert step_secrets["DB_PASSWORD"] == "users/test-user/secrets/db-secret"
+        finally:
+            Path(config_path).unlink()
+
+    def test_empty_step_secrets_handling(self):
+        """Test that empty step secrets are handled gracefully."""
+        config = {
+            "pipeline": {
+                "id": "test-pipeline",
+                "user_id": "test-user",
+                "app_id": "test-app",
+                "step_directories": [],
+                "orchestration_spec": {
+                    "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates: []
+"""
+                },
+            }
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+            yaml.dump(config, f)
+            config_path = f.name
+
+        try:
+            with patch('clarifai.runners.pipelines.pipeline_builder.BaseClient'):
+                builder = PipelineBuilder(config_path)
+
+            # Generate lockfile data without secrets
+            lockfile_data = builder.prepare_lockfile_with_step_versions()
+
+            # Verify no step_version_secrets in lockfile when not provided
+            assert "step_version_secrets" not in lockfile_data["pipeline"]["orchestration_spec"]
+        finally:
+            Path(config_path).unlink()

From 348414fd08116d3d7f646d70dfd10f17f341dbbe Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 30 Oct 2025 06:45:49 +0000
Subject: [PATCH 3/6] Add pipeline step secrets documentation and usage guide

Co-authored-by: ydixit-clarifai <219138670+ydixit-clarifai@users.noreply.github.com>
---
 docs/pipeline_step_secrets.md | 230 ++++++++++++++++++++++++++++++++++
 1 file changed, 230 insertions(+)
 create mode 100644 docs/pipeline_step_secrets.md

diff --git a/docs/pipeline_step_secrets.md b/docs/pipeline_step_secrets.md
new file mode 100644
index 00000000..1f88d387
--- /dev/null
+++ b/docs/pipeline_step_secrets.md
@@ -0,0 +1,230 @@
+# Pipeline Step Secrets Usage Guide
+
+This guide explains how to use pipeline step secrets in the Clarifai Python SDK.
+
+## Overview
+
+Pipeline step secrets allow different pipeline steps to access distinct sets of secrets with step-level isolation. Each step can have its own set of secret environment variables that are mounted securely at runtime.
+
+## Configuration Format
+
+Define step secrets in your pipeline `config.yaml` file within the `orchestration_spec` section:
+
+```yaml
+pipeline:
+  id: "my-pipeline"
+  user_id: "user123"
+  app_id: "app456"
+  step_directories:
+    - step1
+    - step2
+  orchestration_spec:
+    argo_orchestration_spec: |
+      apiVersion: argoproj.io/v1alpha1
+      kind: Workflow
+      spec:
+        entrypoint: sequence
+        templates:
+        - name: sequence
+          steps:
+            - - name: step-0
+                templateRef:
+                  name: users/user123/apps/app456/pipeline_steps/step1
+                  template: users/user123/apps/app456/pipeline_steps/step1
+            - - name: step-1
+                templateRef:
+                  name: users/user123/apps/app456/pipeline_steps/step2
+                  template: users/user123/apps/app456/pipeline_steps/step2
+    
+    # Define secrets for each step
+    step_version_secrets:
+      step-0:
+        secrets:
+          API_KEY: users/user123/secrets/my-api-key
+          DB_PASSWORD: users/user123/secrets/db-secret
+      step-1:
+        secrets:
+          EMAIL_TOKEN: users/user123/secrets/email-token
+```
+
+## Secret Reference Format
+
+Secret references follow the pattern: `users/{user_id}/secrets/{secret_name}`
+
+- The secrets must already exist in your Clarifai account
+- Only references are stored in the config; actual values are injected at runtime
+- Each step can only access its explicitly configured secrets
+
+## CLI Usage
+
+### 1. Initialize a Pipeline with Secrets
+
+```bash
+clarifai pipeline init my-pipeline
+```
+
+Then edit the generated `config.yaml` to add the `step_version_secrets` section as shown above.
+
+### 2. Upload Pipeline with Secrets
+
+```bash
+clarifai pipeline upload config.yaml
+```
+
+This will:
+1. Upload all pipeline steps from `step_directories`
+2. Create the pipeline with step secrets configuration
+3. Generate a `config-lock.yaml` file that includes the secrets config
+
+### 3. Run Pipeline
+
+```bash
+clarifai pipeline run \
+  --pipeline_id my-pipeline \
+  --user_id user123 \
+  --app_id app456 \
+  --compute_cluster_id cluster-id \
+  --nodepool_id nodepool-id
+```
+
+## Python SDK Usage
+
+### Get Pipeline Version with Secrets
+
+```python
+from clarifai.client.pipeline import Pipeline
+
+# Initialize pipeline
+pipeline = Pipeline(
+    pipeline_id="my-pipeline",
+    pipeline_version_id="version-123",
+    user_id="user123",
+    app_id="app456",
+    pat="your-pat"
+)
+
+# Get pipeline version details including secrets
+version = pipeline.get_pipeline_version()
+print(version['config']['step_version_secrets'])
+```
+
+### List Step Secrets
+
+```python
+# List all step secrets
+all_secrets = pipeline.list_step_secrets()
+print(all_secrets)
+# Output: {
+#   'step-0': {'API_KEY': 'users/user123/secrets/my-api-key', ...},
+#   'step-1': {'EMAIL_TOKEN': 'users/user123/secrets/email-token'}
+# }
+
+# List secrets for specific step
+step0_secrets = pipeline.list_step_secrets(step_ref='step-0')
+print(step0_secrets)
+# Output: {
+#   'step-0': {'API_KEY': 'users/user123/secrets/my-api-key', ...}
+# }
+```
+
+### Create Pipeline Version with Secrets
+
+```python
+from clarifai.client.pipeline import Pipeline
+
+pipeline = Pipeline(
+    pipeline_id="my-pipeline",
+    user_id="user123",
+    app_id="app456",
+    pat="your-pat"
+)
+
+# Define orchestration spec
+orchestration_spec = {
+    "argo_orchestration_spec": """
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+spec:
+  entrypoint: sequence
+  templates:
+  - name: sequence
+    steps:
+      - - name: step-0
+          templateRef:
+            name: users/user123/apps/app456/pipeline_steps/step1/versions/v1
+            template: users/user123/apps/app456/pipeline_steps/step1/versions/v1
+"""
+}
+
+# Define step secrets
+step_version_secrets = {
+    "step-0": {
+        "API_KEY": "users/user123/secrets/my-api-key",
+        "DB_PASSWORD": "users/user123/secrets/db-secret"
+    }
+}
+
+# Create new version
+version_id = pipeline.create_pipeline_version(
+    orchestration_spec=orchestration_spec,
+    step_version_secrets=step_version_secrets,
+    description="Pipeline version with step secrets"
+)
+print(f"Created version: {version_id}")
+```
+
+## Security Considerations
+
+- **Step-Level Isolation**: Each step only accesses explicitly configured secrets
+- **Reference-Only Storage**: Only secret references are stored in config files
+- **Runtime Injection**: Actual secret values are injected securely at runtime
+- **Kubernetes Integration**: Backend uses Kubernetes SecretKeyRef for secure mounting
+- **No Value Leakage**: Secret values are never logged or exposed in API responses
+
+## Best Practices
+
+1. **Manage Secrets via Config File**: Always define secrets in `config.yaml` rather than trying to add them programmatically
+2. **Use Descriptive Names**: Give secrets clear, descriptive names like `API_KEY`, `DB_PASSWORD`
+3. **Minimize Access**: Only give each step the secrets it needs
+4. **Version Control**: Use `config-lock.yaml` to track which secrets are configured for each pipeline version
+5. **Regular Rotation**: Rotate secrets regularly and update references in config
+
+## Troubleshooting
+
+### Secret Not Found Error
+
+If you get an error about a secret not existing:
+- Verify the secret exists in your Clarifai account
+- Check that the secret reference path is correct: `users/{user_id}/secrets/{secret_name}`
+- Ensure you have permission to access the secret
+
+### Secret Not Available in Step
+
+If a step can't access a secret:
+- Verify the step reference matches the one in your Argo orchestration spec
+- Check that the secret is defined for that specific step in `step_version_secrets`
+- Ensure the pipeline was uploaded after adding the secrets configuration
+
+## Example Project Structure
+
+```
+my-pipeline/
+├── config.yaml              # Pipeline config with step_version_secrets
+├── config-lock.yaml         # Generated lockfile with versions
+├── step1/
+│   ├── config.yaml
+│   ├── requirements.txt
+│   └── 1/
+│       └── pipeline_step.py
+└── step2/
+    ├── config.yaml
+    ├── requirements.txt
+    └── 1/
+        └── pipeline_step.py
+```
+
+## Related Documentation
+
+- [Clarifai Secrets Management](https://docs.clarifai.com/api-guide/secrets)
+- [Pipeline Orchestration](https://docs.clarifai.com/api-guide/pipelines)
+- [Pipeline Steps](https://docs.clarifai.com/api-guide/pipeline-steps)

From e10c64a53f33cdfdabc14b0769c4674383e9185e Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 30 Oct 2025 06:49:36 +0000
Subject: [PATCH 4/6] Clarify Python API format for step secrets in
 documentation

Add note explaining that Python API uses simplified format without 'secrets' wrapper compared to YAML configuration

Co-authored-by: ydixit-clarifai <219138670+ydixit-clarifai@users.noreply.github.com>
---
 docs/pipeline_step_secrets.md | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/docs/pipeline_step_secrets.md b/docs/pipeline_step_secrets.md
index 1f88d387..e87dfbf7 100644
--- a/docs/pipeline_step_secrets.md
+++ b/docs/pipeline_step_secrets.md
@@ -89,6 +89,8 @@ clarifai pipeline run \
 
 ## Python SDK Usage
 
+**Note on Format**: The Python SDK uses a simplified format for step secrets compared to the YAML configuration. In YAML, secrets are nested under a `secrets` key, but in the Python API you provide them directly as a dictionary.
+
 ### Get Pipeline Version with Secrets
 
 ```python
@@ -156,7 +158,9 @@ spec:
 """
 }
 
-# Define step secrets
+# Define step secrets (simplified format for Python API)
+# Note: The Python API accepts a simplified format where you provide
+# the secrets directly without the 'secrets' wrapper key
 step_version_secrets = {
     "step-0": {
         "API_KEY": "users/user123/secrets/my-api-key",

From 5a9a7fe697e999d567a483c18b3ff5f951e3d37c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 30 Oct 2025 06:56:35 +0000
Subject: [PATCH 5/6] Improve logging security for step secrets handling

- Change info logging to debug level to reduce verbosity
- Rename 'secrets' variable to 'secret_refs' for clarity
- Add comments explaining that only references (not values) are handled
- Remove per-secret logging to minimize exposure surface
- Address CodeQL security alerts (false positives documented)

Co-authored-by: ydixit-clarifai <219138670+ydixit-clarifai@users.noreply.github.com>
---
 clarifai/runners/pipelines/pipeline_builder.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/clarifai/runners/pipelines/pipeline_builder.py b/clarifai/runners/pipelines/pipeline_builder.py
index 055bed0c..91b9d9c1 100644
--- a/clarifai/runners/pipelines/pipeline_builder.py
+++ b/clarifai/runners/pipelines/pipeline_builder.py
@@ -386,22 +386,24 @@ def _add_step_version_secrets(
             step_version_secrets: Dictionary mapping step references to their secret configs
                                  Format: {step_ref: {secrets: {secret_name: secret_path}}}
         """
-        logger.info(f"Adding step version secrets for {len(step_version_secrets)} steps")
+        logger.debug(f"Processing step version secrets for {len(step_version_secrets)} steps")
 
         for step_ref, step_config in step_version_secrets.items():
-            secrets = step_config.get("secrets", {})
-            if not secrets:
-                logger.warning(f"No secrets found for step {step_ref}, skipping")
+            # Note: 'secret_refs' contains only secret reference paths (not actual values)
+            # Secret references are like "users/user123/secrets/my-api-key" 
+            secret_refs = step_config.get("secrets", {})
+            if not secret_refs:
+                logger.debug(f"No secret references found for step {step_ref}, skipping")
                 continue
 
             # Create StepSecretConfig proto
             step_secret_config = resources_pb2.StepSecretConfig()
-            for secret_name, secret_ref in secrets.items():
+            for secret_name, secret_ref in secret_refs.items():
                 step_secret_config.secrets[secret_name] = secret_ref
-                logger.info(f"Added secret {secret_name} for step {step_ref}")
 
             # Add to pipeline version config
             pipeline_version.config.step_version_secrets[step_ref].CopyFrom(step_secret_config)
+            logger.debug(f"Configured secret references for step {step_ref}")
 
     def create_pipeline(self) -> tuple[bool, str]:
         """Create the pipeline using PostPipelines RPC.

From 9cdadbdb5de1bfcc13c4ed45f9097008fb78ac74 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 30 Oct 2025 08:14:14 +0000
Subject: [PATCH 6/6] Move step_version_secrets to pipeline config section

Per review feedback, step_version_secrets should be part of the pipeline config section (matching the proto structure PipelineVersion.config.step_version_secrets) rather than nested inside orchestration_spec.

Changes:
- Move step_version_secrets from orchestration_spec to config section in YAML structure
- Update pipeline_builder.py to read from config section
- Update lockfile generation to preserve secrets in config section
- Update templates and documentation to show correct structure
- Update all tests to use new config structure

The new YAML structure is:
```yaml
pipeline:
  orchestration_spec:
    argo_orchestration_spec: |
      ...
  config:
    step_version_secrets:
      step-0:
        secrets:
          API_KEY: users/user123/secrets/my-api-key
```

Co-authored-by: ydixit-clarifai <219138670+ydixit-clarifai@users.noreply.github.com>
---
 clarifai/cli/templates/pipeline_templates.py  | 19 +++++++------
 .../runners/pipelines/pipeline_builder.py     | 28 ++++++++++---------
 docs/pipeline_step_secrets.md                 |  7 +++--
 tests/test_pipeline_step_secrets.py           | 26 +++++++++++------
 4 files changed, 46 insertions(+), 34 deletions(-)

diff --git a/clarifai/cli/templates/pipeline_templates.py b/clarifai/cli/templates/pipeline_templates.py
index 513df4fc..e0cedb88 100644
--- a/clarifai/cli/templates/pipeline_templates.py
+++ b/clarifai/cli/templates/pipeline_templates.py
@@ -45,15 +45,16 @@ def get_pipeline_config_template(
         - name: sequence
           steps:
 {steps_yaml}
-    # Optional: Define secrets for pipeline steps
-    # step_version_secrets:
-    #   step-0:
-    #     secrets:
-    #       API_KEY: users/{user_id}/secrets/my-api-key
-    #       DB_PASSWORD: users/{user_id}/secrets/db-secret
-    #   step-1:
-    #     secrets:
-    #       EMAIL_TOKEN: users/{user_id}/secrets/email-token
+  # Optional: Define secrets for pipeline steps
+  # config:
+  #   step_version_secrets:
+  #     step-0:
+  #       secrets:
+  #         API_KEY: users/{user_id}/secrets/my-api-key
+  #         DB_PASSWORD: users/{user_id}/secrets/db-secret
+  #     step-1:
+  #       secrets:
+  #         EMAIL_TOKEN: users/{user_id}/secrets/email-token
 """
 
 
diff --git a/clarifai/runners/pipelines/pipeline_builder.py b/clarifai/runners/pipelines/pipeline_builder.py
index 91b9d9c1..7438e759 100644
--- a/clarifai/runners/pipelines/pipeline_builder.py
+++ b/clarifai/runners/pipelines/pipeline_builder.py
@@ -202,12 +202,12 @@ def prepare_lockfile_with_step_versions(self) -> Dict[str, Any]:
             }
         }
 
-        # Include step_version_secrets if present in config
-        step_version_secrets = orchestration_spec.get("step_version_secrets", {})
+        # Include step_version_secrets if present in pipeline config (not orchestration_spec)
+        step_version_secrets = pipeline_config.get("config", {}).get("step_version_secrets", {})
         if step_version_secrets:
-            lockfile_data["pipeline"]["orchestration_spec"]["step_version_secrets"] = (
-                step_version_secrets
-            )
+            if "config" not in lockfile_data["pipeline"]:
+                lockfile_data["pipeline"]["config"] = {}
+            lockfile_data["pipeline"]["config"]["step_version_secrets"] = step_version_secrets
 
         return lockfile_data
 
@@ -253,12 +253,12 @@ def generate_lockfile_data(
             }
         }
 
-        # Include step_version_secrets if present in config
-        step_version_secrets = orchestration_spec.get("step_version_secrets", {})
+        # Include step_version_secrets if present in pipeline config (not orchestration_spec)
+        step_version_secrets = pipeline_config.get("config", {}).get("step_version_secrets", {})
         if step_version_secrets:
-            lockfile_data["pipeline"]["orchestration_spec"]["step_version_secrets"] = (
-                step_version_secrets
-            )
+            if "config" not in lockfile_data["pipeline"]:
+                lockfile_data["pipeline"]["config"] = {}
+            lockfile_data["pipeline"]["config"]["step_version_secrets"] = step_version_secrets
 
         return lockfile_data
 
@@ -390,7 +390,7 @@ def _add_step_version_secrets(
 
         for step_ref, step_config in step_version_secrets.items():
             # Note: 'secret_refs' contains only secret reference paths (not actual values)
-            # Secret references are like "users/user123/secrets/my-api-key" 
+            # Secret references are like "users/user123/secrets/my-api-key"
             secret_refs = step_config.get("secrets", {})
             if not secret_refs:
                 logger.debug(f"No secret references found for step {step_ref}, skipping")
@@ -447,8 +447,10 @@ def create_pipeline(self) -> tuple[bool, str]:
             )
             pipeline_version.orchestration_spec.CopyFrom(orchestration_spec_proto)
 
-            # Add step_version_secrets if present in config
-            step_version_secrets = orchestration_spec.get("step_version_secrets", {})
+            # Add step_version_secrets if present in pipeline config (not orchestration_spec)
+            step_version_secrets = pipeline_config.get("config", {}).get(
+                "step_version_secrets", {}
+            )
             if step_version_secrets:
                 self._add_step_version_secrets(pipeline_version, step_version_secrets)
 
diff --git a/docs/pipeline_step_secrets.md b/docs/pipeline_step_secrets.md
index e87dfbf7..1e3540a9 100644
--- a/docs/pipeline_step_secrets.md
+++ b/docs/pipeline_step_secrets.md
@@ -8,7 +8,7 @@ Pipeline step secrets allow different pipeline steps to access distinct sets of
 
 ## Configuration Format
 
-Define step secrets in your pipeline `config.yaml` file within the `orchestration_spec` section:
+Define step secrets in your pipeline `config.yaml` file in a `config` section at the pipeline level:
 
 ```yaml
 pipeline:
@@ -35,8 +35,9 @@ pipeline:
                 templateRef:
                   name: users/user123/apps/app456/pipeline_steps/step2
                   template: users/user123/apps/app456/pipeline_steps/step2
-    
-    # Define secrets for each step
+  
+  # Define secrets for each step in the config section
+  config:
     step_version_secrets:
       step-0:
         secrets:
diff --git a/tests/test_pipeline_step_secrets.py b/tests/test_pipeline_step_secrets.py
index 1a4881f4..dc0d0137 100644
--- a/tests/test_pipeline_step_secrets.py
+++ b/tests/test_pipeline_step_secrets.py
@@ -37,6 +37,8 @@ def test_pipeline_builder_with_step_secrets(self):
             name: users/test-user/apps/test-app/pipeline_steps/step1
             template: users/test-user/apps/test-app/pipeline_steps/step1
 """,
+                },
+                "config": {
                     "step_version_secrets": {
                         "step-0": {
                             "secrets": {
@@ -61,9 +63,10 @@ def test_pipeline_builder_with_step_secrets(self):
             with patch('clarifai.runners.pipelines.pipeline_builder.BaseClient'):
                 builder = PipelineBuilder(config_path)
 
-            # Verify config was loaded with step secrets
-            assert "step_version_secrets" in builder.config["pipeline"]["orchestration_spec"]
-            step_secrets = builder.config["pipeline"]["orchestration_spec"]["step_version_secrets"]
+            # Verify config was loaded with step secrets in the config section
+            assert "config" in builder.config["pipeline"]
+            assert "step_version_secrets" in builder.config["pipeline"]["config"]
+            step_secrets = builder.config["pipeline"]["config"]["step_version_secrets"]
             assert "step-0" in step_secrets
             assert "step-1" in step_secrets
             assert (
@@ -147,6 +150,8 @@ def test_lockfile_includes_step_secrets(self):
   entrypoint: sequence
   templates: []
 """,
+                },
+                "config": {
                     "step_version_secrets": {
                         "step-0": {"secrets": {"API_KEY": "users/test-user/secrets/my-api-key"}}
                     },
@@ -165,10 +170,10 @@ def test_lockfile_includes_step_secrets(self):
             # Generate lockfile data
             lockfile_data = builder.prepare_lockfile_with_step_versions()
 
-            # Verify step secrets are in lockfile
-            assert "orchestration_spec" in lockfile_data["pipeline"]
-            assert "step_version_secrets" in lockfile_data["pipeline"]["orchestration_spec"]
-            step_secrets = lockfile_data["pipeline"]["orchestration_spec"]["step_version_secrets"]
+            # Verify step secrets are in lockfile config section
+            assert "config" in lockfile_data["pipeline"]
+            assert "step_version_secrets" in lockfile_data["pipeline"]["config"]
+            step_secrets = lockfile_data["pipeline"]["config"]["step_version_secrets"]
             assert "step-0" in step_secrets
             assert (
                 step_secrets["step-0"]["secrets"]["API_KEY"]
@@ -199,6 +204,8 @@ def test_create_pipeline_with_step_secrets(self):
             name: users/test-user/apps/test-app/pipeline_steps/step1/versions/v1
             template: users/test-user/apps/test-app/pipeline_steps/step1/versions/v1
 """,
+                },
+                "config": {
                     "step_version_secrets": {
                         "step-0": {
                             "secrets": {
@@ -294,7 +301,8 @@ def test_empty_step_secrets_handling(self):
             # Generate lockfile data without secrets
             lockfile_data = builder.prepare_lockfile_with_step_versions()
 
-            # Verify no step_version_secrets in lockfile when not provided
-            assert "step_version_secrets" not in lockfile_data["pipeline"]["orchestration_spec"]
+            # Verify no config or step_version_secrets in lockfile when not provided
+            # The config key should not exist if there are no secrets
+            assert "config" not in lockfile_data["pipeline"]
         finally:
             Path(config_path).unlink()