feat(automations): implement model eligibility checks for automation creation

- Added model eligibility checks to ensure automations can only use billable models (premium or BYOK).
- Introduced new API endpoint to report model eligibility status for search spaces.
- Updated frontend components to display eligibility alerts and disable creation options when models are not billable.
- Enhanced automation creation forms to reflect model eligibility, preventing users from submitting invalid configurations.
- Implemented server-side logic to capture and preserve model preferences across automation edits, ensuring consistent behavior during execution.

surfsense_backend/app/automations/actions/agent_task/dependencies.py

@@ -8,6 +8,12 @@ from typing import Any
 from langgraph.checkpoint.memory import InMemorySaver
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from app.automations.services.model_policy import (
+    AutomationModelPolicyError,
+    assert_automation_models_billable,
+    assert_models_billable,
+)
+from app.db import SearchSpace
 from app.tasks.chat.streaming.flows.shared.llm_bundle import load_llm_bundle
 from app.tasks.chat.streaming.flows.shared.pre_stream_setup import (
     setup_connector_and_firecrawl,
@@ -33,17 +39,48 @@ async def build_dependencies(
     *,
     session: AsyncSession,
     search_space_id: int,
+    agent_llm_id: int | None = None,
+    image_generation_config_id: int | None = None,
+    vision_llm_config_id: int | None = None,
 ) -> AgentDependencies:
     """Load the LLM bundle, connector service, and a per-invoke in-memory checkpointer.
 
-    Uses the search space's default LLM config (``config_id=-1``). Per-step
-    model overrides land in a future iteration alongside the ``model`` param.
+    Resolves the agent LLM from the automation's *captured* model snapshot
+    (``agent_llm_id``) so runs are insulated from later chat/search-space model
+    changes. The model policy is enforced here as a runtime backstop: a captured
+    model that is no longer billable (e.g. a premium global config was removed)
+    fails the run clearly instead of silently consuming a free model.
+
+    When ``agent_llm_id`` is ``None`` (no captured snapshot — defensive fallback),
+    fall back to the live search space's ``agent_llm_id`` and validate that.
     """
+    if agent_llm_id is not None:
+        try:
+            assert_models_billable(
+                agent_llm_id=agent_llm_id,
+                image_generation_config_id=image_generation_config_id,
+                vision_llm_config_id=vision_llm_config_id,
+            )
+        except AutomationModelPolicyError as exc:
+            raise DependencyError(str(exc)) from exc
+        resolved_agent_llm_id = agent_llm_id or 0
+    else:
+        search_space = await session.get(SearchSpace, search_space_id)
+        if search_space is None:
+            raise DependencyError(f"search space {search_space_id} not found")
+        try:
+            assert_automation_models_billable(search_space)
+        except AutomationModelPolicyError as exc:
+            raise DependencyError(str(exc)) from exc
+        resolved_agent_llm_id = search_space.agent_llm_id or 0
+
     llm, agent_config, err = await load_llm_bundle(
-        session, config_id=-1, search_space_id=search_space_id
+        session,
+        config_id=resolved_agent_llm_id,
+        search_space_id=search_space_id,
     )
     if err is not None or llm is None:
-        raise DependencyError(err or "failed to load default LLM config")
+        raise DependencyError(err or "failed to load agent LLM config")
 
     connector_service, firecrawl_api_key = await setup_connector_and_firecrawl(
         session, search_space_id=search_space_id

surfsense_backend/app/automations/actions/agent_task/invoke.py

@@ -147,6 +147,9 @@ async def run_agent_task(
         deps = await build_dependencies(
             session=agent_session,
             search_space_id=ctx.search_space_id,
+            agent_llm_id=ctx.agent_llm_id,
+            image_generation_config_id=ctx.image_generation_config_id,
+            vision_llm_config_id=ctx.vision_llm_config_id,
         )
 
         agent = await create_multi_agent_chat_deep_agent(
@@ -161,6 +164,7 @@ async def run_agent_task(
             firecrawl_api_key=deps.firecrawl_api_key,
             thread_visibility=ChatVisibility.PRIVATE,
             mentioned_document_ids=mentioned_document_ids,
+            image_generation_config_id=ctx.image_generation_config_id,
         )
 
         agent_query, runtime_context = await _resolve_mention_context(

surfsense_backend/app/automations/actions/types.py

@@ -20,6 +20,12 @@ class ActionContext:
     step_id: str
     search_space_id: int
     creator_user_id: UUID | None
+    # Captured model snapshot from the automation definition (``definition.models``),
+    # resolved per run instead of the live search space. ``None`` falls back to the
+    # search space's current prefs (defensive; should not happen post-capture).
+    agent_llm_id: int | None = None
+    image_generation_config_id: int | None = None
+    vision_llm_config_id: int | None = None
 
 
 ActionHandler = Callable[[dict[str, Any]], Awaitable[Any]]

surfsense_backend/app/automations/api/automation.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from fastapi import APIRouter, Depends, Query, status
+from pydantic import BaseModel
 
 from app.automations.schemas.api import (
     AutomationCreate,
@@ -16,6 +17,17 @@ from app.automations.services import AutomationService, get_automation_service
 router = APIRouter()
 
 
+class ModelEligibilityViolation(BaseModel):
+    kind: str
+    config_id: int | None
+    reason: str
+
+
+class ModelEligibility(BaseModel):
+    allowed: bool
+    violations: list[ModelEligibilityViolation]
+
+
 @router.post(
     "/automations",
     response_model=AutomationDetail,
@@ -47,6 +59,23 @@ async def list_automations(
     )
 
 
+@router.get("/automations/model-eligibility", response_model=ModelEligibility)
+async def get_automation_model_eligibility(
+    search_space_id: int = Query(...),
+    service: AutomationService = Depends(get_automation_service),
+) -> ModelEligibility:
+    """Report whether a search space's models are billable for automations.
+
+    Used by the frontend to gate creation: automations may only use premium
+    global models or user BYOK models (free models and Auto mode are blocked).
+
+    NOTE: declared before ``/automations/{automation_id}`` so the literal path
+    isn't captured by the int-typed ``{automation_id}`` route.
+    """
+    result = await service.model_eligibility(search_space_id=search_space_id)
+    return ModelEligibility.model_validate(result)
+
+
 @router.get("/automations/{automation_id}", response_model=AutomationDetail)
 async def get_automation(
     automation_id: int,

surfsense_backend/app/automations/runtime/executor.py

@@ -9,7 +9,10 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from app.automations.actions.types import ActionContext
 from app.automations.persistence.enums.run_status import RunStatus
 from app.automations.persistence.models.run import AutomationRun
-from app.automations.schemas.definition.envelope import AutomationDefinition
+from app.automations.schemas.definition.envelope import (
+    AutomationDefinition,
+    AutomationModels,
+)
 from app.automations.schemas.definition.plan_step import PlanStep
 from app.automations.templating import build_run_context
 
@@ -47,7 +50,7 @@ async def execute_run(session: AsyncSession, run_id: int) -> None:
 
     for step in definition.plan:
         template_ctx = _build_template_ctx(run, step_outputs)
-        action_ctx = _build_action_ctx(session, run, step)
+        action_ctx = _build_action_ctx(session, run, step, definition.models)
         result = await execute_step(
             step=step,
             template_context=template_ctx,
@@ -82,7 +85,7 @@ async def _run_on_failure(
         return
     template_ctx = _build_template_ctx(run, step_outputs={})
     for step in definition.execution.on_failure:
-        action_ctx = _build_action_ctx(session, run, step)
+        action_ctx = _build_action_ctx(session, run, step, definition.models)
         result = await execute_step(
             step=step,
             template_context=template_ctx,
@@ -117,7 +120,10 @@ def _build_template_ctx(
 
 
 def _build_action_ctx(
-    session: AsyncSession, run: AutomationRun, step: PlanStep
+    session: AsyncSession,
+    run: AutomationRun,
+    step: PlanStep,
+    models: AutomationModels | None,
 ) -> ActionContext:
     automation = run.automation
     return ActionContext(
@@ -126,4 +132,9 @@ def _build_action_ctx(
         step_id=step.step_id,
         search_space_id=automation.search_space_id,
         creator_user_id=automation.created_by_user_id,
+        agent_llm_id=models.agent_llm_id if models else None,
+        image_generation_config_id=(
+            models.image_generation_config_id if models else None
+        ),
+        vision_llm_config_id=models.vision_llm_config_id if models else None,
     )

surfsense_backend/app/automations/schemas/definition/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from .envelope import AutomationDefinition
+from .envelope import AutomationDefinition, AutomationModels
 from .execution import Execution
 from .inputs import Inputs
 from .metadata import Metadata
@@ -11,6 +11,7 @@ from .trigger_spec import TriggerSpec
 
 __all__ = [
     "AutomationDefinition",
+    "AutomationModels",
     "Execution",
     "Inputs",
     "Metadata",

surfsense_backend/app/automations/schemas/definition/envelope.py

@@ -11,6 +11,21 @@ from .plan_step import PlanStep
 from .trigger_spec import TriggerSpec
 
 
+class AutomationModels(BaseModel):
+    """Captured model profile for an automation.
+
+    Snapshotted from the search space's preferences at create time so runs are
+    insulated from later chat/search-space model changes. Config-id conventions
+    match the shared scheme (``0`` Auto, ``< 0`` global, ``> 0`` BYOK).
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    agent_llm_id: int = 0
+    image_generation_config_id: int = 0
+    vision_llm_config_id: int = 0
+
+
 class AutomationDefinition(BaseModel):
     """Top-level shape of an automation."""
 
@@ -24,3 +39,7 @@ class AutomationDefinition(BaseModel):
     plan: list[PlanStep] = Field(..., min_length=1)
     execution: Execution = Field(default_factory=Execution)
     metadata: Metadata = Field(default_factory=Metadata)
+    # Captured server-side at create() and preserved across update(); resolved
+    # at runtime instead of the live search space. Optional so drafts/builder
+    # payloads validate without it.
+    models: AutomationModels | None = None

surfsense_backend/app/automations/services/__init__.py

@@ -3,14 +3,26 @@
 from __future__ import annotations
 
 from .automation import AutomationService, get_automation_service
+from .model_policy import (
+    AutomationModelPolicyError,
+    assert_automation_models_billable,
+    assert_models_billable,
+    get_automation_model_eligibility,
+    get_model_eligibility,
+)
 from .run import RunService, get_run_service
 from .trigger import TriggerService, get_trigger_service
 
 __all__ = [
+    "AutomationModelPolicyError",
     "AutomationService",
     "RunService",
     "TriggerService",
+    "assert_automation_models_billable",
+    "assert_models_billable",
+    "get_automation_model_eligibility",
     "get_automation_service",
+    "get_model_eligibility",
     "get_run_service",
     "get_trigger_service",
 ]

surfsense_backend/app/automations/services/automation.py

@@ -18,9 +18,15 @@ from app.automations.schemas.api import (
     AutomationUpdate,
     TriggerCreate,
 )
+from app.automations.schemas.definition.envelope import AutomationModels
+from app.automations.services.model_policy import (
+    AutomationModelPolicyError,
+    assert_automation_models_billable,
+    get_automation_model_eligibility,
+)
 from app.automations.triggers import get_trigger
 from app.automations.triggers.schedule import compute_next_fire_at
-from app.db import Permission, User, get_async_session
+from app.db import Permission, SearchSpace, User, get_async_session
 from app.users import current_active_user
 from app.utils.rbac import check_permission
 
@@ -37,6 +43,16 @@ class AutomationService:
         await self._authorize(
             payload.search_space_id, Permission.AUTOMATIONS_CREATE.value
         )
+        search_space = await self._assert_models_billable(payload.search_space_id)
+
+        # Snapshot the search space's current (already-validated) model prefs onto
+        # the definition so runs are insulated from later chat/search-space model
+        # changes. Captured ids are guaranteed billable by the check above.
+        payload.definition.models = AutomationModels(
+            agent_llm_id=search_space.agent_llm_id or 0,
+            image_generation_config_id=search_space.image_generation_config_id or 0,
+            vision_llm_config_id=search_space.vision_llm_config_id or 0,
+        )
 
         automation = Automation(
             search_space_id=payload.search_space_id,
@@ -105,9 +121,15 @@ class AutomationService:
         if "status" in data:
             automation.status = data["status"]
         if "definition" in data:
-            automation.definition = patch.definition.model_dump(
-                mode="json", by_alias=True
-            )
+            new_def = patch.definition.model_dump(mode="json", by_alias=True)
+            # Preserve the captured model snapshot across edits so a definition
+            # change never silently re-binds the automation to the current chat
+            # model selection. Backend-managed; survives whether or not the
+            # client round-trips ``models``.
+            existing_models = (automation.definition or {}).get("models")
+            if existing_models is not None:
+                new_def["models"] = existing_models
+            automation.definition = new_def
             automation.version += 1
 
         await self.session.commit()
@@ -143,6 +165,40 @@ class AutomationService:
             )
         return automation
 
+    async def model_eligibility(self, *, search_space_id: int) -> dict:
+        """Return whether a search space's models are billable for automations.
+
+        ``{"allowed": bool, "violations": [{kind, config_id, reason}, ...]}``.
+        """
+        await self._authorize(search_space_id, Permission.AUTOMATIONS_READ.value)
+        search_space = await self.session.get(SearchSpace, search_space_id)
+        if search_space is None:
+            raise HTTPException(
+                status_code=404, detail=f"search space {search_space_id} not found"
+            )
+        return get_automation_model_eligibility(search_space)
+
+    async def _assert_models_billable(self, search_space_id: int) -> SearchSpace:
+        """Reject creation when the search space's models aren't billable.
+
+        Automations may only use premium global models or user BYOK models; free
+        global models and Auto mode are blocked. Mirrors the runtime backstop in
+        ``agent_task`` so users can't save an automation that would fail to run.
+
+        Returns the loaded :class:`SearchSpace` so the caller can capture its
+        model prefs without a second DB read.
+        """
+        search_space = await self.session.get(SearchSpace, search_space_id)
+        if search_space is None:
+            raise HTTPException(
+                status_code=404, detail=f"search space {search_space_id} not found"
+            )
+        try:
+            assert_automation_models_billable(search_space)
+        except AutomationModelPolicyError as exc:
+            raise HTTPException(status_code=422, detail=str(exc)) from exc
+        return search_space
+
     async def _authorize(self, search_space_id: int, permission: str) -> None:
         await check_permission(
             self.session,

surfsense_backend/app/automations/services/model_policy.py

@@ -0,0 +1,173 @@
+"""Model-billing policy for automations.
+
+Automations run unattended, so every run must be **billable**: it may only use
+either a premium global model (``billing_tier == "premium"``) or a user-provided
+BYOK model (a positive config id pointing at a per-user/per-space DB row). Free
+global models and Auto mode are blocked, because Auto can dispatch to a free
+deployment and free models aren't metered in premium credits.
+
+Config id conventions (shared across chat / image / vision):
+- ``id == 0``  → Auto mode (``AUTO_MODE_ID`` / ``IMAGE_GEN_AUTO_MODE_ID`` /
+  ``VISION_AUTO_MODE_ID``). Blocked.
+- ``id < 0``   → global YAML/OpenRouter config. Allowed only if premium.
+- ``id > 0``   → user BYOK DB row. Always allowed.
+
+This module is the single source of truth used by both creation-time enforcement
+(``AutomationService.create`` and the ``create_automation`` chat tool) and the
+runtime backstop (``agent_task`` dependencies).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from app.db import SearchSpace
+
+ModelKind = Literal["llm", "image", "vision"]
+
+_KIND_LABEL: dict[ModelKind, str] = {
+    "llm": "agent LLM",
+    "image": "image generation model",
+    "vision": "vision model",
+}
+
+
+def _is_premium_global(kind: ModelKind, config_id: int) -> bool:
+    """Return True if a negative (global) config id is a premium tier model."""
+    from app.config import config as app_config
+
+    cfg: dict | None = None
+    if kind == "llm":
+        from app.agents.new_chat.llm_config import load_global_llm_config_by_id
+
+        cfg = load_global_llm_config_by_id(config_id)
+    elif kind == "image":
+        cfg = next(
+            (
+                c
+                for c in app_config.GLOBAL_IMAGE_GEN_CONFIGS
+                if c.get("id") == config_id
+            ),
+            None,
+        )
+    else:  # vision
+        cfg = next(
+            (
+                c
+                for c in app_config.GLOBAL_VISION_LLM_CONFIGS
+                if c.get("id") == config_id
+            ),
+            None,
+        )
+
+    if not cfg:
+        return False
+    return str(cfg.get("billing_tier", "free")).lower() == "premium"
+
+
+def _classify(kind: ModelKind, config_id: int | None) -> tuple[bool, str]:
+    """Classify a resolved config id as allowed or blocked.
+
+    Returns ``(allowed, reason)``; ``reason`` is empty when allowed.
+    """
+    label = _KIND_LABEL[kind]
+
+    if config_id is None or config_id == 0:
+        return (
+            False,
+            f"The {label} is set to Auto mode. Automations require an explicit "
+            "premium model or your own (BYOK) model so every run is billable.",
+        )
+
+    if config_id > 0:
+        # Positive id → user-owned BYOK config. Always allowed.
+        return True, ""
+
+    # Negative id → global config. Allowed only if premium.
+    if _is_premium_global(kind, config_id):
+        return True, ""
+
+    return (
+        False,
+        f"The {label} is a free model. Automations can only use premium models "
+        "or your own (BYOK) models so every run is billable.",
+    )
+
+
+def get_model_eligibility(
+    *,
+    agent_llm_id: int | None,
+    image_generation_config_id: int | None,
+    vision_llm_config_id: int | None,
+) -> dict:
+    """Return ``{"allowed": bool, "violations": [...]}`` for explicit config ids.
+
+    The ID-based core shared by both the search-space path (creation/eligibility)
+    and the captured-snapshot path (runtime backstop). Each violation is
+    ``{"kind", "config_id", "reason"}``.
+    """
+    checks: list[tuple[ModelKind, int | None]] = [
+        ("llm", agent_llm_id),
+        ("image", image_generation_config_id),
+        ("vision", vision_llm_config_id),
+    ]
+
+    violations: list[dict] = []
+    for kind, config_id in checks:
+        allowed, reason = _classify(kind, config_id)
+        if not allowed:
+            violations.append({"kind": kind, "config_id": config_id, "reason": reason})
+
+    return {"allowed": not violations, "violations": violations}
+
+
+def get_automation_model_eligibility(search_space: SearchSpace) -> dict:
+    """Return ``{"allowed": bool, "violations": [...]}`` for a search space.
+
+    Used by the eligibility endpoint and the chat tool's early check. Thin
+    wrapper over :func:`get_model_eligibility`.
+    """
+    return get_model_eligibility(
+        agent_llm_id=search_space.agent_llm_id,
+        image_generation_config_id=search_space.image_generation_config_id,
+        vision_llm_config_id=search_space.vision_llm_config_id,
+    )
+
+
+class AutomationModelPolicyError(Exception):
+    """Raised when a search space's models are not billable for automations."""
+
+    def __init__(self, violations: list[dict]) -> None:
+        self.violations = violations
+        reasons = "; ".join(v["reason"] for v in violations)
+        super().__init__(
+            reasons or "Automations require premium or BYOK models for all model slots."
+        )
+
+
+def assert_models_billable(
+    *,
+    agent_llm_id: int | None,
+    image_generation_config_id: int | None,
+    vision_llm_config_id: int | None,
+) -> None:
+    """Raise :class:`AutomationModelPolicyError` if any explicit id is not billable.
+
+    The ID-based core used by the runtime backstop against an automation's
+    captured model snapshot.
+    """
+    result = get_model_eligibility(
+        agent_llm_id=agent_llm_id,
+        image_generation_config_id=image_generation_config_id,
+        vision_llm_config_id=vision_llm_config_id,
+    )
+    if not result["allowed"]:
+        raise AutomationModelPolicyError(result["violations"])
+
+
+def assert_automation_models_billable(search_space: SearchSpace) -> None:
+    """Raise :class:`AutomationModelPolicyError` if any model slot is not billable."""
+    result = get_automation_model_eligibility(search_space)
+    if not result["allowed"]:
+        raise AutomationModelPolicyError(result["violations"])