ryan
5809ad3adb
Shared runtime - One agent entry point (runtime.invoke_agent_for_user / prepare_agent_turn) for web, email, iMessage, WhatsApp, scheduled and cron, so identity, permissions, memories, summaries and metrics match across channels. - Serialize turns per checkpoint thread; /query validates input and returns 409 while a durable run owns the conversation. Reconnectable chat (durable runs) - POST /runs creates an idempotent, owner-scoped run; GET .../events replays an ordered SSE log by cursor, so generation survives refresh and device switches. Stop cancels; disconnect only detaches. /stream-query is a compatibility adapter. - Coalesce streamed deltas (150 ms) and wake same-process subscribers instead of polling; keep buffered text when a run is stopped. - Drain runs on shutdown; serve via Hypercorn with a shutdown timeout covering the drain (and no longer in debug mode). CHAT_RUN_* settings are forwarded. Scheduled messages and cron - Atomic delivery claims, attempt history and explicit retries; ambiguous sends become "uncertain" and are never resent automatically. - Cron triggers need a run_id/Idempotency-Key; legacy jobs need an explicit owner. - Cancelling never rewrites a started delivery; series stop via tombstones. Legacy sent rows are not forked into new series. Malformed bodies return 400. History, integrations and retrieval - Message-content search, pinned conversations, JSON/Markdown export and a memory editor. - Integration status panel and YNAB data-freshness provenance. - Date-filtered document/note search with stable citations and a source viewer. Migration 20 is generated by aerich. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
577 lines
21 KiB
Python
577 lines
21 KiB
Python
"""A browser connection is only a subscriber, never the owner of generation."""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
import time
|
||
import uuid
|
||
import weakref
|
||
from datetime import datetime, timedelta, timezone
|
||
from contextlib import aclosing
|
||
from contextvars import Context
|
||
|
||
from quart import Blueprint, current_app, jsonify, make_response, request
|
||
from quart_jwt_extended import get_jwt_identity
|
||
from tortoise.exceptions import IntegrityError
|
||
from tortoise.transactions import in_transaction
|
||
|
||
from blueprints.users.decorators import jwt_access_required
|
||
from blueprints.users.models import User
|
||
from .models import Conversation
|
||
from .logic import add_message_to_conversation
|
||
from . import runtime
|
||
from .run_models import ChatRun, ChatRunEvent
|
||
|
||
runs_blueprint = Blueprint("chat_runs", __name__, url_prefix="/api/conversation")
|
||
LEASE_SECONDS = max(10, float(os.getenv("CHAT_RUN_LEASE_SECONDS", "30")))
|
||
MAX_SECONDS = max(1, float(os.getenv("CHAT_RUN_MAX_SECONDS", "1800")))
|
||
_tasks = {} # strong refs, keyed by run id; DB is the authority across processes
|
||
# Same-process subscribers wait on these instead of polling; other processes
|
||
# still see new events through the polling fallback.
|
||
# Weak values: an entry lives only while some subscriber is waiting on it.
|
||
_wakeups: weakref.WeakValueDictionary[str, asyncio.Event] = (
|
||
weakref.WeakValueDictionary()
|
||
)
|
||
# Streamed deltas are coalesced so each write carries ~this much text.
|
||
FLUSH_SECONDS = 0.15
|
||
POLL_SECONDS = 0.5
|
||
_DELTA_TYPES = ("content", "thinking_delta")
|
||
|
||
|
||
def _notify(run_id):
|
||
event = _wakeups.pop(str(run_id), None)
|
||
if event is not None:
|
||
event.set()
|
||
|
||
|
||
def _waiter(run_id):
|
||
event = _wakeups.get(str(run_id))
|
||
if event is None:
|
||
event = _wakeups[str(run_id)] = asyncio.Event()
|
||
return event
|
||
|
||
|
||
_draining_apps = set()
|
||
|
||
|
||
def now():
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
class RunConflict(Exception):
|
||
def __init__(self, message, status=409):
|
||
super().__init__(message)
|
||
self.status = status
|
||
|
||
|
||
async def owned_run(owner, conversation_id, run_id):
|
||
try:
|
||
return await ChatRun.get_or_none(
|
||
id=run_id,
|
||
owner_id=owner,
|
||
conversation_id=conversation_id,
|
||
conversation__user_id=owner,
|
||
)
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
|
||
async def _locked(run_id, db):
|
||
# UPDATE takes a row lock on PostgreSQL and a write lock on SQLite, before reads.
|
||
await ChatRun.filter(id=run_id).using_db(db).update(updated_at=now())
|
||
return await ChatRun.get_or_none(id=run_id).using_db(db)
|
||
|
||
|
||
async def _save_answer(run, text):
|
||
if text and not run.answer_saved:
|
||
conversation = await Conversation.get(id=run.conversation_id)
|
||
user = await User.get(id=run.owner_id)
|
||
await add_message_to_conversation(conversation, text, "simba", user)
|
||
run.answer_saved = True
|
||
|
||
|
||
async def append(run_id, token, payload):
|
||
async with in_transaction() as db:
|
||
run = await _locked(run_id, db)
|
||
if (
|
||
run.status != "running"
|
||
or run.cancel_requested
|
||
or str(run.worker_token) != str(token)
|
||
):
|
||
raise asyncio.CancelledError()
|
||
run.sequence += 1
|
||
if payload.get("type") == "content":
|
||
run.partial += payload.get("delta", "")
|
||
await ChatRunEvent.create(
|
||
run_id=run.id, sequence=run.sequence, payload=payload, using_db=db
|
||
)
|
||
await run.save(using_db=db)
|
||
_notify(run_id)
|
||
|
||
|
||
async def finish(run_id, status, payload=None, token=None, tail=""):
|
||
"""Terminalize a run; `tail` is buffered content never appended as an event."""
|
||
async with in_transaction() as db:
|
||
run = await _locked(run_id, db)
|
||
if (
|
||
run is None
|
||
or run.status != "running"
|
||
or (token and str(run.worker_token) != str(token))
|
||
):
|
||
return
|
||
run.partial += tail
|
||
await _save_answer(run, run.partial)
|
||
if run.cancel_requested and status == "completed":
|
||
status = "cancelled"
|
||
run.status = status
|
||
run.active_key = None
|
||
for event in ([payload] if payload else []) + [
|
||
{"type": "done", "status": status}
|
||
]:
|
||
run.sequence += 1
|
||
await ChatRunEvent.create(
|
||
run_id=run.id, sequence=run.sequence, payload=event, using_db=db
|
||
)
|
||
await run.save(using_db=db)
|
||
_notify(run_id)
|
||
|
||
|
||
async def expire_runs(conversation_id=None):
|
||
query = ChatRun.filter(status="running", lease_until__lt=now())
|
||
if conversation_id:
|
||
query = query.filter(conversation_id=conversation_id)
|
||
for candidate in await query:
|
||
# Recheck under lock; a heartbeat may have arrived after the initial scan.
|
||
async with in_transaction() as db:
|
||
run = await _locked(candidate.id, db)
|
||
if run is None or run.status != "running" or run.lease_until >= now():
|
||
continue
|
||
await finish(
|
||
run.id,
|
||
"cancelled" if run.cancel_requested else "interrupted",
|
||
{
|
||
"type": "error",
|
||
"message": "Generation was interrupted by a server restart. Please retry.",
|
||
},
|
||
)
|
||
|
||
|
||
async def create_run(app, owner, data):
|
||
if not isinstance(data, dict):
|
||
raise RunConflict("JSON object required", 400)
|
||
conversation_id = data.get("conversation_id")
|
||
try:
|
||
conversation = await Conversation.get_or_none(id=conversation_id, user_id=owner)
|
||
except (ValueError, TypeError):
|
||
conversation = None
|
||
if conversation is None:
|
||
raise RunConflict("conversation not found", 404)
|
||
key = data.get("request_id", str(uuid.uuid4()))
|
||
if not isinstance(key, str) or not key or len(key) > 128:
|
||
raise RunConflict("request_id must be 1–128 characters", 400)
|
||
payload = {k: data.get(k) for k in ("query", "image_key", "file_key")}
|
||
if not isinstance(payload["query"], str):
|
||
raise RunConflict("query must be a string", 400)
|
||
for name, prefix in (("image_key", ""), ("file_key", "files/")):
|
||
value = payload[name]
|
||
# Matches the worker's checks, so a run is never created only to fail.
|
||
if value and (
|
||
not isinstance(value, str)
|
||
or ".." in value
|
||
or not value.startswith(f"conversations/{conversation.id}/{prefix}")
|
||
):
|
||
raise RunConflict(f"invalid {name}", 400)
|
||
await expire_runs(str(conversation.id))
|
||
|
||
async def existing():
|
||
run = await ChatRun.get_or_none(
|
||
owner_id=owner, conversation_id=conversation.id, request_key=key
|
||
)
|
||
if run and run.payload != payload:
|
||
raise RunConflict("request_id already used with different input")
|
||
return run
|
||
|
||
run = await existing()
|
||
if run:
|
||
return run, False
|
||
if id(app) in _draining_apps:
|
||
raise RunConflict("Server is draining; retry creation shortly", 503)
|
||
token = uuid.uuid4()
|
||
try:
|
||
run = await ChatRun.create(
|
||
owner_id=owner,
|
||
conversation_id=conversation.id,
|
||
request_key=key,
|
||
payload=payload,
|
||
active_key=str(conversation.id),
|
||
worker_token=token,
|
||
lease_until=now() + timedelta(seconds=LEASE_SECONDS),
|
||
)
|
||
except IntegrityError:
|
||
run = await existing()
|
||
if run:
|
||
return run, False
|
||
raise RunConflict("A run is already active in this conversation")
|
||
task = asyncio.create_task(_work(app, run.id, token), context=Context())
|
||
_tasks[str(run.id)] = (app, task)
|
||
|
||
def completed(task):
|
||
_tasks.pop(str(run.id), None)
|
||
if not task.cancelled() and task.exception():
|
||
logging.error(
|
||
"Chat run worker exited before finalization", exc_info=task.exception()
|
||
)
|
||
|
||
task.add_done_callback(completed)
|
||
return run, True
|
||
|
||
|
||
class _DeltaBatch:
|
||
"""Coalesce streamed deltas so a run writes ~FLUSH_SECONDS of text per event."""
|
||
|
||
def __init__(self, run_id, token):
|
||
self.run_id, self.token = run_id, token
|
||
self.frames: asyncio.Queue = asyncio.Queue() # raw SSE frames not yet handled
|
||
self.pending = None
|
||
self.last_flush = 0.0
|
||
|
||
def due(self):
|
||
"""Seconds until the pending delta must be written; None if nothing waits."""
|
||
if self.pending is None:
|
||
return None
|
||
return max(0.0, FLUSH_SECONDS - (time.monotonic() - self.last_flush))
|
||
|
||
def unsent_content(self):
|
||
"""Answer text that was produced but never written as an event."""
|
||
parts = []
|
||
if self.pending is not None and self.pending["type"] == "content":
|
||
parts.append(self.pending["delta"])
|
||
while not self.frames.empty():
|
||
item = self.frames.get_nowait()
|
||
if isinstance(item, str) and item.startswith("data: {"):
|
||
payload = json.loads(item.removeprefix("data: "))
|
||
if payload.get("type") == "content":
|
||
parts.append(payload.get("delta", ""))
|
||
return "".join(parts)
|
||
|
||
async def flush(self):
|
||
if self.pending is not None:
|
||
await append(self.run_id, self.token, self.pending)
|
||
self.pending = None
|
||
self.last_flush = time.monotonic()
|
||
|
||
async def add(self, payload):
|
||
if payload.get("type") not in _DELTA_TYPES:
|
||
await self.flush()
|
||
await append(self.run_id, self.token, payload)
|
||
elif self.pending is not None and self.pending["type"] == payload["type"]:
|
||
self.pending["delta"] += payload.get("delta", "")
|
||
else:
|
||
await self.flush()
|
||
self.pending = {**payload, "delta": payload.get("delta", "")}
|
||
|
||
|
||
async def _heartbeat(run_id, token, worker):
|
||
try:
|
||
while True:
|
||
await asyncio.sleep(min(1, LEASE_SECONDS / 3))
|
||
changed = await ChatRun.filter(
|
||
id=run_id,
|
||
worker_token=token,
|
||
status="running",
|
||
cancel_requested=False,
|
||
lease_until__gte=now(),
|
||
).update(lease_until=now() + timedelta(seconds=LEASE_SECONDS))
|
||
if not changed:
|
||
worker.cancel()
|
||
return
|
||
except Exception:
|
||
logging.exception("Chat run heartbeat failed")
|
||
worker.cancel()
|
||
|
||
|
||
async def _work(app, run_id, token):
|
||
heartbeat = asyncio.create_task(_heartbeat(run_id, token, asyncio.current_task()))
|
||
batch = _DeltaBatch(run_id, token)
|
||
async with app.app_context():
|
||
try:
|
||
async with asyncio.timeout(MAX_SECONDS):
|
||
from . import _chat_run_stream
|
||
|
||
run = await ChatRun.get(id=run_id)
|
||
conversation = await Conversation.get(id=run.conversation_id)
|
||
user = await User.get(id=run.owner_id)
|
||
|
||
async def save_user(text, image_key):
|
||
async with in_transaction() as db:
|
||
locked = await _locked(run_id, db)
|
||
if (
|
||
locked is None
|
||
or locked.status != "running"
|
||
or locked.cancel_requested
|
||
):
|
||
raise asyncio.CancelledError()
|
||
if not locked.user_saved:
|
||
await add_message_to_conversation(
|
||
conversation, text, "user", user, image_key=image_key
|
||
)
|
||
locked.user_saved = True
|
||
await locked.save(using_db=db)
|
||
|
||
async def save_answer(text):
|
||
async with in_transaction() as db:
|
||
locked = await _locked(run_id, db)
|
||
if locked is not None and locked.status == "running":
|
||
await _save_answer(locked, text)
|
||
await locked.save(using_db=db)
|
||
|
||
failed = False
|
||
async with (
|
||
runtime.thread_lock(str(conversation.id)),
|
||
aclosing(
|
||
_chat_run_stream(
|
||
{**run.payload, "conversation_id": str(conversation.id)},
|
||
user,
|
||
conversation,
|
||
save_user,
|
||
save_answer,
|
||
)
|
||
) as stream,
|
||
):
|
||
# The generator runs in one task (keeping its context across
|
||
# steps) so buffered text can be flushed on a timer even while
|
||
# the model emits nothing, e.g. while writing tool arguments.
|
||
frames = batch.frames
|
||
|
||
async def produce():
|
||
try:
|
||
async for frame in stream:
|
||
frames.put_nowait(frame)
|
||
frames.put_nowait(None)
|
||
except BaseException as exc:
|
||
frames.put_nowait(exc)
|
||
raise
|
||
|
||
producer = asyncio.create_task(produce())
|
||
try:
|
||
while True:
|
||
try:
|
||
item = await asyncio.wait_for(frames.get(), batch.due())
|
||
except TimeoutError:
|
||
await batch.flush()
|
||
continue
|
||
if item is None:
|
||
break
|
||
if isinstance(item, BaseException):
|
||
raise item
|
||
raw = item.removeprefix("data: ").strip()
|
||
if raw == "[DONE]":
|
||
continue
|
||
payload = json.loads(raw)
|
||
failed |= payload.get("type") == "error"
|
||
await batch.add(payload)
|
||
await batch.flush()
|
||
finally:
|
||
producer.cancel()
|
||
await asyncio.gather(producer, return_exceptions=True)
|
||
await finish(run_id, "failed" if failed else "completed", token=token)
|
||
except asyncio.CancelledError:
|
||
run = await ChatRun.get_or_none(id=run_id)
|
||
if run is None:
|
||
return
|
||
status = "cancelled" if run.cancel_requested else "interrupted"
|
||
await finish(
|
||
run_id,
|
||
status,
|
||
{
|
||
"type": "error",
|
||
"message": (
|
||
"Generation stopped."
|
||
if status == "cancelled"
|
||
else "Generation interrupted. Please retry."
|
||
),
|
||
},
|
||
token=token,
|
||
tail=batch.unsent_content(),
|
||
)
|
||
except Exception:
|
||
logging.exception("Chat generation failed")
|
||
await finish(
|
||
run_id,
|
||
"failed",
|
||
{
|
||
"type": "error",
|
||
"message": "Generation failed or exceeded its time limit. Please retry.",
|
||
},
|
||
token=token,
|
||
)
|
||
finally:
|
||
heartbeat.cancel()
|
||
await asyncio.gather(heartbeat, return_exceptions=True)
|
||
|
||
|
||
async def replay_response(run, cursor):
|
||
async def events():
|
||
nonlocal cursor
|
||
started = time.monotonic()
|
||
last_expiry = 0.0
|
||
while time.monotonic() - started < MAX_SECONDS + LEASE_SECONDS:
|
||
waiter = _waiter(run.id) # taken before reading, so no wakeup is lost
|
||
if time.monotonic() - last_expiry >= POLL_SECONDS:
|
||
await expire_runs(str(run.conversation_id))
|
||
last_expiry = time.monotonic()
|
||
batch = (
|
||
await ChatRunEvent.filter(run_id=run.id, sequence__gt=cursor)
|
||
.order_by("sequence")
|
||
.limit(200)
|
||
)
|
||
for event in batch:
|
||
cursor = event.sequence
|
||
if event.payload.get("type") == "done":
|
||
yield f"id: {cursor}\ndata: [DONE]\n\n"
|
||
return
|
||
yield f"id: {cursor}\ndata: {json.dumps(event.payload)}\n\n"
|
||
latest = await ChatRun.get_or_none(id=run.id)
|
||
if latest is None or (
|
||
latest.status != "running" and cursor >= latest.sequence
|
||
):
|
||
yield "data: [DONE]\n\n"
|
||
return
|
||
if not batch:
|
||
yield ": keepalive\n\n"
|
||
try:
|
||
await asyncio.wait_for(waiter.wait(), POLL_SECONDS)
|
||
except TimeoutError:
|
||
pass
|
||
# Finite subscriber lifetime; client reattaches with its last cursor.
|
||
|
||
response = await make_response(
|
||
events(),
|
||
200,
|
||
{
|
||
"Content-Type": "text/event-stream",
|
||
"Cache-Control": "no-cache",
|
||
"X-Accel-Buffering": "no",
|
||
"X-Chat-Run-ID": str(run.id),
|
||
},
|
||
)
|
||
response.timeout = None
|
||
return response
|
||
|
||
|
||
def serialize(run):
|
||
return {
|
||
"id": str(run.id),
|
||
"conversation_id": str(run.conversation_id),
|
||
"status": run.status,
|
||
"cursor": run.sequence,
|
||
"created_at": run.created_at.isoformat(),
|
||
"query": run.payload.get("query", ""),
|
||
"image_key": run.payload.get("image_key"),
|
||
}
|
||
|
||
|
||
@runs_blueprint.post("/runs")
|
||
@jwt_access_required
|
||
async def post_run():
|
||
try:
|
||
run, created = await create_run(
|
||
current_app._get_current_object(),
|
||
get_jwt_identity(),
|
||
await request.get_json(),
|
||
)
|
||
except RunConflict as exc:
|
||
return jsonify({"error": str(exc)}), exc.status
|
||
return jsonify(serialize(run)), 201 if created else 200
|
||
|
||
|
||
@runs_blueprint.get("/<conversation_id>/runs")
|
||
@jwt_access_required
|
||
async def discover(conversation_id):
|
||
owner = get_jwt_identity()
|
||
try:
|
||
conversation = await Conversation.get_or_none(id=conversation_id, user_id=owner)
|
||
except ValueError:
|
||
conversation = None
|
||
if conversation is None:
|
||
return jsonify({"error": "conversation not found"}), 404
|
||
await expire_runs(conversation_id)
|
||
# Latest terminal run included so restart/interruption can be shown after refresh.
|
||
run = (
|
||
await ChatRun.filter(conversation_id=conversation.id, owner_id=owner)
|
||
.order_by("-created_at")
|
||
.first()
|
||
)
|
||
return jsonify({"run": serialize(run) if run else None})
|
||
|
||
|
||
@runs_blueprint.get("/<conversation_id>/runs/<run_id>/events")
|
||
@jwt_access_required
|
||
async def get_events(conversation_id, run_id):
|
||
run = await owned_run(get_jwt_identity(), conversation_id, run_id)
|
||
if run is None:
|
||
return jsonify({"error": "run not found"}), 404
|
||
try:
|
||
cursor = int(
|
||
request.args.get("cursor", request.headers.get("Last-Event-ID", "0"))
|
||
)
|
||
if cursor < 0 or cursor > run.sequence:
|
||
raise ValueError()
|
||
except ValueError:
|
||
return jsonify({"error": "invalid cursor"}), 400
|
||
return await replay_response(run, cursor)
|
||
|
||
|
||
@runs_blueprint.post("/<conversation_id>/runs/<run_id>/cancel")
|
||
@jwt_access_required
|
||
async def cancel_run(conversation_id, run_id):
|
||
run = await owned_run(get_jwt_identity(), conversation_id, run_id)
|
||
if run is None:
|
||
return jsonify({"error": "run not found"}), 404
|
||
await ChatRun.filter(id=run.id, status="running").update(cancel_requested=True)
|
||
return jsonify(
|
||
{"status": "cancel_requested" if run.status == "running" else run.status}
|
||
)
|
||
|
||
|
||
def begin_drain(app):
|
||
"""Call when ingress/health starts draining, before awaiting run completion."""
|
||
_draining_apps.add(id(app))
|
||
|
||
|
||
def is_draining(app):
|
||
return id(app) in _draining_apps
|
||
|
||
|
||
async def drain_runs(app, timeout=None):
|
||
"""Root calls BEFORE closing agents/checkpointers/database, after stopping ingress."""
|
||
begin_drain(app)
|
||
claimed_ids = {
|
||
run_id for run_id, (task_app, _) in list(_tasks.items()) if task_app is app
|
||
}
|
||
tasks = [task for task_app, task in list(_tasks.values()) if task_app is app]
|
||
if tasks:
|
||
_, pending = await asyncio.wait(
|
||
tasks,
|
||
timeout=(
|
||
timeout
|
||
if timeout is not None
|
||
else float(os.getenv("CHAT_RUN_DRAIN_SECONDS", "1800"))
|
||
),
|
||
)
|
||
for task in pending:
|
||
task.cancel()
|
||
await asyncio.gather(*pending, return_exceptions=True)
|
||
# Covers a task cancelled before its coroutine first ran.
|
||
async with app.app_context():
|
||
for run in await ChatRun.filter(id__in=claimed_ids, status="running"):
|
||
await finish(
|
||
run.id,
|
||
"interrupted",
|
||
{
|
||
"type": "error",
|
||
"message": "Generation interrupted during shutdown. Please retry.",
|
||
},
|
||
)
|