fix(serverless): inject runtime vars on update to survive template env overwrite

Extract _inject_runtime_template_vars() from _do_deploy so both initial
deploy and update() paths inject RUNPOD_API_KEY and FLASH_MODULE_PATH
into template.env. Without this, runtime vars set during _do_deploy
were silently dropped when update() overwrote the template env on
config drift. Also preserve explicit template.env entries when env
dict is empty on both sides.

Author: deanq

src/runpod_flash/core/resources/serverless.py

@@ -654,6 +654,41 @@ def _inject_template_env(self, key: str, value: str) -> None:
         if key not in existing_keys:
             self.template.env.append(KeyValuePair(key=key, value=value))
 
+    def _inject_runtime_template_vars(self) -> None:
+        """Inject runtime env vars into template.env without mutating self.env.
+
+        For QB endpoints making remote calls: injects RUNPOD_API_KEY.
+        For LB endpoints: injects FLASH_MODULE_PATH.
+
+        Called by both _do_deploy (initial) and update (env changes) so
+        runtime vars survive template updates.
+        """
+        env_dict = self.env or {}
+
+        if self.type == ServerlessType.QB:
+            if self._check_makes_remote_calls():
+                if "RUNPOD_API_KEY" not in env_dict:
+                    from runpod_flash.core.credentials import get_api_key
+
+                    api_key = get_api_key()
+                    if api_key:
+                        self._inject_template_env("RUNPOD_API_KEY", api_key)
+                        log.debug(
+                            f"{self.name}: Injected RUNPOD_API_KEY for remote calls "
+                            f"(makes_remote_calls=True)"
+                        )
+                    else:
+                        log.warning(
+                            f"{self.name}: makes_remote_calls=True but RUNPOD_API_KEY not set. "
+                            f"Remote calls to other endpoints will fail."
+                        )
+
+        elif self.type == ServerlessType.LB:
+            module_path = self._get_module_path()
+            if module_path and "FLASH_MODULE_PATH" not in env_dict:
+                self._inject_template_env("FLASH_MODULE_PATH", module_path)
+                log.debug(f"{self.name}: Injected FLASH_MODULE_PATH={module_path}")
+
     async def _do_deploy(self) -> "DeployableResource":
         """
         Deploys the serverless resource using the provided configuration.
@@ -669,37 +704,7 @@ async def _do_deploy(self) -> "DeployableResource":
                 log.debug(f"{self} exists")
                 return self
 
-            # Inject API key for queue-based endpoints that make remote calls.
-            # Injected into template.env (not self.env) to avoid false config drift.
-            if self.type == ServerlessType.QB:
-                makes_remote_calls = self._check_makes_remote_calls()
-
-                if makes_remote_calls:
-                    env_dict = self.env or {}
-                    if "RUNPOD_API_KEY" not in env_dict:
-                        from runpod_flash.core.credentials import get_api_key
-
-                        api_key = get_api_key()
-                        if api_key:
-                            self._inject_template_env("RUNPOD_API_KEY", api_key)
-                            log.debug(
-                                f"{self.name}: Injected RUNPOD_API_KEY for remote calls "
-                                f"(makes_remote_calls=True)"
-                            )
-                        else:
-                            log.warning(
-                                f"{self.name}: makes_remote_calls=True but RUNPOD_API_KEY not set. "
-                                f"Remote calls to other endpoints will fail."
-                            )
-
-            # Inject module path for load-balanced endpoints.
-            # Injected into template.env (not self.env) to avoid false config drift.
-            elif self.type == ServerlessType.LB:
-                env_dict = self.env or {}
-                module_path = self._get_module_path()
-                if module_path and "FLASH_MODULE_PATH" not in env_dict:
-                    self._inject_template_env("FLASH_MODULE_PATH", module_path)
-                    log.debug(f"{self.name}: Injected FLASH_MODULE_PATH={module_path}")
+            self._inject_runtime_template_vars()
 
             # Ensure network volume is deployed first
             await self._ensure_network_volume_deployed()
@@ -773,11 +778,25 @@ async def update(self, new_config: "ServerlessResource") -> "ServerlessResource"
                         # hasn't changed.  This lets the platform keep vars it
                         # injected (e.g. PORT, PORT_HEALTH on LB endpoints)
                         # and avoids a spurious rolling release.
+                        #
+                        # Also check template.env: if env is empty but the
+                        # caller provided explicit template env entries, those
+                        # must not be silently dropped.
                         env_unchanged = self.env == new_config.env
+                        has_explicit_template_env = (
+                            not new_config.env and new_config.template.env is not None
+                        )
+                        skip_env = env_unchanged and not has_explicit_template_env
+
+                        if not skip_env:
+                            # Inject runtime vars (RUNPOD_API_KEY, FLASH_MODULE_PATH)
+                            # so they survive the template env overwrite.
+                            new_config._inject_runtime_template_vars()
+
                         template_payload = self._build_template_update_payload(
                             new_config.template,
                             resolved_template_id,
-                            skip_env=env_unchanged,
+                            skip_env=skip_env,
                         )
                         await client.update_template(template_payload)
                         log.debug(

tests/unit/resources/test_serverless.py

@@ -1693,3 +1693,181 @@ async def test_update_includes_env_when_changed(self):
         assert mock_client.update_template.called
         template_payload = mock_client.update_template.call_args.args[0]
         assert "env" in template_payload
+
+    @pytest.mark.asyncio
+    async def test_update_injects_runtime_vars_when_env_changed(self):
+        """update() injects RUNPOD_API_KEY into template.env when env changed.
+
+        Without this, runtime-injected vars (set during _do_deploy) would be
+        lost when update() overwrites the template env.
+        """
+        old_resource = ServerlessEndpoint(
+            name="update-inject-test",
+            imageName="test:latest",
+            env={"LOG_LEVEL": "INFO"},
+            flashboot=False,
+        )
+        old_resource.id = "ep-inject"
+        old_resource.templateId = "tpl-inject"
+
+        new_resource = ServerlessEndpoint(
+            name="update-inject-test",
+            imageName="test:latest",
+            env={"LOG_LEVEL": "DEBUG"},
+            flashboot=False,
+        )
+
+        mock_client = AsyncMock()
+        mock_client.save_endpoint = AsyncMock(
+            return_value={
+                "id": "ep-inject",
+                "name": "update-inject-test",
+                "templateId": "tpl-inject",
+                "gpuIds": "AMPERE_48",
+                "allowedCudaVersions": "",
+            }
+        )
+        mock_client.update_template = AsyncMock(return_value={})
+
+        with patch(
+            "runpod_flash.core.resources.serverless.RunpodGraphQLClient"
+        ) as mock_client_class:
+            mock_client_class.return_value.__aenter__.return_value = mock_client
+            mock_client_class.return_value.__aexit__.return_value = None
+
+            with patch.object(
+                ServerlessResource,
+                "_ensure_network_volume_deployed",
+                new=AsyncMock(),
+            ):
+                with patch.object(
+                    ServerlessResource,
+                    "_check_makes_remote_calls",
+                    return_value=True,
+                ):
+                    with patch.dict(os.environ, {"RUNPOD_API_KEY": "inject-key"}):
+                        await old_resource.update(new_resource)
+
+        template_payload = mock_client.update_template.call_args.args[0]
+        env_entries = template_payload.get("env", [])
+        api_key_entries = [e for e in env_entries if e["key"] == "RUNPOD_API_KEY"]
+        assert len(api_key_entries) == 1
+        assert api_key_entries[0]["value"] == "inject-key"
+
+    @pytest.mark.asyncio
+    async def test_update_skips_runtime_injection_when_env_unchanged(self):
+        """update() does not inject runtime vars when env is unchanged.
+
+        When skip_env=True, the template env payload is omitted entirely,
+        so runtime vars already on the platform are preserved as-is.
+        """
+        env = {"LOG_LEVEL": "INFO"}
+        old_resource = ServerlessEndpoint(
+            name="update-no-inject",
+            imageName="test:latest",
+            env=env,
+            flashboot=False,
+        )
+        old_resource.id = "ep-no-inject"
+        old_resource.templateId = "tpl-no-inject"
+
+        new_resource = ServerlessEndpoint(
+            name="update-no-inject",
+            imageName="test:latest",
+            env=env,
+            flashboot=False,
+        )
+
+        mock_client = AsyncMock()
+        mock_client.save_endpoint = AsyncMock(
+            return_value={
+                "id": "ep-no-inject",
+                "name": "update-no-inject",
+                "templateId": "tpl-no-inject",
+                "gpuIds": "AMPERE_48",
+                "allowedCudaVersions": "",
+            }
+        )
+        mock_client.update_template = AsyncMock(return_value={})
+
+        with patch(
+            "runpod_flash.core.resources.serverless.RunpodGraphQLClient"
+        ) as mock_client_class:
+            mock_client_class.return_value.__aenter__.return_value = mock_client
+            mock_client_class.return_value.__aexit__.return_value = None
+
+            with patch.object(
+                ServerlessResource,
+                "_ensure_network_volume_deployed",
+                new=AsyncMock(),
+            ):
+                with patch.object(
+                    ServerlessResource,
+                    "_check_makes_remote_calls",
+                    return_value=True,
+                ):
+                    with patch.dict(os.environ, {"RUNPOD_API_KEY": "inject-key"}):
+                        await old_resource.update(new_resource)
+
+        # env should be omitted from template payload (skip_env=True)
+        template_payload = mock_client.update_template.call_args.args[0]
+        assert "env" not in template_payload
+
+    @pytest.mark.asyncio
+    async def test_update_includes_env_for_explicit_template_env(self):
+        """update() sends env when caller provides explicit template.env with empty env.
+
+        Even if self.env == new_config.env (both empty), explicit template.env
+        entries must not be silently dropped.
+        """
+        old_resource = ServerlessEndpoint(
+            name="update-tpl-env",
+            imageName="test:latest",
+            env={},
+            flashboot=False,
+        )
+        old_resource.id = "ep-tpl-env"
+        old_resource.templateId = "tpl-tpl-env"
+
+        new_resource = ServerlessEndpoint(
+            name="update-tpl-env",
+            imageName="test:latest",
+            env={},
+            flashboot=False,
+            template=PodTemplate(
+                name="explicit-tpl",
+                imageName="test:latest",
+                env=[KeyValuePair(key="EXPLICIT_VAR", value="explicit_val")],
+            ),
+        )
+
+        mock_client = AsyncMock()
+        mock_client.save_endpoint = AsyncMock(
+            return_value={
+                "id": "ep-tpl-env",
+                "name": "update-tpl-env",
+                "templateId": "tpl-tpl-env",
+                "gpuIds": "AMPERE_48",
+                "allowedCudaVersions": "",
+            }
+        )
+        mock_client.update_template = AsyncMock(return_value={})
+
+        with patch(
+            "runpod_flash.core.resources.serverless.RunpodGraphQLClient"
+        ) as mock_client_class:
+            mock_client_class.return_value.__aenter__.return_value = mock_client
+            mock_client_class.return_value.__aexit__.return_value = None
+
+            with patch.object(
+                ServerlessResource,
+                "_ensure_network_volume_deployed",
+                new=AsyncMock(),
+            ):
+                await old_resource.update(new_resource)
+
+        template_payload = mock_client.update_template.call_args.args[0]
+        assert "env" in template_payload
+        env_entries = template_payload["env"]
+        explicit = [e for e in env_entries if e["key"] == "EXPLICIT_VAR"]
+        assert len(explicit) == 1