mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-20 11:42:23 +02:00
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
"""Blocking bridge for interactive tool-ui components: AskUI parks here until the user
|
|
answers in the transcript (or the wait times out). Keyed by (session_id, component props.id),
|
|
so the frontend can respond without ever learning a server-side request id."""
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import Any, Dict, Literal, Optional, Tuple
|
|
|
|
from pydantic import BaseModel, ConfigDict, InstanceOf
|
|
from typeguard import typechecked
|
|
|
|
MAX_PENDING = 50
|
|
MAX_WAIT_SECONDS = 600.0
|
|
# The card is clickable the moment its tool_call broadcasts, but the wait only parks after the CLI hook round-trip, stdio dispatch and an HTTP hop; an instant click landing in that gap must not be dropped (ENG-232 D4).
|
|
EARLY_ANSWER_TTL_SECONDS = 45.0
|
|
MAX_EARLY = 50
|
|
|
|
RespondOutcome = Literal["delivered", "buffered", "gone"]
|
|
|
|
|
|
class PendingUiRequest(BaseModel):
|
|
model_config = ConfigDict(validate_assignment=True)
|
|
event: InstanceOf[asyncio.Event]
|
|
response: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
class EarlyAnswer(BaseModel):
|
|
model_config = ConfigDict(validate_assignment=True)
|
|
stamp: float
|
|
response: Dict[str, Any]
|
|
|
|
|
|
p_pending: Dict[Tuple[str, str], PendingUiRequest] = {}
|
|
p_early: Dict[Tuple[str, str], EarlyAnswer] = {}
|
|
|
|
|
|
@typechecked
|
|
def p_prune_early(now: float) -> None:
|
|
for key in [k for k, v in p_early.items() if now - v.stamp > EARLY_ANSWER_TTL_SECONDS]:
|
|
p_early.pop(key, None)
|
|
|
|
|
|
@typechecked
|
|
async def wait_for_ui_response(session_id: str, component_id: str, timeout_s: float) -> Optional[Dict[str, Any]]:
|
|
"""Registers the request and blocks until respond_to_ui_request fires it; None on timeout."""
|
|
key = (session_id, component_id)
|
|
p_prune_early(time.monotonic())
|
|
early = p_early.pop(key, None)
|
|
if early is not None:
|
|
return early.response
|
|
if len(p_pending) >= MAX_PENDING:
|
|
raise ValueError("too many pending UI requests")
|
|
# A retried tool call for the same component replaces the stale wait; the old waiter times out.
|
|
pending = PendingUiRequest(event=asyncio.Event())
|
|
p_pending[key] = pending
|
|
try:
|
|
await asyncio.wait_for(pending.event.wait(), timeout=min(timeout_s, MAX_WAIT_SECONDS))
|
|
return pending.response
|
|
except asyncio.TimeoutError:
|
|
return None
|
|
finally:
|
|
if p_pending.get(key) is pending:
|
|
p_pending.pop(key, None)
|
|
|
|
|
|
@typechecked
|
|
def respond_to_ui_request(session_id: str, component_id: str, response: Dict[str, Any]) -> RespondOutcome:
|
|
"""Delivers the user's answer to the parked wait, or holds it briefly for a wait still en route."""
|
|
pending = p_pending.get((session_id, component_id))
|
|
if pending is not None:
|
|
pending.response = response
|
|
pending.event.set()
|
|
return "delivered"
|
|
now = time.monotonic()
|
|
p_prune_early(now)
|
|
if len(p_early) >= MAX_EARLY:
|
|
return "gone"
|
|
p_early[(session_id, component_id)] = EarlyAnswer(stamp=now, response=response)
|
|
return "buffered"
|
|
|
|
|
|
@typechecked
|
|
def reset_ui_bridge() -> None:
|
|
p_pending.clear()
|
|
p_early.clear()
|
|
|
|
|
|
@typechecked
|
|
def cancel_session_waits(session_id: str) -> int:
|
|
"""Releases every parked wait for a stopped session so its cards can't eat later clicks (ENG-232 D5)."""
|
|
released = 0
|
|
for key, pending in list(p_pending.items()):
|
|
if key[0] == session_id:
|
|
pending.event.set()
|
|
released += 1
|
|
for key in [k for k in p_early if k[0] == session_id]:
|
|
p_early.pop(key, None)
|
|
return released
|