Skip to content

Commit 66abcc6

Browse files
committed
Merge PR #24: Implement acb-compose-shell and message sync mechanism
- Add reply token mechanism to prevent message replay attacks - Add sequence synchronization for conflict detection - Implement /api/threads/{thread_id}/sync-context endpoint - Add new config options: REPLY_TOKEN_LEASE_SECONDS, SEQ_TOLERANCE, SEQ_MISMATCH_MAX_MESSAGES - Update shared-chat.js to use sync-context before posting messages - Fix test files to use new msg_post signature with expected_last_seq and reply_token - Resolve merge conflicts with main branch (templates API + sync API coexist)
2 parents e7e93e5 + f0abaff commit 66abcc6

19 files changed

Lines changed: 1135 additions & 559 deletions

src/config.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,11 @@
4040
MSG_WAIT_TIMEOUT = int(os.getenv("AGENTCHATBUS_WAIT_TIMEOUT", config_data.get("MSG_WAIT_TIMEOUT", "300")))
4141
BUS_VERSION = "0.1.0"
4242

43+
# Strict message sync mode (mandatory)
44+
REPLY_TOKEN_LEASE_SECONDS = int(os.getenv("AGENTCHATBUS_REPLY_TOKEN_LEASE_SECONDS", "10"))
45+
SEQ_TOLERANCE = int(os.getenv("AGENTCHATBUS_SEQ_TOLERANCE", "5"))
46+
SEQ_MISMATCH_MAX_MESSAGES = int(os.getenv("AGENTCHATBUS_SEQ_MISMATCH_MAX_MESSAGES", "20"))
47+
4348
# Rate limiting: max messages per minute per author identity (0 = disabled)
4449
RATE_LIMIT_MSG_PER_MINUTE = int(os.getenv("AGENTCHATBUS_RATE_LIMIT", "30"))
4550
RATE_LIMIT_ENABLED = RATE_LIMIT_MSG_PER_MINUTE > 0
@@ -62,6 +67,9 @@ def get_config_dict():
6267
"PORT": PORT,
6368
"AGENT_HEARTBEAT_TIMEOUT": AGENT_HEARTBEAT_TIMEOUT,
6469
"MSG_WAIT_TIMEOUT": MSG_WAIT_TIMEOUT,
70+
"REPLY_TOKEN_LEASE_SECONDS": REPLY_TOKEN_LEASE_SECONDS,
71+
"SEQ_TOLERANCE": SEQ_TOLERANCE,
72+
"SEQ_MISMATCH_MAX_MESSAGES": SEQ_MISMATCH_MAX_MESSAGES,
6573
"EXPOSE_THREAD_RESOURCES": EXPOSE_THREAD_RESOURCES,
6674
}
6775

src/db/crud.py

Lines changed: 165 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,15 @@
1313
import aiosqlite
1414

1515
from 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+
)
1825
from src.content_filter import check_content, ContentFilterError
1926

