175 lines
12 KiB
Python
175 lines
12 KiB
Python
|
"""Local daemon + two CLI clients + mock provider. No paid model calls."""
|
||
|
|
import json, os, signal, socket, sqlite3, subprocess, tempfile, threading, time
|
||
|
|
from pathlib import Path
|
||
|
|
from http.server import ThreadingHTTPServer
|
||
|
|
from runtime_flow import Provider
|
||
|
|
from cli_flow import BINARY, tool, response
|
||
|
|
|
||
|
|
def wait(fn, seconds=15):
|
||
|
|
end=time.monotonic()+seconds
|
||
|
|
while time.monotonic()<end:
|
||
|
|
try:value=fn()
|
||
|
|
except (ConnectionRefusedError,FileNotFoundError):value=None
|
||
|
|
if value:return value
|
||
|
|
time.sleep(.05)
|
||
|
|
raise AssertionError('timed out')
|
||
|
|
|
||
|
|
def main():
|
||
|
|
with tempfile.TemporaryDirectory(prefix='gb-team-',dir='/tmp') as root, ThreadingHTTPServer(('127.0.0.1',0),Provider) as server:
|
||
|
|
root=Path(root);data=root/'data';workspace=root/'work';workspace.mkdir();server.errors=[]
|
||
|
|
def answer(text):return response({'role':'assistant','content':text})
|
||
|
|
def callback(req):
|
||
|
|
msgs=req['messages'];system=msgs[0]['content'];last=msgs[-1]
|
||
|
|
if system.startswith('Extract memory'):
|
||
|
|
payload=json.loads(msgs[-1]['content'])
|
||
|
|
expertise='research analysis' if 'teach research' in payload['conversation'] else payload['previous_expertise']
|
||
|
|
return answer(json.dumps({'expertise':expertise,'memories':[{'kind':'user_statement','content':'private example note','supersedes':''}]}))
|
||
|
|
if 'persistent GrokBoy main agent' in system:
|
||
|
|
user=next(m['content'] for m in reversed(msgs) if m['role']=='user')
|
||
|
|
if 'Background task result' in user:return answer('Background report received and assessed.')
|
||
|
|
if user=='teach research':return answer('Learned your research context.')
|
||
|
|
if user=='2':
|
||
|
|
if last['role']=='user':
|
||
|
|
pending=next(t for t in rpc('tasks','a') if t['state']=='waiting_input')
|
||
|
|
return response(tool('answer_task',{'task_id':pending['id']}))
|
||
|
|
assert json.loads(last['content']).get('forwarded'),last
|
||
|
|
return answer('Forwarded your answer to the task.')
|
||
|
|
if user=='ping':return answer('pong while background work continues')
|
||
|
|
if user.startswith('start '):
|
||
|
|
goal=user.removeprefix('start ')
|
||
|
|
if last['role']=='user':return response(tool('find_agents',{'query':'research'}))
|
||
|
|
if last['role']=='tool':
|
||
|
|
data=json.loads(last['content'])
|
||
|
|
if isinstance(data,list):return response(tool('delegate_task',{'target':data[0]['id'],'goal':goal}))
|
||
|
|
return answer('Delegated; still available to chat.')
|
||
|
|
return answer('hello')
|
||
|
|
goal=next(m['content'] for m in msgs if m['role']=='user');results=[json.loads(m['content']) for m in msgs if m['role']=='tool']
|
||
|
|
if goal=='ROOT_TASK':
|
||
|
|
if not results:return response(tool('search_memory',{'query':'private'}))
|
||
|
|
if len(results)==1:
|
||
|
|
assert results[0] and results[0][0]['content']=='private example note',results
|
||
|
|
return response(tool('spawn_agent',{'goal':'CHILD_TASK'}))
|
||
|
|
child=results[1]['id']
|
||
|
|
if len(results)==2:return response(tool('wait_task',{'task_id':child,'timeout_ms':20000}))
|
||
|
|
assert any('done' in json.dumps(x) for x in results[2:]),results
|
||
|
|
return response(tool('report_done',{'message':'Verified child artifact result.txt.'}))
|
||
|
|
if goal=='CHILD_TASK':
|
||
|
|
if not results:return response(tool('exec_command',{'cmd':'sleep 1; printf artifact > result.txt','yield_time_ms':10000}))
|
||
|
|
if len(results)==1:return response(tool('read_file',{'path':'result.txt'}))
|
||
|
|
assert results[-1]['content']=='artifact',results
|
||
|
|
return response(tool('report_done',{'message':'Read result.txt and verified artifact.'}))
|
||
|
|
if goal=='QUESTION':
|
||
|
|
if not results:return response(tool('request_user_input',{'question':'Choose a label','options':['alpha','beta']}))
|
||
|
|
assert results[-1].get('answer')=='beta' or any('beta' in m.get('content','') and 'unanswered question' in m.get('content','') for m in msgs),results
|
||
|
|
return response(tool('report_done',{'message':'selected beta'}))
|
||
|
|
if goal=='STEER':
|
||
|
|
if not results:
|
||
|
|
first=tool('exec_command',{'cmd':'echo started > steer.started; sleep 0.8','yield_time_ms':10000})
|
||
|
|
second=tool('write_file',{'path':'forbidden.txt','content':'bad'})['tool_calls'][0];second['id']='second';first['tool_calls'].append(second)
|
||
|
|
return response(first)
|
||
|
|
assert not (workspace/'forbidden.txt').exists()
|
||
|
|
assert any('do not write forbidden' in m.get('content','') for m in msgs)
|
||
|
|
return response(tool('report_done',{'message':'steering applied; no forbidden file'}))
|
||
|
|
if goal in ('BROWSER_SET','BROWSER_CHECK','BROWSER_OTHER'):
|
||
|
|
if not results:return response(tool('browser_navigate',{'url':f'http://127.0.0.1:{server.server_port}/fixture'}))
|
||
|
|
if len(results)==1:
|
||
|
|
expression="document.cookie='private_cookie=one'; localStorage.setItem('private','one'); 'set'" if goal=='BROWSER_SET' else "({cookie:document.cookie,storage:localStorage.getItem('private')})"
|
||
|
|
return response(tool('browser_eval',{'expression':expression}))
|
||
|
|
assert 'error' not in results[-1],results
|
||
|
|
if goal=='BROWSER_SET' and len(results)==2:return response(tool('browser_release',{}))
|
||
|
|
if goal=='BROWSER_CHECK':assert results[-1]['result']=={'cookie':'private_cookie=one','storage':'one'},results[-1]
|
||
|
|
if goal=='BROWSER_OTHER':assert results[-1]['result']=={'cookie':'','storage':None},results[-1]
|
||
|
|
return response(tool('report_done',{'message':'browser profile verified'}))
|
||
|
|
if goal=='LONG':
|
||
|
|
if not results:return response(tool('exec_command',{'cmd':'echo $$ > running.pid; sleep 60','yield_time_ms':10000}))
|
||
|
|
return response(tool('report_done',{'message':'should not finish'}))
|
||
|
|
raise AssertionError((goal,results))
|
||
|
|
server.callback=callback;threading.Thread(target=server.serve_forever,daemon=True).start()
|
||
|
|
env={**os.environ,'GROKBOY_DATA_DIR':str(data),'GROKBOY_API_KEY':'mock','GROKBOY_MODEL':'mock','GROKBOY_BASE_URL':f'http://127.0.0.1:{server.server_port}/v1','GROKBOY_MAX_ROUNDS_TOTAL':'48','GROKBOY_BROWSER_HEADED':'0'}
|
||
|
|
logs=[];clients=[];daemons=[]
|
||
|
|
def start():
|
||
|
|
p=subprocess.Popen([str(BINARY),'serve'],cwd=workspace,env=env,stdout=subprocess.PIPE,stderr=subprocess.PIPE,text=True)
|
||
|
|
daemons.append(p)
|
||
|
|
threading.Thread(target=lambda:[logs.append(l) for l in p.stderr],daemon=True).start()
|
||
|
|
wait(lambda:(data/'service.sock').exists() and p.poll() is None)
|
||
|
|
wait(lambda:rpc('ping'))
|
||
|
|
return p
|
||
|
|
def rpc(op,agent=None,**kw):
|
||
|
|
with socket.socket(socket.AF_UNIX) as s:
|
||
|
|
s.settimeout(5);s.connect(str(data/'service.sock'));s.sendall((json.dumps({'op':op,'agent':agent,**kw})+'\n').encode());f=s.makefile();v=json.loads(f.readline())
|
||
|
|
assert 'error' not in v,v
|
||
|
|
return v
|
||
|
|
def tasks():return rpc('tasks','a')
|
||
|
|
def latest(goal):return next((t for t in reversed(tasks()) if t['goal']==goal),None)
|
||
|
|
def terminal(goal):
|
||
|
|
t=latest(goal);return t if t and t['state']=='terminal' else None
|
||
|
|
def event_reply(agent,text):return any(text in e['payload'].get('message','') for e in rpc('events',agent,after=0)['events'] if e['kind']=='reply')
|
||
|
|
def client(name):
|
||
|
|
p=subprocess.Popen([str(BINARY),'agent','--name',name],cwd=workspace,env=env,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.PIPE,text=True)
|
||
|
|
out=[];threading.Thread(target=lambda:[out.append(l) for l in p.stdout],daemon=True).start();threading.Thread(target=lambda:[logs.append(l) for l in p.stderr],daemon=True).start();clients.append(p);return p,out
|
||
|
|
def send(p,text):p.stdin.write(text+'\n');p.stdin.flush()
|
||
|
|
daemon=None
|
||
|
|
try:
|
||
|
|
daemon=start()
|
||
|
|
for name in ['a','b']:rpc('create',name=name,cwd=str(workspace))
|
||
|
|
a,aout=client('a');b,bout=client('b')
|
||
|
|
send(b,'teach research');wait(lambda:rpc('expertise','b')['expertise']=='research analysis')
|
||
|
|
assert rpc('memory','a')==[]
|
||
|
|
assert 'private example note' not in json.dumps(rpc('agents'))
|
||
|
|
print('PASS automatic expertise extraction and private memory isolation',flush=True)
|
||
|
|
send(a,'start ROOT_TASK');wait(lambda:latest('CHILD_TASK'))
|
||
|
|
send(a,'ping');send(b,'ping')
|
||
|
|
wait(lambda:any('pong' in x for x in aout));wait(lambda:any('pong' in x for x in bout))
|
||
|
|
root_task=wait(lambda:terminal('ROOT_TASK'))
|
||
|
|
assert root_task['verdict']=='done',root_task
|
||
|
|
assert (workspace/'result.txt').read_text()=='artifact'
|
||
|
|
wait(lambda:any('Background report' in x for x in aout))
|
||
|
|
print('PASS expertise routing, nested delegation, verified artifact and two live chats',flush=True)
|
||
|
|
send(a,'start QUESTION');wait(lambda:latest('QUESTION') and latest('QUESTION')['state']=='waiting_input')
|
||
|
|
q=latest('QUESTION');send(a,'/exit');a.wait(timeout=5)
|
||
|
|
assert latest('QUESTION')['state']=='waiting_input'
|
||
|
|
send(b,'ping');wait(lambda:event_reply('b','pong'))
|
||
|
|
rpc('chat','a',message='2')
|
||
|
|
assert wait(lambda:terminal('QUESTION'))['verdict']=='done'
|
||
|
|
a,aout=client('a');wait(lambda:any('Background report' in x for x in aout))
|
||
|
|
print('PASS durable question, CLI disconnect and reconnect notification',flush=True)
|
||
|
|
send(a,'start LONG');wait(lambda:(workspace/'running.pid').exists());long=latest('LONG')
|
||
|
|
rpc('stop','a',task=long['id']);wait(lambda:terminal('LONG'))
|
||
|
|
pid=int((workspace/'running.pid').read_text())
|
||
|
|
def dead():
|
||
|
|
try:os.kill(pid,0);return False
|
||
|
|
except ProcessLookupError:return True
|
||
|
|
wait(dead)
|
||
|
|
send(b,'ping');wait(lambda:event_reply('b','pong'))
|
||
|
|
print('PASS task cancellation kills process group without stopping peer chat',flush=True)
|
||
|
|
send(a,'start STEER');wait(lambda:(workspace/'steer.started').exists())
|
||
|
|
steering=latest('STEER');rpc('say','a',task=steering['id'],message='do not write forbidden.txt')
|
||
|
|
assert wait(lambda:terminal('STEER'))['verdict']=='done'
|
||
|
|
assert not (workspace/'forbidden.txt').exists()
|
||
|
|
print('PASS task-specific steering skips unstarted actions',flush=True)
|
||
|
|
for goal in ('BROWSER_SET','BROWSER_CHECK'):
|
||
|
|
send(a,'start '+goal);result=wait(lambda:terminal(goal),30);assert result['verdict']=='done',result
|
||
|
|
send(b,'start BROWSER_OTHER');assert wait(lambda:next((t for t in rpc('tasks','b') if t['goal']=='BROWSER_OTHER' and t['state']=='terminal'),None),30)['verdict']=='done'
|
||
|
|
profiles=list((data/'profiles').glob('owner-*.browser'));assert len(profiles)==2,profiles
|
||
|
|
print('PASS real Chromium login persists across tasks; different owners remain isolated',flush=True)
|
||
|
|
send(a,'start QUESTION');wait(lambda:latest('QUESTION')['state']=='waiting_input')
|
||
|
|
q=latest('QUESTION');daemon.kill();daemon.wait(timeout=5)
|
||
|
|
daemon=start();t=latest('QUESTION');assert t['verdict']=='interrupted' and t['question'],t
|
||
|
|
rpc('resume','a',task=q['id']);wait(lambda:latest('QUESTION')['state']=='waiting_input')
|
||
|
|
rpc('say','a',task=q['id'],message='2')
|
||
|
|
assert wait(lambda:terminal('QUESTION'))['verdict']=='done'
|
||
|
|
print('PASS daemon crash, unknown-state recovery and explicit resume',flush=True)
|
||
|
|
db=sqlite3.connect(data/'team.sqlite3');
|
||
|
|
for (raw,) in db.execute('select data from tasks'):
|
||
|
|
t=json.loads(raw);assert t['requests']<=t['limit']
|
||
|
|
assert not server.errors,server.errors
|
||
|
|
except Exception:
|
||
|
|
print('daemon/client logs:\n'+''.join(logs));raise
|
||
|
|
finally:
|
||
|
|
for p in clients:
|
||
|
|
if p.poll() is None:p.terminate()
|
||
|
|
for p in daemons:
|
||
|
|
if p.poll() is None:p.send_signal(signal.SIGINT);p.wait(timeout=10)
|
||
|
|
server.shutdown()
|
||
|
|
if __name__=='__main__':main()
|