diff --git a/core/database.py b/core/database.py index a9ad90b8b..2ea7ecf9c 100644 --- a/core/database.py +++ b/core/database.py @@ -457,6 +457,11 @@ class ModelEndpoint(TimestampMixin, Base): # can be toggled per-endpoint in the UI. NULL = unknown, falls # back to the model-name keyword heuristic in agent_loop.py. supports_tools = Column(Boolean, nullable=True, default=None) + # Per-endpoint LLM completion read-timeout override, in seconds. NULL = + # fall back to the global agent_stream_timeout_seconds setting. Useful for + # a single slow/local model that needs more headroom than every other + # configured endpoint. + stream_timeout_seconds = Column(Integer, nullable=True) # Per-user ownership. NULL = legacy/shared (visible to every user) — this # is the historical default. When non-null, the model picker only shows # the endpoint to that user (admins always see everything). @@ -1664,6 +1669,21 @@ def _migrate_add_notifications_enabled(): logging.getLogger(__name__).warning(f"notifications_enabled migration: {e}") +def _migrate_add_endpoint_stream_timeout(): + """Add stream_timeout_seconds column to model_endpoints (per-endpoint LLM + read-timeout override; null falls back to the global + agent_stream_timeout_seconds setting).""" + try: + with engine.connect() as conn: + cols = [r[1] for r in conn.execute(text("PRAGMA table_info(model_endpoints)"))] + if "stream_timeout_seconds" not in cols: + conn.execute(text("ALTER TABLE model_endpoints ADD COLUMN stream_timeout_seconds INTEGER")) + conn.commit() + logging.getLogger(__name__).info("Added stream_timeout_seconds column to model_endpoints") + except Exception as e: + logging.getLogger(__name__).warning(f"model_endpoints stream_timeout_seconds migration: {e}") + + def _migrate_add_crew_member_id(): """Add crew_member_id column to sessions and scheduled_tasks tables if missing.""" try: @@ -1972,6 +1992,7 @@ def init_db(): _migrate_encrypt_signatures() _migrate_encrypt_endpoint_keys() _migrate_backfill_task_folders() + _migrate_add_endpoint_stream_timeout() def _migrate_backfill_task_folders(): diff --git a/routes/model_routes.py b/routes/model_routes.py index 600150a66..27271f008 100644 --- a/routes/model_routes.py +++ b/routes/model_routes.py @@ -1971,6 +1971,7 @@ def setup_model_routes(model_discovery): "model_refresh_mode": _endpoint_refresh_mode(r, kind), "model_refresh_interval": getattr(r, "model_refresh_interval", None), "model_refresh_timeout": getattr(r, "model_refresh_timeout", None), + "stream_timeout_seconds": getattr(r, "stream_timeout_seconds", None), }) if upgraded_legacy_pins: db.commit() @@ -2569,6 +2570,10 @@ def setup_model_routes(model_discovery): if "model_refresh_timeout" in body: timeout = _parse_positive_int(body.get("model_refresh_timeout"), minimum=1, maximum=60) ep.model_refresh_timeout = timeout + if "stream_timeout_seconds" in body: + ep.stream_timeout_seconds = _parse_positive_int( + body.get("stream_timeout_seconds"), minimum=30, maximum=3600 + ) # Rotating an API key used to require DELETE+POST, which wiped # endpoint_url/model from every session referencing the old base # URL. Allow in-place updates so the admin can change the key @@ -2602,6 +2607,7 @@ def setup_model_routes(model_discovery): "model_refresh_mode": getattr(ep, "model_refresh_mode", None) or "auto", "model_refresh_interval": getattr(ep, "model_refresh_interval", None), "model_refresh_timeout": getattr(ep, "model_refresh_timeout", None), + "stream_timeout_seconds": getattr(ep, "stream_timeout_seconds", None), } finally: db.close() diff --git a/src/agent_loop.py b/src/agent_loop.py index cca93fe56..5d38d8c22 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -957,6 +957,38 @@ def _endpoint_lookup_keys(endpoint_url: str) -> List[str]: pass return keys + +def _resolve_endpoint_stream_timeout(endpoint_url: str) -> Optional[int]: + """Per-endpoint LLM stream timeout override (ModelEndpoint.stream_timeout_seconds), + or None if unset/unresolvable — caller falls back to the global setting.""" + try: + from core.database import SessionLocal as _SL, ModelEndpoint as _ME + db = _SL() + try: + for key in _endpoint_lookup_keys(endpoint_url): + ep = db.query(_ME).filter(_ME.base_url == key).first() + if ep is not None: + return getattr(ep, "stream_timeout_seconds", None) + finally: + db.close() + except Exception: + pass + return None + + +def resolve_stream_timeout(endpoint_url: str) -> int: + """Single source of truth for the LLM stream read-timeout: the target + ModelEndpoint's stream_timeout_seconds if set, else the global + agent_stream_timeout_seconds setting. Every caller of + stream_llm(_with_fallback) for a model completion should route through + this instead of reading get_setting("agent_stream_timeout_seconds", ...) + directly, so a new per-endpoint override applies everywhere without + touching each call site.""" + endpoint_override = _resolve_endpoint_stream_timeout(endpoint_url) + if endpoint_override: + return int(endpoint_override) + return int(get_setting("agent_stream_timeout_seconds", 300) or 300) + # Admin tool keywords — if the last user message contains any of these, include admin tools _ADMIN_KEYWORDS = [ "session", "sessions", "chat", "chats", "conversation", "conversations", @@ -3237,7 +3269,7 @@ async def stream_agent_loop( max_tokens=min(max_tokens or 128, 128), prompt_type=None, tools=None, - timeout=int(get_setting("agent_stream_timeout_seconds", 300) or 300), + timeout=resolve_stream_timeout(endpoint_url), session_id=session_id, workload=workload, ): @@ -3942,7 +3974,7 @@ async def stream_agent_loop( _last_content = _last_user.lower() _wants_mcp = any(kw in _last_content for kw in _MCP_KEYWORDS) all_tool_schemas = mcp_schemas if (_wants_mcp and mcp_schemas) else [] - agent_stream_timeout = int(get_setting("agent_stream_timeout_seconds", 300) or 300) + agent_stream_timeout = resolve_stream_timeout(endpoint_url) _tool_names_sent = [t.get("function", {}).get("name") for t in (all_tool_schemas or []) if t.get("function")] logger.info(f"[agent-debug] round={round_num} model={model} _is_api_model={_is_api_model} tools_sent={len(_tool_names_sent)} tool_names={_tool_names_sent[:15]} relevant_tools={sorted(_relevant_tools)[:15] if _relevant_tools else 'ALL'}")