2027
logger = 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+
3380
GLOBAL_SYSTEM_PROMPT = """**SYSTEM DIRECTIVE: ACTIVE AGENT COLLABORATION WORKSPACE**
3481
3582
Welcome 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+
258368
def _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)

src/db/database.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,24 @@ async def init_schema(db: aiosqlite.Connection) -> None:
212212
CREATE INDEX IF NOT EXISTS idx_messages_thread_seq
213213
ON messages(thread_id, seq);
214214
215+
-- ----------------------------------------------------------------
216+
-- Reply token lease: mandatory sync token for msg_post in strict mode
217+
-- ----------------------------------------------------------------
218+
CREATE TABLE IF NOT EXISTS reply_tokens (
219+
token TEXT PRIMARY KEY,
220+
thread_id TEXT NOT NULL REFERENCES threads(id),
221+
agent_id TEXT,
222+
issued_at TEXT NOT NULL,
223+
expires_at TEXT NOT NULL,
224+
consumed_at TEXT,
225+
status TEXT NOT NULL CHECK (status IN ('issued', 'consumed', 'expired'))
226+
);
227+
228+
CREATE INDEX IF NOT EXISTS idx_reply_tokens_thread_status
229+
ON reply_tokens(thread_id, status);
230+
CREATE INDEX IF NOT EXISTS idx_reply_tokens_expires_at
231+
ON reply_tokens(expires_at);
232+
215233
-- ----------------------------------------------------------------
216234
-- Sequence counter: single-row table for thread-safe seq increment
217235
-- ----------------------------------------------------------------

src/main.py

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,14 @@
2828
from src.config import HOST, PORT, get_config_dict, save_config_dict
2929
from src.db.database import get_db, close_db
3030
from src.db import crud
31-
from src.db.crud import RateLimitExceeded
31+
from src.db.crud import (
32+
RateLimitExceeded,
33+
MissingSyncFieldsError,
34+
SeqMismatchError,
35+
ReplyTokenInvalidError,
36+
ReplyTokenExpiredError,
37+
ReplyTokenReplayError,
38+
)
3239
from src.config import THREAD_TIMEOUT_ENABLED, THREAD_TIMEOUT_MINUTES, THREAD_TIMEOUT_SWEEP_INTERVAL, RELOAD_ENABLED
3340
from src.mcp_server import server as mcp_server, _session_language
3441
from src.content_filter import ContentFilterError
@@ -397,10 +404,16 @@ class MessageCreate(BaseModel):
397404
author: str = "human"
398405
role: Literal["user", "assistant", "system"] = "user"
399406
content: str
407+
expected_last_seq: int
408+
reply_token: str
400409
mentions: list[str] | None = None
401410
metadata: dict | None = None
402411
images: list[dict] | None = None # [{url: str, name: str}, ...]
403412

413+
class SyncContextRequest(BaseModel):
414+
agent_id: str | None = None
415+
416+
404417
@app.get("/api/templates")
405418
async def api_list_templates():
406419
try:
@@ -475,6 +488,23 @@ async def api_delete_template(template_id: str):
475488
raise HTTPException(status_code=403, detail=err)
476489

477490

491+
@app.post("/api/threads/{thread_id}/sync-context")
492+
async def api_sync_context(thread_id: str, body: SyncContextRequest | None = None):
493+
try:
494+
db = await asyncio.wait_for(get_db(), timeout=DB_TIMEOUT)
495+
t = await asyncio.wait_for(crud.thread_get(db, thread_id), timeout=DB_TIMEOUT)
496+
except asyncio.TimeoutError:
497+
raise HTTPException(status_code=503, detail="Database operation timeout")
498+
if t is None:
499+
raise HTTPException(status_code=404, detail="Thread not found")
500+
501+
agent_id = body.agent_id if body else None
502+
sync = await asyncio.wait_for(
503+
crud.issue_reply_token(db, thread_id=thread_id, agent_id=agent_id),
504+
timeout=DB_TIMEOUT,
505+
)
506+
return sync
507+
478508
@app.post("/api/threads", status_code=201)
479509
async def api_create_thread(body: ThreadCreate):
480510
try:
@@ -509,10 +539,44 @@ async def api_post_message(thread_id: str, body: MessageCreate):
509539
try:
510540
m = await asyncio.wait_for(
511541
crud.msg_post(db, thread_id=thread_id, author=body.author,
512-
content=body.content, role=body.role,
542+
content=body.content,
543+
expected_last_seq=body.expected_last_seq,
544+
reply_token=body.reply_token,
545+
role=body.role,
513546
metadata=msg_metadata if msg_metadata else None),
514547
timeout=DB_TIMEOUT
515548
)
549+
except MissingSyncFieldsError as e:
550+
raise HTTPException(status_code=400, detail={
551+
"error": "MISSING_SYNC_FIELDS",
552+
"missing_fields": e.missing_fields,
553+
"action": "CALL_SYNC_CONTEXT_THEN_RETRY",
554+
})
555+
except SeqMismatchError as e:
556+
raise HTTPException(status_code=409, detail={
557+
"error": "SEQ_MISMATCH",
558+
"expected_last_seq": e.expected_last_seq,
559+
"current_seq": e.current_seq,
560+
"new_messages": e.new_messages,
561+
"action": "RE_READ_AND_RETRY",
562+
})
563+
except ReplyTokenInvalidError:
564+
raise HTTPException(status_code=400, detail={
565+
"error": "TOKEN_INVALID",
566+
"action": "CALL_SYNC_CONTEXT_THEN_RETRY",
567+
})
568+
except ReplyTokenExpiredError as e:
569+
raise HTTPException(status_code=400, detail={
570+
"error": "TOKEN_EXPIRED",
571+
"expires_at": e.expires_at,
572+
"action": "CALL_SYNC_CONTEXT_THEN_RETRY",
573+
})
574+
except ReplyTokenReplayError as e:
575+
raise HTTPException(status_code=400, detail={
576+
"error": "TOKEN_REPLAY",
577+
"consumed_at": e.consumed_at,
578+
"action": "CALL_SYNC_CONTEXT_THEN_RETRY",
579+
})
516580
except ContentFilterError as e:
517581
raise HTTPException(status_code=400, detail={"error": "Content blocked by filter", "pattern": e.pattern_name})
518582
except asyncio.TimeoutError:

src/mcp_server.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,14 @@ async def list_tools() -> list[types.Tool]:
160160
"thread_id": {"type": "string"},
161161
"author": {"type": "string", "description": "Agent ID, 'system', or 'human'."},
162162
"content": {"type": "string"},
163+
"expected_last_seq": {
164+
"type": "integer",
165+
"description": "Strict sync field. Thread seq the sender used as context baseline.",
166+
},
167+
"reply_token": {
168+
"type": "string",
169+
"description": "Strict sync field. Unconsumed reply token from msg_wait.",
170+
},
163171
"role": {"type": "string", "enum": ["user", "assistant", "system"], "default": "user"},
164172
"mentions": {
165173
"type": "array",
@@ -187,7 +195,7 @@ async def list_tools() -> list[types.Tool]:
187195
},
188196
},
189197
},
190-
"required": ["thread_id", "author", "content"],
198+
"required": ["thread_id", "author", "content", "expected_last_seq", "reply_token"],
191199
},
192200
),
193201
types.Tool(
@@ -223,6 +231,8 @@ async def list_tools() -> list[types.Tool]:
223231
description=(
224232
"Block until at least one new message arrives in the thread after `after_seq`. "
225233
"Returns immediately if messages are already available. "
234+
"Always includes sync context (`current_seq`, `reply_token`, `reply_window`) "
235+
"for the next strict `msg_post` call. "
226236
"If this tool returns an empty list (timeout), avoid spammy waiting messages, "
227237
"but after repeated timeouts you SHOULD send a concise, meaningful progress update "
228238
"(status/blocker/next action) and optionally @mention a relevant online agent."

0 commit comments

Comments
 (0)