feat(server): make agent env and timeout configurable in helm chart (#1281)

Signed-off-by: Radek Ježek <radek.jezek@ibm.com>

Author: jezekra1

apps/beeai-sdk/src/beeai_sdk/platform/provider.py

@@ -27,7 +27,7 @@ class EnvVar(pydantic.BaseModel):
 
 class Provider(pydantic.BaseModel):
     id: str
-    auto_stop_timeout: timedelta
+    auto_stop_timeout: timedelta | None = None
     source: str
     registry: str | None = None
     auto_remove: bool = False

apps/beeai-server/src/beeai_server/api/routes/providers.py

@@ -1,5 +1,6 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
+from datetime import timedelta
 from typing import Annotated
 from uuid import UUID
 
@@ -37,6 +38,7 @@ async def create_provider(
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Auto remove functionality is disabled")
     return await provider_service.create_provider(
         user=user.user,
+        auto_stop_timeout=timedelta(seconds=request.auto_stop_timeout_sec),
         location=request.location,
         agent_card=request.agent_card,
         auto_remove=auto_remove,

apps/beeai-server/src/beeai_server/api/schema/provider.py

@@ -1,8 +1,10 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
 
+from datetime import timedelta
+
 from a2a.types import AgentCard
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from beeai_server.domain.models.provider import ProviderLocation
 
@@ -11,3 +13,12 @@ class CreateProviderRequest(BaseModel):
     location: ProviderLocation
     agent_card: AgentCard | None = None
     variables: dict[str, str] | None = None
+    auto_stop_timeout_sec: int = Field(
+        default=int(timedelta(minutes=5).total_seconds()),
+        gt=0,
+        le=600,
+        description=(
+            "Timeout after which the agent provider will be automatically downscaled if unused."
+            "Contact administrator if you need to increase this value."
+        ),
+    )

apps/beeai-server/src/beeai_server/application.py

@@ -2,15 +2,17 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+import time
 from collections.abc import Iterable
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, suppress
 
 import procrastinate
 from fastapi import APIRouter, FastAPI, HTTPException
 from fastapi.exception_handlers import http_exception_handler
 from fastapi.responses import ORJSONResponse
 from kink import Container, di, inject
 from opentelemetry.metrics import CallbackOptions, Observation, get_meter
+from procrastinate.exceptions import AlreadyEnqueued
 from starlette.requests import Request
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_500_INTERNAL_SERVER_ERROR
 
@@ -34,6 +36,7 @@
     ManifestLoadError,
     PlatformError,
 )
+from beeai_server.jobs.crons.provider import check_registry
 from beeai_server.run_workers import run_workers
 from beeai_server.service_layer.services.mcp import McpService
 from beeai_server.telemetry import INSTRUMENTATION_NAME, shutdown_telemetry
@@ -149,6 +152,9 @@ async def lifespan(_app: FastAPI, procrastinate_app: procrastinate.App, mcp_serv
         try:
             register_telemetry()
             async with procrastinate_app.open_async(), run_workers(app=procrastinate_app), mcp_service:
+                with suppress(AlreadyEnqueued):
+                    # Force initial sync of the registry immediately
+                    await check_registry.defer_async(timestamp=int(time.time()))
                 try:
                     yield
                 finally:

apps/beeai-server/src/beeai_server/domain/models/registry.py

@@ -1,12 +1,12 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
-
+from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 import httpx
 import yaml
 from anyio import Path
-from pydantic import BaseModel, FileUrl, HttpUrl, RootModel
+from pydantic import BaseModel, Field, FileUrl, HttpUrl, RootModel, computed_field, field_validator
 
 from beeai_server.utils.github import GithubUrl
 
@@ -16,22 +16,42 @@
     from beeai_server.domain.models.provider import ProviderLocation
 
 
-def parse_providers_manifest(content: dict[str, Any]) -> list["ProviderLocation"]:
-    from beeai_server.domain.models.provider import ProviderLocation
+class ProviderRegistryRecord(BaseModel, extra="allow"):
+    location: "ProviderLocation"
+    auto_stop_timeout_sec: int | None = Field(default=int(timedelta(minutes=5).total_seconds()), ge=0)
+    variables: dict[str, str] = {}
+
+    @computed_field
+    @property
+    def auto_stop_timeout(self) -> timedelta | None:
+        return timedelta(seconds=self.auto_stop_timeout_sec) if self.auto_stop_timeout_sec else None
+
+    @field_validator("variables", mode="before")
+    @classmethod
+    def convert_variables_to_str(cls, v: Any | None):
+        if v is None:
+            return {}
+        if not isinstance(v, dict):
+            raise ValueError("env must be a dictionary")
+        return {str(k): str(v) for k, v in v.items()}
 
-    class ProviderRegistryRecord(BaseModel, extra="allow"):
-        location: ProviderLocation
 
-    class RegistryManifest(BaseModel):
-        providers: list[ProviderRegistryRecord]
+class RegistryManifest(BaseModel):
+    providers: list[ProviderRegistryRecord]
+
+
+def parse_providers_manifest(content: dict[str, Any]) -> list[ProviderRegistryRecord]:
+    from beeai_server.domain.models.provider import ProviderLocation
+
+    _ = ProviderLocation  # make sure this is imported
 
-    return [p.location for p in RegistryManifest.model_validate(content).providers]
+    return RegistryManifest.model_validate(content).providers
 
 
 class NetworkRegistryLocation(RootModel):
     root: HttpUrl
 
-    async def load(self) -> list["ProviderLocation"]:
+    async def load(self) -> list[ProviderRegistryRecord]:
         async with httpx.AsyncClient(
             headers={"Cache-Control": "no-cache, no-store, must-revalidate", "Pragma": "no-cache", "Expires": "0"}
         ) as client:
@@ -42,7 +62,7 @@ async def load(self) -> list["ProviderLocation"]:
 class GithubRegistryLocation(RootModel):
     root: GithubUrl
 
-    async def load(self) -> list["ProviderLocation"]:
+    async def load(self) -> list[ProviderRegistryRecord]:
         resolved_url = await self.root.resolve_version()
         url = await resolved_url.get_raw_url()
         network_location = NetworkRegistryLocation(root=HttpUrl(url))
@@ -52,7 +72,7 @@ async def load(self) -> list["ProviderLocation"]:
 class FileSystemRegistryLocation(RootModel):
     root: FileUrl
 
-    async def load(self) -> list["ProviderLocation"]:
+    async def load(self) -> list[ProviderRegistryRecord]:
         content = await Path(self.root.path).read_text()
         return parse_providers_manifest(yaml.safe_load(content))
 

apps/beeai-server/src/beeai_server/domain/repositories/env.py

@@ -37,3 +37,5 @@ async def update(
         parent_entity_id: UUID,
         variables: dict[str, str | None] | dict[str, str],
     ) -> None: ...
+
+    async def delete(self, parent_entity: EnvStoreEntity, parent_entity_id: UUID) -> None: ...

apps/beeai-server/src/beeai_server/infrastructure/persistence/repositories/env.py

@@ -139,6 +139,10 @@ async def update(
                 )
             )
 
+    async def delete(self, parent_entity: EnvStoreEntity, parent_entity_id: UUID):
+        query = variables_table.delete().where(self._parent_filter(parent_entity, parent_entity_id))
+        await self.connection.execute(query)
+
     async def get(
         self,
         *,

apps/beeai-server/src/beeai_server/jobs/crons/provider.py

@@ -4,6 +4,7 @@
 import logging
 from contextlib import suppress
 from datetime import timedelta
+from uuid import UUID
 
 import anyio
 import httpx
@@ -15,6 +16,7 @@
 from beeai_server import get_configuration
 from beeai_server.configuration import Configuration
 from beeai_server.domain.models.provider import ProviderLocation
+from beeai_server.domain.models.registry import ProviderRegistryRecord, RegistryLocation
 from beeai_server.exceptions import EntityNotFoundError
 from beeai_server.jobs.queues import Queues
 from beeai_server.service_layer.services.providers import ProviderService
@@ -49,8 +51,8 @@ async def check_registry(
 
     user = await user_service.get_user_by_email("admin@beeai.dev")
 
-    registry_by_provider_id = {}
-    desired_providers = {}
+    registry_by_provider_id: dict[UUID, RegistryLocation] = {}
+    desired_providers: dict[UUID, ProviderRegistryRecord] = {}
     errors = []
 
     try:
@@ -59,10 +61,10 @@ async def check_registry(
         errors.extend(ex.exceptions if isinstance(ex, ExceptionGroup) else [ex])
 
     for registry in configuration.agent_registry.locations.values():
-        for provider_location in await registry.load():
+        for provider_record in await registry.load():
             try:
-                provider_id = RootModel[ProviderLocation](root=provider_location).root.provider_id
-                desired_providers[provider_id] = provider_location
+                provider_id = RootModel[ProviderLocation](root=provider_record.location).root.provider_id
+                desired_providers[provider_id] = provider_record
                 registry_by_provider_id[provider_id] = registry
             except ValueError as e:
                 errors.append(e)
@@ -79,31 +81,38 @@ async def check_registry(
     for provider_id in old_providers:
         provider = managed_providers[provider_id]
         try:
-            await provider_service.delete_provider(provider_id=provider.id)
+            await provider_service.delete_provider(provider_id=provider.id, user=user)
             logger.info(f"Removed provider {provider.source}")
         except Exception as ex:
             errors.append(RuntimeError(f"[{provider.source}]: Failed to remove provider: {ex}"))
 
     for provider_id in new_providers:
-        provider_location = desired_providers[provider_id]
+        provider_record = desired_providers[provider_id]
         try:
             await provider_service.create_provider(
                 user=user,
-                location=provider_location,
+                location=provider_record.location,
                 registry=registry_by_provider_id[provider_id],
+                auto_stop_timeout=provider_record.auto_stop_timeout,
+                variables=provider_record.variables,
             )
-            logger.info(f"Added provider {provider_location}")
+            logger.info(f"Added provider {provider_record}")
         except Exception as ex:
-            errors.append(RuntimeError(f"[{provider_location}]: Failed to add provider: {ex}"))
+            errors.append(RuntimeError(f"[{provider_record}]: Failed to add provider: {ex}"))
 
     for provider_id in existing_providers:
-        provider_location = desired_providers[provider_id]
+        provider_record = desired_providers[provider_id]
         try:
-            result = await provider_service.upgrade_provider(provider_id=provider_id, location=provider_location)
+            result = await provider_service.upgrade_provider(
+                provider_id=provider_id,
+                location=provider_record.location,
+                auto_stop_timeout=provider_record.auto_stop_timeout,
+                env=provider_record.variables,
+            )
             if managed_providers[provider_id].source.root != result.source.root:
-                logger.info(f"Upgraded provider {provider_location}")
+                logger.info(f"Upgraded provider {provider_record}")
         except Exception as ex:
-            errors.append(RuntimeError(f"[{provider_location}]: Failed to add provider: {ex}"))
+            errors.append(RuntimeError(f"[{provider_record}]: Failed to add provider: {ex}"))
 
     if errors:
         raise ExceptionGroup("Exceptions occurred when reloading providers", errors)

apps/beeai-server/src/beeai_server/service_layer/services/providers.py

@@ -8,6 +8,7 @@
 import logging
 import uuid
 from collections.abc import AsyncIterator, Callable
+from datetime import timedelta
 from uuid import UUID
 
 from a2a.types import AgentCard
@@ -46,6 +47,7 @@ async def create_provider(
         *,
         user: User,
         location: ProviderLocation,
+        auto_stop_timeout: timedelta | None,
         registry: RegistryLocation | None = None,
         auto_remove: bool = False,
         agent_card: AgentCard | None = None,
@@ -55,7 +57,12 @@ async def create_provider(
             if not agent_card:
                 agent_card = await location.load_agent_card()
             provider = Provider(
-                source=location, registry=registry, auto_remove=auto_remove, agent_card=agent_card, created_by=user.id
+                source=location,
+                registry=registry,
+                auto_remove=auto_remove,
+                agent_card=agent_card,
+                created_by=user.id,
+                auto_stop_timeout=auto_stop_timeout,
             )
         except ValueError as ex:
             raise ManifestLoadError(location=location, message=str(ex), status_code=HTTP_400_BAD_REQUEST) from ex
@@ -73,13 +80,33 @@ async def create_provider(
         return provider_response
 
     async def upgrade_provider(
-        self, *, provider_id: UUID, location: ProviderLocation, force: bool = False
+        self,
+        *,
+        provider_id: UUID,
+        location: ProviderLocation,
+        force: bool = False,
+        auto_stop_timeout: timedelta | None,
+        env: dict[str, str] | None = None,
     ) -> ProviderWithState:
+        env = env or {}
         async with self._uow() as uow:
             provider = await uow.providers.get(provider_id=provider_id)
+            old_env = (
+                await uow.env.get_all(
+                    parent_entity=EnvStoreEntity.PROVIDER,
+                    parent_entity_ids=[provider.id],
+                )
+            )[provider.id]
 
-        if provider.source.root == location.root and not force:
-            return (await self._get_providers_with_state([provider]))[0]
+        should_update = (
+            provider.source.root != location.root
+            or provider.auto_stop_timeout != auto_stop_timeout
+            or env != old_env
+            or force
+        )
+
+        if not should_update:
+            return (await self._get_providers_with_state(providers=[provider]))[0]
 
         try:
             agent_card = await location.load_agent_card()
@@ -90,12 +117,16 @@ async def upgrade_provider(
 
         provider.source = location
         provider.agent_card = agent_card
+        provider.auto_stop_timeout = auto_stop_timeout
 
         async with self._uow() as uow:
             await uow.providers.update(provider=provider)
-            env = await uow.env.get_all(parent_entity=EnvStoreEntity.PROVIDER, parent_entity_ids=[provider.id])
+
+            if old_env != env:
+                await uow.env.delete(parent_entity=EnvStoreEntity.PROVIDER, parent_entity_id=provider.id)
+                await uow.env.update(parent_entity=EnvStoreEntity.PROVIDER, parent_entity_id=provider.id, variables=env)
             # Rotate the provider (inside the transaction)
-            await self._rotate_provider(provider=provider, env=env[provider.id])
+            await self._rotate_provider(provider=provider, env=env)
             await uow.commit()
         [provider_response] = await self._get_providers_with_state(providers=[provider])
         return provider_response
@@ -216,7 +247,9 @@ async def _rotate_provider(self, provider: Provider, env: dict[str, str]):
         ):
             await self._deployment_manager.create_or_replace(provider=provider, env=env)
 
-    async def update_provider_env(self, *, provider_id: UUID, env: dict[str, str | None], user: User) -> None:
+    async def update_provider_env(
+        self, *, provider_id: UUID, env: dict[str, str | None] | dict[str, str], user: User
+    ) -> None:
         user_id = user.id if user.role != UserRole.ADMIN else None
         provider = None
         try:

helm/templates/config/providers.yaml

@@ -1,17 +1,18 @@
 {{- $root := . }}
 apiVersion: v1
-kind: ConfigMap
+kind: Secret
 metadata:
-  name: beeai-platform-providers-cm
-data:
+  name: beeai-platform-providers-secret
+type: Opaque
+stringData:
   registry.yaml: |
     {{- if or .Values.providers .Values.unmanagedProviders }}
     providers:
       {{- range $idx, $p := $root.Values.unmanagedProviders }}
       - location: "http://{{ include "agent.fullname" (dict "root" $root "image" $p.location) }}:{{ $root.Values.agent.service.port }}"
       {{- end }}
       {{- range $idx, $p := $root.Values.providers }}
-      - location: {{ $p.location }}
+      - {{ $p | toYaml | nindent 8 | trim }}
       {{- end }}
     {{- else }}
     providers: [ ]