380 lines
21 KiB
Python
380 lines
21 KiB
Python
"""Behavioral regression for the independent agent coordination service."""
|
|
import concurrent.futures
|
|
import http.client
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
from agent_control.store import Store, Invalid, Conflict
|
|
from agent_control.server import Server, restricted_token
|
|
|
|
|
|
class StoreTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.path = Path(self.tmp.name) / 'db.sqlite3'
|
|
self.store = Store(self.path)
|
|
for sid in ('a','b'):
|
|
self.store.session(dict(id=sid,name=sid,wp='WP1'))
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def send(self, sid='a', key='one'):
|
|
return self.store.send(dict(session_id=sid,text='question',idempotency_key=key))
|
|
|
|
def claim(self, sid='a', worker='w'):
|
|
return self.store.claim(dict(session_id=sid,worker_id=worker))['message']
|
|
|
|
def completion(self, message, **overrides):
|
|
data=dict(session_id=message['session_id'],worker_id='w',message_id=message['id'],lease_id=message['lease_id'],text='real output',thread_id='thread-a')
|
|
data.update(overrides)
|
|
return data
|
|
|
|
def test_persistence_and_idempotency(self):
|
|
one = self.send()
|
|
self.assertEqual(one['id'], self.send()['id'])
|
|
reopened = Store(self.path)
|
|
self.assertEqual(reopened.messages('a')[0]['id'], one['id'])
|
|
with self.assertRaises(Conflict):
|
|
reopened.send(dict(session_id='a',text='different',idempotency_key='one'))
|
|
self.assertEqual(len(reopened.messages('a')),1)
|
|
|
|
def test_late_history_final_reconciles_without_duplicate_reply(self):
|
|
self.store.session(dict(id='a',name='a',thread_id='thread-a'))
|
|
original=self.send()
|
|
lease=self.claim()
|
|
owner=dict(session_id='a',worker_id='w',message_id=original['id'],lease_id=lease['lease_id'])
|
|
self.store.start(owner)
|
|
proof='exact submitted input'
|
|
self.store.submission(dict(owner,state='submitting',thread_id='thread-a',input_sha256=hashlib.sha256(proof.encode()).hexdigest()))
|
|
self.store.uncertain(owner)
|
|
batch=dict(session_id='a',thread_id='thread-a',items=[
|
|
dict(id='external-user',source='human',text='question',input_text=proof,related_message_id=original['id'],phase='input',turn_id='turn-a',turn_status='completed'),
|
|
dict(id='external-final',source='agent',text='real output',input_text=proof,related_message_id=original['id'],phase='final_answer',turn_id='turn-a',turn_status='completed')])
|
|
self.store.mirror(batch)
|
|
self.store.mirror(batch)
|
|
self.assertEqual(len(self.store.messages('a')),2)
|
|
self.assertEqual(self.store.messages('a')[0]['status'],'completed')
|
|
self.store.complete(self.completion(lease))
|
|
self.assertEqual(len(self.store.messages('a')),2)
|
|
with self.store.db() as db:
|
|
link=json.loads(db.execute('SELECT data FROM history_links WHERE external_id=?',('external-final',)).fetchone()['data'])
|
|
self.assertEqual(link['turn_id'],'turn-a')
|
|
|
|
def test_submission_is_bound_persistent_and_fenced(self):
|
|
self.store.session(dict(id='a',name='a',thread_id='thread-a'))
|
|
self.send()
|
|
lease=self.claim()
|
|
owner=self.completion(lease)
|
|
digest=hashlib.sha256(b'exact input').hexdigest()
|
|
data=dict(owner,state='submitting',input_sha256=digest)
|
|
with self.assertRaises(Conflict): self.store.submission(data)
|
|
self.store.start(owner)
|
|
self.store.submission(data)
|
|
self.store.submission(dict(data,state='accepted',receipt_id='receipt-a'))
|
|
self.store.submission(dict(data,state='observed',turn_id='turn-a'))
|
|
value=Store(self.path).messages('a')[0]['submission']
|
|
self.assertEqual(value['receipt_id'],'receipt-a')
|
|
self.assertEqual(value['turn_id'],'turn-a')
|
|
self.assertEqual(self.store.submission(data)['submission']['state'],'observed')
|
|
for change in ({'worker_id':'other'},{'thread_id':'wrong'},{'input_sha256':'b'*64},
|
|
{'state':'accepted','receipt_id':'other'},{'state':'observed','turn_id':'other'}):
|
|
with self.assertRaises(Conflict): self.store.submission(dict(data,**change))
|
|
with patch('agent_control.store.time.time',return_value=lease['lease_until']+1):
|
|
with self.assertRaises(Conflict): self.store.submission(data)
|
|
|
|
def test_mirror_requires_registered_full_input_and_turn(self):
|
|
self.store.session(dict(id='a',name='a',thread_id='thread-a'))
|
|
original=self.send(); lease=self.claim(); owner=self.completion(lease)
|
|
self.store.start(owner)
|
|
base=dict(session_id='a',thread_id='thread-a')
|
|
item=dict(id='f',source='agent',text='result',related_message_id=original['id'],
|
|
input_text='exact input',phase='final_answer',turn_id='turn-a',turn_status='completed')
|
|
with self.assertRaises(Conflict): self.store.mirror(dict(base,items=[item]))
|
|
self.store.submission(dict(owner,state='observed',turn_id='turn-a',input_sha256=hashlib.sha256(b'exact input').hexdigest()))
|
|
for change in ({'input_text':'forged exact input'},{'turn_id':'other-turn'}):
|
|
with self.assertRaises(Conflict): self.store.mirror(dict(base,items=[dict(item,**change)]))
|
|
self.assertIsNone(self.store.messages('a')[0]['response_id'])
|
|
|
|
def test_long_input_final_deferred_then_completed_without_duplicates(self):
|
|
self.store.session(dict(id='a',name='a',thread_id='thread-a'))
|
|
original=self.store.send(dict(session_id='a',text='q'*8001,idempotency_key='long'))
|
|
lease=self.claim(); owner=self.completion(lease); self.store.start(owner)
|
|
proof='envelope\n'+original['text']; answer='r'*8001
|
|
self.store.submission(dict(owner,state='accepted',receipt_id='r',input_sha256=hashlib.sha256(proof.encode()).hexdigest()))
|
|
items=[dict(id='u',source='human',text=original['text'],phase='input'),
|
|
dict(id='f',source='agent',text=answer,phase='final_answer')]
|
|
for i in items: i.update(related_message_id=original['id'],input_text=proof,turn_id='t',turn_status='completed')
|
|
batch=dict(session_id='a',thread_id='thread-a',items=items)
|
|
self.assertEqual(self.store.mirror(batch)['deferred'],['f'])
|
|
self.assertEqual(len(self.store.messages('a')),1)
|
|
self.store.complete(dict(owner,text=answer))
|
|
self.assertEqual(self.store.mirror(batch)['deferred'],[])
|
|
self.assertEqual(len(self.store.messages('a')),2)
|
|
self.assertEqual(self.store.messages('a')[1]['text'],answer)
|
|
|
|
def test_long_late_final_recovers_uncertain_and_rejects_oversize(self):
|
|
self.store.session(dict(id='a',name='a',thread_id='thread-a'))
|
|
original=self.send(); lease=self.claim(); owner=self.completion(lease); self.store.start(owner)
|
|
self.store.submission(dict(owner,state='submitting',input_sha256=hashlib.sha256(b'proof').hexdigest()))
|
|
self.store.uncertain(owner)
|
|
item=dict(id='f',source='agent',text='r'*8001,phase='final_answer',related_message_id=original['id'],
|
|
input_text='proof',turn_id='late',turn_status='completed')
|
|
batch=dict(session_id='a',thread_id='thread-a',items=[item])
|
|
self.store.mirror(batch); self.store.mirror(batch)
|
|
self.assertEqual(self.store.messages('a')[0]['status'],'completed')
|
|
self.assertEqual(len(self.store.messages('a')),2)
|
|
self.store.complete(dict(owner,text='r'*8001))
|
|
with self.assertRaises(Invalid): self.store.mirror(dict(batch,items=[dict(item,text='x'*262145)]))
|
|
with self.assertRaises(Invalid): self.store.complete(dict(owner,text='x'*262145))
|
|
self.assertEqual(len(self.store.messages('a')),2)
|
|
|
|
def test_live_visible_commentary_updates_same_item(self):
|
|
self.store.session(dict(id='a',name='a',thread_id='thread-a'))
|
|
payload=dict(session_id='a',thread_id='thread-a',items=[dict(id='comment',source='agent',text='first',phase='commentary',allow_update=True)])
|
|
self.store.mirror(payload)
|
|
payload['items'][0]['text']='updated'
|
|
self.store.mirror(payload)
|
|
self.assertEqual(len(self.store.messages('a')),1)
|
|
self.assertEqual(self.store.messages('a')[0]['text'],'updated')
|
|
|
|
def test_claim_is_atomic_and_serial_per_session(self):
|
|
self.send()
|
|
self.send(key='two')
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
|
|
results=list(pool.map(lambda n:self.claim(worker='w'+str(n)),range(8)))
|
|
self.assertEqual(sum(value is not None for value in results),1)
|
|
self.assertIsNone(self.claim('b'))
|
|
|
|
def test_cross_session_replies_and_completion_rejected(self):
|
|
one=self.send()
|
|
with self.assertRaises(Invalid):
|
|
self.store.send(dict(session_id='b',text='bad',source='agent',reply_to=one['id'],idempotency_key='reply'))
|
|
claimed=self.claim()
|
|
with self.assertRaises(Conflict):
|
|
self.store.complete(self.completion(claimed,session_id='b'))
|
|
self.assertEqual(self.store.messages('b'),[])
|
|
self.assertEqual(self.store.messages('a')[0]['status'],'claimed')
|
|
|
|
def test_expired_lease_fencing_and_renew(self):
|
|
self.send()
|
|
with patch('agent_control.store.time.time', return_value=100):
|
|
first=self.claim()
|
|
with patch('agent_control.store.time.time', return_value=200):
|
|
result=self.store.renew(dict(self.completion(first),lease_seconds=300))
|
|
self.assertEqual(result['lease_until'],500)
|
|
with patch('agent_control.store.time.time', return_value=501):
|
|
with self.assertRaises(Conflict):
|
|
self.store.renew(self.completion(first))
|
|
second=self.claim()
|
|
self.assertNotEqual(first['lease_id'],second['lease_id'])
|
|
with self.assertRaises(Conflict):
|
|
self.store.complete(self.completion(first))
|
|
self.store.complete(self.completion(second))
|
|
self.assertEqual(len(self.store.messages('a')),2)
|
|
|
|
def test_completion_is_atomic_idempotent_and_failure_terminal(self):
|
|
self.send()
|
|
claim=self.claim()
|
|
response=self.store.complete(self.completion(claim))
|
|
self.assertEqual(response['id'],self.store.complete(self.completion(claim))['id'])
|
|
self.assertEqual(len(Store(self.path).messages('a')),2)
|
|
self.assertIsNone(self.claim())
|
|
self.send(key='two')
|
|
next_claim=self.claim()
|
|
self.store.complete(self.completion(next_claim,outcome='failed',thread_id=None,text='authentication unavailable'))
|
|
self.assertEqual(self.store.messages('a')[-2]['status'],'failed')
|
|
self.assertIsNone(self.claim())
|
|
|
|
def test_task_validation_is_atomic(self):
|
|
self.store.task(dict(id='WP0',name='总控',status='running',phase='planning',eta_min_minutes=2,eta_max_minutes=5,eta_basis='remaining checks'))
|
|
before=self.store.state()['tasks']
|
|
for update in ({'phase':'imaginary'},{'eta_max_minutes':1},{'eta_min_minutes':float('nan')},{'eta_basis':''},{'status':'completed'}):
|
|
with self.assertRaises(Invalid):
|
|
self.store.task(update,'WP0')
|
|
self.assertEqual(before,self.store.state()['tasks'])
|
|
self.store.task(dict(status='completed',phase='done'),'WP0')
|
|
|
|
def test_started_expiry_blocks_session_until_reconciliation(self):
|
|
self.send()
|
|
self.send(key='next')
|
|
with patch('agent_control.store.time.time', return_value=100):
|
|
first=self.claim()
|
|
self.store.start(self.completion(first))
|
|
with patch('agent_control.store.time.time', return_value=401):
|
|
result=self.store.claim(dict(session_id='a',worker_id='other'))
|
|
self.assertIsNone(result['message'])
|
|
self.assertTrue(result['needs_reconciliation'])
|
|
with self.assertRaises(Conflict):
|
|
self.store.complete(self.completion(first))
|
|
self.assertEqual(self.store.messages('a')[0]['outcome'],'uncertain')
|
|
self.assertNotIn('lease_id',self.store.messages('a')[0])
|
|
self.store.reconcile(dict(session_id='a',message_id=first['id'],action='failed',resolution='人工确认已产生副作用，禁止重跑'))
|
|
second=self.claim()
|
|
self.assertNotEqual(first['id'],second['id'])
|
|
|
|
def test_uncertain_signal_fences_worker_and_requires_explicit_retry(self):
|
|
self.send()
|
|
first=self.claim()
|
|
self.store.start(self.completion(first))
|
|
with self.assertRaises(Conflict):
|
|
self.store.uncertain(self.completion(first,worker_id='other'))
|
|
self.store.uncertain(self.completion(first))
|
|
self.assertIsNone(self.claim())
|
|
with self.assertRaises(Invalid):
|
|
self.store.reconcile(dict(session_id='a',message_id=first['id'],action='retry',resolution=''))
|
|
self.store.reconcile(dict(session_id='a',message_id=first['id'],action='retry',resolution='人工确认未产生副作用，允许重试'))
|
|
second=self.claim()
|
|
self.assertEqual(first['id'],second['id'])
|
|
self.assertNotEqual(first['lease_id'],second['lease_id'])
|
|
with self.assertRaises(Conflict):
|
|
self.store.complete(self.completion(first))
|
|
|
|
def test_mirror_is_bound_idempotent_and_never_executable(self):
|
|
self.store.session(dict(id='a',name='bound',thread_id='actual-thread'))
|
|
batch=dict(session_id='a',thread_id='actual-thread',items=[
|
|
dict(id='external-1',source='human',text='historical question',created=10),
|
|
dict(id='external-1',source='agent',text='historical reply',created=11)])
|
|
self.assertEqual(self.store.mirror(batch)['imported'],2)
|
|
self.assertEqual(self.store.mirror(batch)['imported'],0)
|
|
self.assertEqual(self.store.state()['sessions'][0]['connection'],'waiting')
|
|
self.assertIsNone(self.claim())
|
|
rows=self.store.messages('a')
|
|
self.assertTrue(all(r['status']=='delivered' and r['origin']=='session_history' for r in rows))
|
|
self.assertEqual(rows[0]['created'],10)
|
|
self.assertEqual(self.store.messages('b'),[])
|
|
with self.assertRaises(Conflict):
|
|
self.store.mirror(dict(batch,thread_id='other-thread'))
|
|
with self.assertRaises(Conflict):
|
|
self.store.mirror(dict(batch,session_id='b'))
|
|
|
|
def test_mirror_batch_failure_rolls_back_and_rejects_changed_items(self):
|
|
self.store.session(dict(id='a',name='bound',thread_id='actual-thread'))
|
|
base=dict(session_id='a',thread_id='actual-thread')
|
|
one=dict(id='one',source='human',text='first')
|
|
self.store.mirror(dict(base,items=[one]))
|
|
with self.assertRaises(Conflict):
|
|
self.store.mirror(dict(base,items=[dict(one,id='new'),dict(one,text='changed')]))
|
|
self.assertEqual(len(self.store.messages('a')),1)
|
|
for items in ([dict(one,created=float('nan'))],[one]*101,[dict(one,source='tool')]):
|
|
with self.assertRaises(Invalid):
|
|
self.store.mirror(dict(base,items=items))
|
|
self.assertEqual(len(self.store.messages('a')),1)
|
|
|
|
def test_requested_is_not_real_session(self):
|
|
self.store.request(dict(name='new',prompt='work'))
|
|
state=self.store.state()
|
|
self.assertEqual(len(state['sessions']),2)
|
|
self.assertEqual(state['requests'][0]['status'],'requested')
|
|
self.assertTrue(all(s['connection']=='waiting' for s in state['sessions']))
|
|
|
|
|
|
class HTTPTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp=tempfile.TemporaryDirectory()
|
|
root=Path(self.tmp.name)
|
|
root.chmod(0o700)
|
|
self.secret='a'*48
|
|
self.token=root/'token'
|
|
self.token.write_text(self.secret)
|
|
self.token.chmod(0o600)
|
|
web=root/'web'
|
|
web.mkdir()
|
|
(web/'index.html').write_text('<title>Agent Workbench</title>')
|
|
(web/'app.js').write_text('console.log("test fixture")')
|
|
self.server=Server(dict(bind='127.0.0.1',port=0,runtime_root=str(root),token_file=str(self.token),web_root=str(web)))
|
|
self.thread=threading.Thread(target=self.server.serve_forever,daemon=True)
|
|
self.thread.start()
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.thread.join()
|
|
self.tmp.cleanup()
|
|
|
|
def request(self,path='/api/state',method='GET',data=None,headers=None):
|
|
conn=http.client.HTTPConnection(*self.server.server_address,timeout=3)
|
|
base={'Authorization':'Bearer '+self.secret,'X-CSRF-Token':'1','Content-Type':'application/json'}
|
|
base.update(headers or {})
|
|
conn.request(method,path,body=json.dumps(data) if data is not None else None,headers=base)
|
|
response=conn.getresponse()
|
|
body=response.read()
|
|
result=(response.status,body,response.getheaders())
|
|
conn.close()
|
|
return result
|
|
|
|
def test_auth_host_origin_csrf_and_static_boundary(self):
|
|
self.assertEqual(self.request()[0],200)
|
|
self.assertEqual(self.request(headers={'Authorization':''})[0],401)
|
|
self.assertEqual(self.request(headers={'Host':'evil.example'})[0],403)
|
|
self.assertEqual(self.request(headers={'Origin':'https://evil.example'})[0],403)
|
|
self.assertEqual(self.request('/api/sessions','POST',{'name':'test'}, {'X-CSRF-Token':''})[0],403)
|
|
self.assertEqual(self.request('/../store.py')[0],404)
|
|
self.assertEqual(self.request('/%2e%2e/server.py')[0],404)
|
|
self.assertEqual(self.request('/?token=secret')[0],404)
|
|
code,body,headers=self.request('/',headers={'Authorization':''})
|
|
self.assertEqual(code,200)
|
|
self.assertIn(b'Agent Workbench',body)
|
|
self.assertIn('Content-Security-Policy',dict(headers))
|
|
|
|
def test_http_roundtrip_and_errors_do_not_leak(self):
|
|
code,body,_=self.request('/api/sessions','POST',{'id':'a','name':'session'})
|
|
self.assertEqual(code,200)
|
|
code,body,_=self.request('/api/messages','POST',{'session_id':'a','text':'hello','idempotency_key':'k'})
|
|
self.assertEqual(code,200)
|
|
code,body,_=self.request('/api/worker/claim','POST',{'session_id':'a','worker_id':'w'})
|
|
message=json.loads(body)['message']
|
|
self.assertEqual(code,200)
|
|
code,body,_=self.request('/api/worker/complete','POST',dict(session_id='a',worker_id='w',message_id=message['id'],lease_id=message['lease_id'],thread_id='real-thread',text='output'))
|
|
self.assertEqual(code,200)
|
|
code,body,_=self.request('/api/messages?session_id=a')
|
|
self.assertEqual(len(json.loads(body)['messages']),2)
|
|
with patch.object(self.server.store,'state',side_effect=RuntimeError(self.secret)):
|
|
code,body,_=self.request()
|
|
self.assertEqual(code,500)
|
|
self.assertNotIn(self.secret.encode(),body)
|
|
|
|
def test_http_mirror_accepts_bounded_long_history(self):
|
|
self.request('/api/sessions','POST',dict(id='history',name='history',thread_id='bound-thread'))
|
|
batch=dict(session_id='history',thread_id='bound-thread',items=[dict(id='long',source='agent',text='历史正文'*6000)])
|
|
self.assertGreater(len(json.dumps(batch).encode()),65536)
|
|
code,body,_=self.request('/api/mirror','POST',batch)
|
|
self.assertEqual(code,200)
|
|
self.assertEqual(json.loads(body)['imported'],1)
|
|
code,body,_=self.request('/api/worker/claim','POST',dict(session_id='history',worker_id='w'))
|
|
self.assertIsNone(json.loads(body)['message'])
|
|
|
|
def test_http_submission_and_large_completion(self):
|
|
self.request('/api/sessions','POST',dict(id='large',name='large',thread_id='thread-large'))
|
|
self.request('/api/messages','POST',dict(session_id='large',text='q',idempotency_key='one'))
|
|
_,body,_=self.request('/api/worker/claim','POST',dict(session_id='large',worker_id='w'))
|
|
claimed=json.loads(body)['message']
|
|
owner=dict(session_id='large',worker_id='w',message_id=claimed['id'],lease_id=claimed['lease_id'])
|
|
self.assertEqual(self.request('/api/worker/start','POST',owner)[0],200)
|
|
payload=dict(owner,state='submitting',thread_id='thread-large',input_sha256=hashlib.sha256(b'input').hexdigest())
|
|
self.assertEqual(self.request('/api/worker/submission','POST',payload)[0],200)
|
|
text='result'*20000
|
|
result=dict(owner,thread_id='thread-large',text=text)
|
|
self.assertGreater(len(json.dumps(result)),65536)
|
|
self.assertEqual(self.request('/api/worker/complete','POST',result)[0],200)
|
|
self.assertEqual(self.request('/api/worker/complete','POST',dict(result,text='x'*262145))[0],400)
|
|
|
|
@unittest.skipUnless(os.name=='posix','POSIX permissions')
|
|
def test_permissive_secret_file_rejected(self):
|
|
self.token.chmod(0o644)
|
|
with self.assertRaises(ValueError):
|
|
restricted_token(self.token)
|
|
|
|
|
|
if __name__=='__main__':
|
|
unittest.main()
|