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.
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-05-29 03:13:46 -07:00
parent 5d90fbe99f
commit 409fec94c3
32 changed files with 1451 additions and 67 deletions

View file

@ -57,6 +57,7 @@ async def build_agent_with_cache(
mcp_tools_by_agent: dict[str, list[BaseTool]],
disabled_tools: list[str] | None,
config_id: str | None,
image_generation_config_id_override: int | None = None,
) -> Any:
"""Compile the multi-agent graph, serving from cache when key components are stable."""
@ -91,7 +92,7 @@ async def build_agent_with_cache(
# the key, otherwise a hit will leak state across threads. Bump the schema
# version when the component list changes shape.
cache_key = stable_hash(
"multi-agent-v1",
"multi-agent-v2",
config_id,
thread_id,
user_id,
@ -109,6 +110,10 @@ async def build_agent_with_cache(
system_prompt_hash(final_system_prompt),
max_input_tokens,
sorted(disabled_tools) if disabled_tools else None,
# Bound into the generate_image subagent tool at construction time, so it
# must key the compiled-agent cache to avoid leaking one automation's
# image model into another with the same config_id/search_space.
image_generation_config_id_override,
)
return await get_cache().get_or_build(cache_key, builder=_build)

View file

@ -62,8 +62,14 @@ async def create_multi_agent_chat_deep_agent(
mentioned_document_ids: list[int] | None = None,
anon_session_id: str | None = None,
filesystem_selection: FilesystemSelection | None = None,
image_generation_config_id: int | None = None,
):
"""Deep agent with SurfSense tools/middleware; registry route subagents behind ``task`` when enabled."""
"""Deep agent with SurfSense tools/middleware; registry route subagents behind ``task`` when enabled.
``image_generation_config_id`` overrides the search space's image model for
this invocation (used by automations to run on their captured model). When
``None``, the ``generate_image`` tool resolves the live search-space pref.
"""
_t_agent_total = time.perf_counter()
apply_litellm_prompt_caching(llm, agent_config=agent_config, thread_id=thread_id)
@ -129,6 +135,9 @@ async def create_multi_agent_chat_deep_agent(
"available_document_types": available_document_types,
"max_input_tokens": _max_input_tokens,
"llm": llm,
# Per-invocation image model override (automations run on their captured
# model). Reaches the generate_image subagent tool via subagent_dependencies.
"image_generation_config_id_override": image_generation_config_id,
}
_t0 = time.perf_counter()
@ -285,6 +294,7 @@ async def create_multi_agent_chat_deep_agent(
mcp_tools_by_agent=mcp_tools_by_agent,
disabled_tools=disabled_tools,
config_id=config_id,
image_generation_config_id_override=image_generation_config_id,
)
_perf_log.info(
"[create_agent] Middleware stack + graph compiled in %.3fs",

View file

@ -32,7 +32,8 @@ from app.agents.multi_agent_chat.subagents.shared.hitl.approvals.self_gated impo
)
from app.automations.schemas.api import AutomationCreate
from app.automations.services.automation import AutomationService
from app.db import User, async_session_maker
from app.automations.services.model_policy import get_automation_model_eligibility
from app.db import SearchSpace, User, async_session_maker
from app.utils.content_utils import extract_text_content
from .prompt import build_draft_prompt
@ -98,6 +99,27 @@ def create_create_automation_tool(
declined. Acknowledge once and stop do NOT retry or pitch
variants without a fresh user request.
"""
# --- 0. Eligibility gate (fail fast, before drafting + HITL) ---
# Automations may only use premium or BYOK models. Check up front so we
# don't make the user draft + approve a card that can't be saved.
async with async_session_maker() as session:
search_space = await session.get(SearchSpace, search_space_id)
if search_space is None:
return {
"status": "error",
"message": "search space not found in this session",
}
eligibility = get_automation_model_eligibility(search_space)
if not eligibility["allowed"]:
reasons = " ".join(v["reason"] for v in eligibility["violations"])
return {
"status": "error",
"message": (
f"{reasons} Update the search space's model settings to a "
"premium or your own (BYOK) model, then try again."
),
}
# --- 1. Draft via sub-LLM ---
prompt = build_draft_prompt(search_space_id=search_space_id, intent=intent)
try:

View file

@ -63,8 +63,14 @@ def _get_global_image_gen_config(config_id: int) -> dict | None:
def create_generate_image_tool(
search_space_id: int,
db_session: AsyncSession,
image_generation_config_id_override: int | None = None,
):
"""Create ``generate_image`` with bound search space; DB work uses a per-call session."""
"""Create ``generate_image`` with bound search space; DB work uses a per-call session.
``image_generation_config_id_override``: when set (automations running on a
captured model), use this config id instead of reading the search space's
live ``image_generation_config_id``.
"""
del db_session # use a fresh per-call session, see below
@tool
@ -108,19 +114,27 @@ def create_generate_image_tool(
# task's session is shared across every tool; without isolation,
# autoflushes from a concurrent writer poison this tool too.
async with shielded_async_session() as session:
result = await session.execute(
select(SearchSpace).filter(SearchSpace.id == search_space_id)
)
search_space = result.scalars().first()
if not search_space:
return _failed(
{"error": "Search space not found"},
error="Search space not found",
if image_generation_config_id_override is not None:
# Automation run: use the captured image model, insulated from
# later search-space changes. No search-space read needed.
config_id = (
image_generation_config_id_override or IMAGE_GEN_AUTO_MODE_ID
)
else:
result = await session.execute(
select(SearchSpace).filter(SearchSpace.id == search_space_id)
)
search_space = result.scalars().first()
if not search_space:
return _failed(
{"error": "Search space not found"},
error="Search space not found",
)
config_id = (
search_space.image_generation_config_id or IMAGE_GEN_AUTO_MODE_ID
)
config_id = (
search_space.image_generation_config_id
or IMAGE_GEN_AUTO_MODE_ID
)
# Build generation kwargs
# NOTE: size, quality, and style are intentionally NOT passed.

View file

@ -51,5 +51,8 @@ def load_tools(
create_generate_image_tool(
search_space_id=d["search_space_id"],
db_session=d["db_session"],
image_generation_config_id_override=d.get(
"image_generation_config_id_override"
),
),
]

View file

@ -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

View file

@ -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(

View file

@ -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]]

View file

@ -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,

View file

@ -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,
)

View file

@ -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",

View file

@ -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

View file

@ -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",
]

View file

@ -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,

View file

@ -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"])