LazyBoy2/tests/team_flow.py

"""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 LazyBoy main agent' in system:
                # Delivery reminders are not new human requests.
                msgs=[m for m in msgs if not (m['role']=='user' and m.get('content','').startswith('<system_reminder>'))]
                last=msgs[-1]
                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']}))
                    if last['role']=='tool':
                        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('external_exec_command',{'cmd':'sleep 1; printf artifact > result.txt','yield_time_ms':10000}))
                if len(results)==1:return response(tool('external_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':
                answered = any(r.get('answer') == 'beta' for r in results) or any(
                    ('beta' in m.get('content','') and 'unanswered question' in m.get('content',''))
                    or 'Durable user answer to pending question: 2' in m.get('content','')
                    for m in msgs)
                if not answered:
                    # A crash/resume can rerun the model with the unanswered tool result.
                    return response(tool('request_user_input',{'question':'Choose a label','options':['alpha','beta']}))
                return response(tool('report_done',{'message':'selected beta'}))
            if goal=='STEER':
                if not results:
                    first=tool('external_exec_command',{'cmd':'echo started > steer.started; sleep 0.8','yield_time_ms':10000})
                    second=tool('external_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('external_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,'LAZYBOY_WEB_PORT':'0','LAZYBOY_DATA_DIR':str(data),'LAZYBOY_API_KEY':'mock','LAZYBOY_MODEL':'mock','LAZYBOY_BASE_URL':f'http://127.0.0.1:{server.server_port}/v1','LAZYBOY_MAX_ROUNDS_TOTAL':'48','LAZYBOY_BROWSER_HEADED':'0','LAZYBOY_BROWSER_SURFACE':'local'}
        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));print('provider errors:',server.errors);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()