1313import aiosqlite
1414
1515from src .db .models import Thread , Message , AgentInfo , Event , ThreadTemplate
16- from src .config import AGENT_HEARTBEAT_TIMEOUT , RATE_LIMIT_MSG_PER_MINUTE , RATE_LIMIT_ENABLED
17- from src .config import AGENT_HEARTBEAT_TIMEOUT , CONTENT_FILTER_ENABLED
16+ from src .config import (
17+ AGENT_HEARTBEAT_TIMEOUT ,
18+ RATE_LIMIT_MSG_PER_MINUTE ,
19+ RATE_LIMIT_ENABLED ,
20+ CONTENT_FILTER_ENABLED ,
21+ REPLY_TOKEN_LEASE_SECONDS ,
22+ SEQ_TOLERANCE ,
23+ SEQ_MISMATCH_MAX_MESSAGES ,
24+ )
1825from src .content_filter import check_content , ContentFilterError
1926
2027logger = logging .getLogger (__name__ )
@@ -30,6 +37,46 @@ def __init__(self, limit: int, window: int, retry_after: int, scope: str) -> Non
3037 super ().__init__ (f"Rate limit exceeded: { limit } messages/{ window } s" )
3138
3239
40+ class MissingSyncFieldsError (Exception ):
41+ """Raised when strict sync fields are absent from msg_post."""
42+
43+ def __init__ (self , missing_fields : list [str ]) -> None :
44+ self .missing_fields = missing_fields
45+ super ().__init__ (f"Missing required sync fields: { ', ' .join (missing_fields )} " )
46+
47+
48+ class SeqMismatchError (Exception ):
49+ """Raised when too many unseen messages exist since expected seq."""
50+
51+ def __init__ (self , expected_last_seq : int , current_seq : int , new_messages : list [dict ]) -> None :
52+ self .expected_last_seq = expected_last_seq
53+ self .current_seq = current_seq
54+ self .new_messages = new_messages
55+ super ().__init__ (
56+ f"SEQ_MISMATCH: expected_last_seq={ expected_last_seq } , current_seq={ current_seq } "
57+ )
58+
59+
60+ class ReplyTokenInvalidError (Exception ):
61+ def __init__ (self , token : str ) -> None :
62+ self .token = token
63+ super ().__init__ ("TOKEN_INVALID" )
64+
65+
66+ class ReplyTokenExpiredError (Exception ):
67+ def __init__ (self , token : str , expires_at : str ) -> None :
68+ self .token = token
69+ self .expires_at = expires_at
70+ super ().__init__ ("TOKEN_EXPIRED" )
71+
72+
73+ class ReplyTokenReplayError (Exception ):
74+ def __init__ (self , token : str , consumed_at : Optional [str ]) -> None :
75+ self .token = token
76+ self .consumed_at = consumed_at
77+ super ().__init__ ("TOKEN_REPLAY" )
78+
79+
3380GLOBAL_SYSTEM_PROMPT = """**SYSTEM DIRECTIVE: ACTIVE AGENT COLLABORATION WORKSPACE**
3481
3582Welcome to this Thread. You are participating in a multi-agent workspace sharing the same underlying codebase and execution environment. You MUST collaborate proactively and keep progress moving.
@@ -255,6 +302,69 @@ async def thread_latest_seq(db: aiosqlite.Connection, thread_id: str) -> int:
255302 return row ["max_seq" ] or 0
256303
257304
305+ async def _expire_old_reply_tokens (db : aiosqlite .Connection ) -> None :
306+ now = _now ()
307+ await db .execute (
308+ "UPDATE reply_tokens SET status = 'expired' "
309+ "WHERE status = 'issued' AND expires_at <= ?" ,
310+ (now ,),
311+ )
312+ await db .commit ()
313+
314+
315+ async def issue_reply_token (
316+ db : aiosqlite .Connection ,
317+ thread_id : str ,
318+ agent_id : Optional [str ] = None ,
319+ ) -> dict :
320+ """Issue a short-lived reply token bound to a thread (and optionally an agent)."""
321+ await _expire_old_reply_tokens (db )
322+ token = secrets .token_urlsafe (24 )
323+ issued_at = _now ()
324+ expires_at = (datetime .now (timezone .utc ) + timedelta (seconds = REPLY_TOKEN_LEASE_SECONDS )).isoformat ()
325+ await db .execute (
326+ "INSERT INTO reply_tokens (token, thread_id, agent_id, issued_at, expires_at, consumed_at, status) "
327+ "VALUES (?, ?, ?, ?, ?, NULL, 'issued')" ,
328+ (token , thread_id , agent_id , issued_at , expires_at ),
329+ )
330+ await db .commit ()
331+ current_seq = await thread_latest_seq (db , thread_id )
332+ return {
333+ "reply_token" : token ,
334+ "current_seq" : current_seq ,
335+ "reply_window" : {
336+ "expires_at" : expires_at ,
337+ "max_new_messages" : SEQ_TOLERANCE ,
338+ },
339+ }
340+
341+
342+ async def _get_new_messages_since (
343+ db : aiosqlite .Connection ,
344+ thread_id : str ,
345+ expected_last_seq : int ,
346+ limit : int = SEQ_MISMATCH_MAX_MESSAGES ,
347+ ) -> list [dict ]:
348+ msgs = await msg_list (
349+ db ,
350+ thread_id = thread_id ,
351+ after_seq = expected_last_seq ,
352+ limit = limit ,
353+ include_system_prompt = False ,
354+ )
355+ return [
356+ {
357+ "msg_id" : m .id ,
358+ "seq" : m .seq ,
359+ "author" : m .author ,
360+ "role" : m .role ,
361+ "content" : m .content ,
362+ "created_at" : m .created_at .isoformat (),
363+ }
364+ for m in msgs
365+ ]
366+
367+
258368def _row_to_thread (row : aiosqlite .Row ) -> Thread :
259369 keys = row .keys ()
260370 system_prompt = row ["system_prompt" ] if "system_prompt" in keys else None
@@ -366,6 +476,8 @@ async def msg_post(
366476 thread_id : str ,
367477 author : str ,
368478 content : str ,
479+ expected_last_seq : int ,
480+ reply_token : str ,
369481 role : str = "user" ,
370482 metadata : Optional [dict ] = None ,
371483) -> Message :
@@ -414,6 +526,41 @@ async def msg_post(
414526 scope = scope ,
415527 )
416528
529+ missing_fields : list [str ] = []
530+ if expected_last_seq is None :
531+ missing_fields .append ("expected_last_seq" )
532+ if not reply_token :
533+ missing_fields .append ("reply_token" )
534+ if missing_fields :
535+ raise MissingSyncFieldsError (missing_fields )
536+
537+ await _expire_old_reply_tokens (db )
538+ async with db .execute (
539+ "SELECT token, thread_id, agent_id, expires_at, consumed_at, status "
540+ "FROM reply_tokens WHERE token = ?" ,
541+ (reply_token ,),
542+ ) as cur :
543+ token_row = await cur .fetchone ()
544+
545+ if token_row is None :
546+ raise ReplyTokenInvalidError (reply_token )
547+ if token_row ["thread_id" ] != thread_id :
548+ raise ReplyTokenInvalidError (reply_token )
549+ if token_row ["status" ] == "consumed" :
550+ raise ReplyTokenReplayError (reply_token , token_row ["consumed_at" ])
551+ if token_row ["status" ] == "expired" :
552+ raise ReplyTokenExpiredError (reply_token , token_row ["expires_at" ])
553+
554+ token_agent_id = token_row ["agent_id" ]
555+ if token_agent_id and author_id and token_agent_id != author_id :
556+ raise ReplyTokenInvalidError (reply_token )
557+
558+ current_seq = await thread_latest_seq (db , thread_id )
559+ new_messages_count = current_seq - expected_last_seq
560+ if new_messages_count > SEQ_TOLERANCE :
561+ new_messages = await _get_new_messages_since (db , thread_id , expected_last_seq )
562+ raise SeqMismatchError (expected_last_seq , current_seq , new_messages )
563+
417564 mid = str (uuid .uuid4 ())
418565 now = _now ()
419566 seq = await next_seq (db )
@@ -427,6 +574,22 @@ async def msg_post(
427574 "UPDATE threads SET updated_at = ? WHERE id = ?" ,
428575 (now , thread_id ),
429576 )
577+ async with db .execute (
578+ "UPDATE reply_tokens SET status = 'consumed', consumed_at = ? "
579+ "WHERE token = ? AND status = 'issued'" ,
580+ (now , reply_token ),
581+ ) as cur :
582+ consumed = cur .rowcount
583+ if consumed == 0 :
584+ await db .rollback ()
585+ async with db .execute (
586+ "SELECT consumed_at FROM reply_tokens WHERE token = ?" ,
587+ (reply_token ,),
588+ ) as cur :
589+ row = await cur .fetchone ()
590+ consumed_at = row ["consumed_at" ] if row else None
591+ raise ReplyTokenReplayError (reply_token , consumed_at )
592+
430593 await db .commit ()
431594 if author_id :
432595 await agent_msg_post (db , author_id )
0 commit comments