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>
92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
"""In-process lifetime of one ReAct turn.
|
|
|
|
The dictionaries below used to live for the whole backend process. A finished
|
|
turn drops its scratch memory and its conversation lock only after every
|
|
waiter for that conversation has left.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Dict, Optional
|
|
|
|
REACT_AGENT_MEMORY_CACHE: Dict[str, object] = {}
|
|
# Serialize ReAct runs per conversation. A second message waits on the same
|
|
# lock; explicit cancel uses REACT_CONVERSATION_AGENT_TASKS.
|
|
REACT_CONVERSATION_LOCKS: Dict[str, asyncio.Lock] = {}
|
|
REACT_CONVERSATION_AGENT_TASKS: Dict[str, "asyncio.Task[object]"] = {}
|
|
|
|
_LOCK_GUARD = asyncio.Lock()
|
|
_LOCK_REFS: Dict[str, int] = {}
|
|
|
|
|
|
async def acquire_conversation_run(conv_id: str) -> asyncio.Lock:
|
|
"""Take the per-conversation lock and count this waiter.
|
|
|
|
The refcount is incremented before waiting, so a finishing turn cannot
|
|
delete the lock out from under someone already queued on it.
|
|
"""
|
|
async with _LOCK_GUARD:
|
|
lock = REACT_CONVERSATION_LOCKS.get(conv_id)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
REACT_CONVERSATION_LOCKS[conv_id] = lock
|
|
_LOCK_REFS[conv_id] = _LOCK_REFS.get(conv_id, 0) + 1
|
|
try:
|
|
await lock.acquire()
|
|
except BaseException:
|
|
# Cancelled while queued. The lock was not taken; drop the reservation
|
|
# so a finished turn can still remove it once every other waiter is gone.
|
|
await _drop_reservation(conv_id, lock)
|
|
raise
|
|
return lock
|
|
|
|
|
|
async def release_conversation_run(conv_id: str, lock: asyncio.Lock) -> None:
|
|
"""Release one holder. Drop the lock only when nobody else is waiting."""
|
|
lock.release()
|
|
await _drop_reservation(conv_id, lock)
|
|
|
|
|
|
async def _drop_reservation(conv_id: str, lock: asyncio.Lock) -> None:
|
|
"""Forget one reserved waiter. Keep the lock while anyone else still needs it."""
|
|
async with _LOCK_GUARD:
|
|
refs = _LOCK_REFS.get(conv_id, 1) - 1
|
|
current = REACT_CONVERSATION_LOCKS.get(conv_id)
|
|
if refs <= 0 and current is lock and not lock.locked():
|
|
REACT_CONVERSATION_LOCKS.pop(conv_id, None)
|
|
_LOCK_REFS.pop(conv_id, None)
|
|
else:
|
|
_LOCK_REFS[conv_id] = max(refs, 0)
|
|
|
|
|
|
def release_cached_gpts_memory(conv_id: str) -> None:
|
|
"""Drop the turn's GptsMemory scratch. Persisted history is not deleted."""
|
|
memory = REACT_AGENT_MEMORY_CACHE.pop(conv_id, None)
|
|
if memory is None:
|
|
return
|
|
dispose = getattr(memory, "dispose_runtime", None)
|
|
if callable(dispose):
|
|
dispose(conv_id)
|
|
|
|
|
|
async def end_react_turn(conv_id: str, lock: Optional[asyncio.Lock]) -> None:
|
|
"""Free scratch memory, then release the conversation lock.
|
|
|
|
Call this from the stream ``finally`` after the agent task has finished
|
|
and after ``_react_terminal_events`` has already persisted the round.
|
|
The lock is still held while the cache is removed, so the next waiter
|
|
cannot enter until this turn's scratch is gone.
|
|
"""
|
|
try:
|
|
if conv_id:
|
|
release_cached_gpts_memory(conv_id)
|
|
finally:
|
|
if lock is not None and conv_id:
|
|
await release_conversation_run(conv_id, lock)
|
|
|
|
|
|
def conversation_lock_refcount(conv_id: str) -> int:
|
|
"""Test helper: holders plus waiters that have already reserved the lock."""
|
|
return _LOCK_REFS.get(conv_id, 0)
|