javadandCursor
edbaa2acec
Connector calls, semantic rules, and DuckDB memory keep long runs stable, and stream retry stops chats from stalling. Co-authored-by: Cursor <cursoragent@cursor.com>
122 lines
3.5 KiB
Python
122 lines
3.5 KiB
Python
"""Map engine SSE payloads to product events, and forget finished runs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import threading
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from kamidso.integrations.event_adapter import EventAdapterState
|
|
|
|
logger = logging.getLogger("kamidso.engine")
|
|
|
|
_ADAPTER_CACHE: Dict[str, EventAdapterState] = {}
|
|
_ACTIVE_RUNS: set[str] = set()
|
|
_TOUCHED: Dict[str, float] = {}
|
|
_LOCK = threading.Lock()
|
|
_SWEEPER_STARTED = False
|
|
_DEFAULT_TTL_SECONDS = 900.0
|
|
|
|
|
|
def adapter_cache_ttl_seconds() -> float:
|
|
raw = (os.environ.get("KAMIDSO_ADAPTER_CACHE_TTL_SECONDS") or "").strip()
|
|
if not raw:
|
|
return _DEFAULT_TTL_SECONDS
|
|
try:
|
|
return max(1.0, float(raw))
|
|
except ValueError:
|
|
logger.warning("Ignoring invalid KAMIDSO_ADAPTER_CACHE_TTL_SECONDS=%s", raw)
|
|
return _DEFAULT_TTL_SECONDS
|
|
|
|
|
|
def begin_adapter_run(run_id: str) -> None:
|
|
"""Keep this run's adapter alive for the whole execution, including long tools."""
|
|
with _LOCK:
|
|
_ACTIVE_RUNS.add(run_id)
|
|
_TOUCHED[run_id] = time.monotonic()
|
|
_ensure_sweeper_locked()
|
|
|
|
|
|
def release_adapter_run(run_id: str) -> None:
|
|
"""Drop adapter state after the last event for this run has been produced.
|
|
|
|
Safe on success, failure, cancellation, and generator close. Does not
|
|
remove a different run that is still active.
|
|
"""
|
|
with _LOCK:
|
|
_ACTIVE_RUNS.discard(run_id)
|
|
_ADAPTER_CACHE.pop(run_id, None)
|
|
_TOUCHED.pop(run_id, None)
|
|
|
|
|
|
def sweep_adapter_cache(*, now: Optional[float] = None, keep: Optional[str] = None) -> List[str]:
|
|
"""Forget inactive adapters that have not been touched within the TTL.
|
|
|
|
Active runs are kept even if they have been quiet, so a long SQL or code
|
|
step is not discarded mid-flight.
|
|
"""
|
|
moment = time.monotonic() if now is None else now
|
|
ttl = adapter_cache_ttl_seconds()
|
|
removed: List[str] = []
|
|
with _LOCK:
|
|
stale = [
|
|
run_id
|
|
for run_id, touched in _TOUCHED.items()
|
|
if run_id != keep and run_id not in _ACTIVE_RUNS and moment - touched >= ttl
|
|
]
|
|
for run_id in stale:
|
|
_ADAPTER_CACHE.pop(run_id, None)
|
|
_TOUCHED.pop(run_id, None)
|
|
removed.append(run_id)
|
|
return removed
|
|
|
|
|
|
def map_engine_event(
|
|
run_id: str, engine_event: Dict[str, Any], locale: str = "fa-IR"
|
|
) -> List[Dict[str, Any]]:
|
|
with _LOCK:
|
|
_TOUCHED[run_id] = time.monotonic()
|
|
adapter = _ADAPTER_CACHE.get(run_id)
|
|
if adapter is None:
|
|
adapter = EventAdapterState(run_id, locale=locale)
|
|
_ADAPTER_CACHE[run_id] = adapter
|
|
_ensure_sweeper_locked()
|
|
events = adapter.map(engine_event)
|
|
sweep_adapter_cache(keep=run_id)
|
|
return events
|
|
|
|
|
|
def adapter_cache_size() -> int:
|
|
with _LOCK:
|
|
return len(_ADAPTER_CACHE)
|
|
|
|
|
|
def adapter_run_is_active(run_id: str) -> bool:
|
|
with _LOCK:
|
|
return run_id in _ACTIVE_RUNS
|
|
|
|
|
|
def _ensure_sweeper_locked() -> None:
|
|
global _SWEEPER_STARTED
|
|
if _SWEEPER_STARTED:
|
|
return
|
|
_SWEEPER_STARTED = True
|
|
thread = threading.Thread(
|
|
target=_sweep_loop,
|
|
name="kamidso-adapter-cache",
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
|
|
|
|
def _sweep_loop() -> None:
|
|
while True:
|
|
ttl = adapter_cache_ttl_seconds()
|
|
time.sleep(min(60.0, ttl))
|
|
try:
|
|
sweep_adapter_cache()
|
|
except Exception:
|
|
logger.debug("adapter cache sweep failed", exc_info=True)
|