fix: persist room operation cancellation across lost admission acknowledgments
This commit is contained in:
1 parent
523f6cbbf9
commit
2ad9f36f4f
3 files changed
+62
No files matched your search
@@ -185,5 +185,14 @@ async def voice_cancel(room: str, rid: str, request: Request):
|
|||||||
return {'run': await request.app.state.chat_runs.cancel(account['id'], rid)}
|
return {'run': await request.app.state.chat_runs.cancel(account['id'], rid)}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post('/internal/chat/rooms/{room}/operations/{key}/cancel', dependencies=[Depends(internal_authorized)])
|
||||||
|
async def voice_cancel_operation(room: str, key: str, request: Request):
|
||||||
|
account = await room_owner(request, room)
|
||||||
|
try:
|
||||||
|
return await request.app.state.chat_runs.cancel_operation(account['id'], account['conversation_id'], key)
|
||||||
|
except KeyError:
|
||||||
|
raise HTTPException(404, 'Chat resource not found') from None
|
||||||
|
|
||||||
|
|
||||||
def create_chat_router():
|
def create_chat_router():
|
||||||
return router
|
return router
|
||||||
@@ -108,6 +108,10 @@ class ChatRuns:
|
|||||||
accounts.db.execute('''CREATE TABLE IF NOT EXISTS chat_transport (
|
accounts.db.execute('''CREATE TABLE IF NOT EXISTS chat_transport (
|
||||||
run_id TEXT PRIMARY KEY, owner TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
run_id TEXT PRIMARY KEY, owner TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||||
conversation_id TEXT NOT NULL, metadata TEXT NOT NULL)''')
|
conversation_id TEXT NOT NULL, metadata TEXT NOT NULL)''')
|
||||||
|
accounts.db.execute('''CREATE TABLE IF NOT EXISTS chat_operation_cancellations (
|
||||||
|
owner TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||||
|
conversation_id TEXT NOT NULL, operation_key TEXT NOT NULL,
|
||||||
|
PRIMARY KEY(owner,conversation_id,operation_key))''')
|
||||||
|
|
||||||
def capabilities(self):
|
def capabilities(self):
|
||||||
enabled = bool(getattr(self.adapter, 'available', True))
|
enabled = bool(getattr(self.adapter, 'available', True))
|
||||||
@@ -132,6 +136,28 @@ class ChatRuns:
|
|||||||
row = self.accounts.db.execute('SELECT metadata FROM chat_transport WHERE owner=? AND run_id=?', (owner, rid)).fetchone()
|
row = self.accounts.db.execute('SELECT metadata FROM chat_transport WHERE owner=? AND run_id=?', (owner, rid)).fetchone()
|
||||||
return json.loads(row[0]) if row else None
|
return json.loads(row[0]) if row else None
|
||||||
|
|
||||||
|
def operation_cancelled(self, owner, cid, key):
|
||||||
|
with self.accounts.lock:
|
||||||
|
return self.accounts.db.execute('SELECT 1 FROM chat_operation_cancellations WHERE owner=? AND conversation_id=? AND operation_key=?', (owner, cid, key)).fetchone() is not None
|
||||||
|
|
||||||
|
async def cancel_operation(self, owner, cid, key):
|
||||||
|
if not isinstance(key, str) or not 1 <= len(key) <= 200 or '\x00' in key:
|
||||||
|
raise HTTPException(400, 'Invalid operation key')
|
||||||
|
async with self.admission_locks[owner]:
|
||||||
|
await self.call('get_conversation', owner, cid)
|
||||||
|
def remember():
|
||||||
|
with self.accounts.lock, self.accounts.db:
|
||||||
|
existing = self.accounts.db.execute('SELECT 1 FROM chat_operation_cancellations WHERE owner=? AND conversation_id=? AND operation_key=?', (owner, cid, key)).fetchone()
|
||||||
|
if not existing and self.accounts.db.execute('SELECT count(*) FROM chat_operation_cancellations WHERE owner=?', (owner,)).fetchone()[0] >= 10000:
|
||||||
|
raise HTTPException(429, 'Operation cancellation limit reached')
|
||||||
|
self.accounts.db.execute('INSERT OR IGNORE INTO chat_operation_cancellations VALUES(?,?,?)', (owner, cid, key))
|
||||||
|
row = self.accounts.db.execute('SELECT run_id FROM chat_idempotency WHERE owner=? AND conversation_id=? AND key=?', (owner, cid, key)).fetchone()
|
||||||
|
return row[0] if row else None
|
||||||
|
rid = await asyncio.to_thread(remember)
|
||||||
|
if rid:
|
||||||
|
return {'ok': True, 'run': await self.cancel(owner, rid)}
|
||||||
|
return {'ok': True, 'run': None}
|
||||||
|
|
||||||
def authorized(self, owner, metadata):
|
def authorized(self, owner, metadata):
|
||||||
with self.accounts.lock:
|
with self.accounts.lock:
|
||||||
if metadata.get('room'):
|
if metadata.get('room'):
|
||||||
@@ -176,6 +202,10 @@ class ChatRuns:
|
|||||||
raise HTTPException(429, 'Chat queue full')
|
raise HTTPException(429, 'Chat queue full')
|
||||||
operation = await self.call('send', owner, cid, content, idempotency_key, attachment_ids, mode=mode)
|
operation = await self.call('send', owner, cid, content, idempotency_key, attachment_ids, mode=mode)
|
||||||
rid = operation['run']['id']
|
rid = operation['run']['id']
|
||||||
|
if await asyncio.to_thread(self.operation_cancelled, owner, cid, idempotency_key):
|
||||||
|
if operation['run']['status'] not in TERMINAL:
|
||||||
|
operation['run'] = await self.cancel(owner, rid)
|
||||||
|
return operation
|
||||||
metadata = await asyncio.to_thread(self.metadata, owner, rid)
|
metadata = await asyncio.to_thread(self.metadata, owner, rid)
|
||||||
if metadata is None:
|
if metadata is None:
|
||||||
settings = await asyncio.to_thread(self.accounts.get_model_settings, owner)
|
settings = await asyncio.to_thread(self.accounts.get_model_settings, owner)
|
||||||
|
|||||||
@@ -31,3 +31,26 @@ def test_event_capacity_failure_persists_stop_and_reserved_terminal(tmp_path):
|
|||||||
assert store.events(owner['id'],rid)[-1]['event'] == 'done'
|
assert store.events(owner['id'],rid)[-1]['event'] == 'done'
|
||||||
accounts.close()
|
accounts.close()
|
||||||
asyncio.run(scenario())
|
asyncio.run(scenario())
|
||||||
|
|
||||||
|
|
||||||
|
def test_cancel_key_before_admission_never_starts_and_is_owner_scoped(tmp_path):
|
||||||
|
async def scenario():
|
||||||
|
accounts = Accounts(str(tmp_path/'cancel.db'))
|
||||||
|
user = accounts.create('test','Test','admin','offline-test-password')
|
||||||
|
other = accounts.create('other','Other','user','offline-test-password')
|
||||||
|
token, _ = accounts.issue_session(user['id'],3600,user['password_hash'])
|
||||||
|
store = ChatStore(accounts)
|
||||||
|
cid = store.create_conversation(user['id'])['id']
|
||||||
|
runs = ChatRuns(accounts,store,adapter=object())
|
||||||
|
assert await runs.cancel_operation(user['id'],cid,'uncertain-key') == {'ok':True,'run':None}
|
||||||
|
operation = await runs.send(user['id'],cid,token,content='Hi',idempotency_key='uncertain-key')
|
||||||
|
assert operation['run']['status'] == 'cancelled'
|
||||||
|
assert runs.tasks == {}
|
||||||
|
assert (await runs.send(user['id'],cid,token,content='Hi',idempotency_key='uncertain-key'))['run']['id'] == operation['run']['id']
|
||||||
|
try:
|
||||||
|
await runs.cancel_operation(other['id'],cid,'uncertain-key')
|
||||||
|
assert False, 'Foreign owner cancellation was accepted'
|
||||||
|
except KeyError:
|
||||||
|
pass
|
||||||
|
accounts.close()
|
||||||
|
asyncio.run(scenario())
|
||||||
Reference in new issue
Block a user