diff --git a/control/chat_api.py b/control/chat_api.py index 11ff4c3..36e2785 100644 --- a/control/chat_api.py +++ b/control/chat_api.py @@ -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)} +@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(): return router diff --git a/control/chat_runs.py b/control/chat_runs.py index a459cf1..bdc340d 100644 --- a/control/chat_runs.py +++ b/control/chat_runs.py @@ -108,6 +108,10 @@ class ChatRuns: 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, 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): 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() 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): with self.accounts.lock: if metadata.get('room'): @@ -176,6 +202,10 @@ class ChatRuns: raise HTTPException(429, 'Chat queue full') operation = await self.call('send', owner, cid, content, idempotency_key, attachment_ids, mode=mode) 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) if metadata is None: settings = await asyncio.to_thread(self.accounts.get_model_settings, owner) diff --git a/tests/unit/test_chat_runs_limits.py b/tests/unit/test_chat_runs_limits.py index bb55bb7..201a851 100644 --- a/tests/unit/test_chat_runs_limits.py +++ b/tests/unit/test_chat_runs_limits.py @@ -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' accounts.close() 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())