fix: preserve idempotent admission under capacity and preparation failures
This commit is contained in:
1 parent
33929b8c09
commit
e059660194
2 files changed
+45
-1
No files matched your search
+38
-1
@@ -161,10 +161,18 @@ class ChatRuns:
|
|||||||
parts = await self.feature_store.provider_content(owner, attachment_ids, vision=vision)
|
parts = await self.feature_store.provider_content(owner, attachment_ids, vision=vision)
|
||||||
if not self.capabilities()['text_chat']:
|
if not self.capabilities()['text_chat']:
|
||||||
raise HTTPException(503, 'Chat unavailable')
|
raise HTTPException(503, 'Chat unavailable')
|
||||||
|
operation = None
|
||||||
try:
|
try:
|
||||||
# Coordinate payload admission with execution and recovery metadata.
|
# Coordinate payload admission with execution and recovery metadata.
|
||||||
async with self.admission_locks[owner]:
|
async with self.admission_locks[owner]:
|
||||||
if sum(not task.done() for task in self.tasks.values()) >= 256 or sum(not task.done() and self.task_owners.get(rid) == owner for rid, task in self.tasks.items()) >= 8:
|
def prior_exists():
|
||||||
|
with self.accounts.lock:
|
||||||
|
return self.accounts.db.execute('SELECT 1 FROM chat_idempotency WHERE owner=? AND conversation_id=? AND key=?', (owner, cid, idempotency_key)).fetchone() is not None
|
||||||
|
for done_id in [key for key, task in self.tasks.items() if task.done()]:
|
||||||
|
self.tasks.pop(done_id, None)
|
||||||
|
self.task_owners.pop(done_id, None)
|
||||||
|
if not await asyncio.to_thread(prior_exists):
|
||||||
|
if len(self.tasks) >= 256 or sum(self.task_owners.get(rid) == owner for rid in self.tasks) >= 8:
|
||||||
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']
|
||||||
@@ -202,12 +210,35 @@ class ChatRuns:
|
|||||||
if operation['run']['status'] not in TERMINAL:
|
if operation['run']['status'] not in TERMINAL:
|
||||||
self.schedule(owner, rid)
|
self.schedule(owner, rid)
|
||||||
return operation
|
return operation
|
||||||
|
except HTTPException:
|
||||||
|
if operation and await asyncio.to_thread(self.metadata, owner, operation['run']['id']) is None:
|
||||||
|
rid = operation['run']['id']
|
||||||
|
run = await self.set_status(owner, rid, 'failed', 'Chat admission could not be prepared. Start a new request.')
|
||||||
|
await self.emit(owner, rid, 'done', {'run': run})
|
||||||
|
raise
|
||||||
except ChatConflict:
|
except ChatConflict:
|
||||||
|
if operation and await asyncio.to_thread(self.metadata, owner, operation['run']['id']) is None:
|
||||||
|
run = await self.set_status(owner, operation['run']['id'], 'failed', 'Chat admission could not be prepared. Start a new request.')
|
||||||
|
await self.emit(owner, operation['run']['id'], 'done', {'run': run})
|
||||||
raise HTTPException(409, 'Idempotency or operation conflict') from None
|
raise HTTPException(409, 'Idempotency or operation conflict') from None
|
||||||
except KeyError:
|
except KeyError:
|
||||||
|
if operation and await asyncio.to_thread(self.metadata, owner, operation['run']['id']) is None:
|
||||||
|
try:
|
||||||
|
run = await self.set_status(owner, operation['run']['id'], 'failed', 'Chat admission could not be prepared. Start a new request.')
|
||||||
|
await self.emit(owner, operation['run']['id'], 'done', {'run': run})
|
||||||
|
except KeyError:
|
||||||
|
pass # Owner/conversation was concurrently removed.
|
||||||
raise HTTPException(404, 'Chat resource not found') from None
|
raise HTTPException(404, 'Chat resource not found') from None
|
||||||
except ValueError:
|
except ValueError:
|
||||||
|
if operation and await asyncio.to_thread(self.metadata, owner, operation['run']['id']) is None:
|
||||||
|
run = await self.set_status(owner, operation['run']['id'], 'failed', 'Chat admission could not be prepared. Start a new request.')
|
||||||
|
await self.emit(owner, operation['run']['id'], 'done', {'run': run})
|
||||||
raise HTTPException(400, 'Invalid chat request') from None
|
raise HTTPException(400, 'Invalid chat request') from None
|
||||||
|
except Exception:
|
||||||
|
if operation and await asyncio.to_thread(self.metadata, owner, operation['run']['id']) is None:
|
||||||
|
run = await self.set_status(owner, operation['run']['id'], 'failed', 'Chat admission could not be prepared. Start a new request.')
|
||||||
|
await self.emit(owner, operation['run']['id'], 'done', {'run': run})
|
||||||
|
raise HTTPException(503, 'Chat admission unavailable') from None
|
||||||
|
|
||||||
def schedule(self, owner, rid):
|
def schedule(self, owner, rid):
|
||||||
if not self.closed and (rid not in self.tasks or self.tasks[rid].done()):
|
if not self.closed and (rid not in self.tasks or self.tasks[rid].done()):
|
||||||
@@ -302,6 +333,12 @@ class ChatRuns:
|
|||||||
await asyncio.to_thread(self.metadata, owner, rid, metadata)
|
await asyncio.to_thread(self.metadata, owner, rid, metadata)
|
||||||
watcher = asyncio.create_task(self.monitor_authorization(owner, rid, metadata))
|
watcher = asyncio.create_task(self.monitor_authorization(owner, rid, metadata))
|
||||||
if not metadata['upstream_id']:
|
if not metadata['upstream_id']:
|
||||||
|
if not metadata.get('submission_attempted'):
|
||||||
|
session_state = await self.call('get_session', owner, metadata['conversation_id'])
|
||||||
|
if session_state['hermes_seeded']:
|
||||||
|
metadata['payload'].pop('conversation_history', None)
|
||||||
|
metadata['submission_attempted'] = True
|
||||||
|
await asyncio.to_thread(self.metadata, owner, rid, metadata)
|
||||||
while True:
|
while True:
|
||||||
try:
|
try:
|
||||||
started = await self.adapter.start(owner, rid, metadata['payload'])
|
started = await self.adapter.start(owner, rid, metadata['payload'])
|
||||||
|
|||||||
@@ -90,6 +90,13 @@ def test_idempotent_send_final_output_and_redacted_replay(env):
|
|||||||
assert 'event: done' in stream.text
|
assert 'event: done' in stream.text
|
||||||
assert 'hidden-secret' not in stream.text and 'tool-secret' not in stream.text
|
assert 'hidden-secret' not in stream.text and 'tool-secret' not in stream.text
|
||||||
assert 'secret-never-public' not in stream.text
|
assert 'secret-never-public' not in stream.text
|
||||||
|
# Capacity limits never prevent retrieval of an already accepted operation.
|
||||||
|
async def fill():
|
||||||
|
app.state.chat_runs.tasks.update({str(i): asyncio.create_task(asyncio.sleep(100)) for i in range(256)})
|
||||||
|
client.portal.call(fill)
|
||||||
|
replay = client.post(url, headers=headers, json=value)
|
||||||
|
assert replay.status_code == 202
|
||||||
|
assert replay.json()['run']['id'] == rid
|
||||||
|
|
||||||
def test_body_validation_and_honest_capabilities(env):
|
def test_body_validation_and_honest_capabilities(env):
|
||||||
client, app, adapter, headers = env
|
client, app, adapter, headers = env
|
||||||
|
|||||||
Reference in new issue
Block a user