LazyBoy2/tests/research_flow.py

"""Offline daemon regression: bounded research, direct delivery and timing. No paid calls."""
import json
import os
import socket
import subprocess
import tempfile
import threading
from pathlib import Path
from http.server import ThreadingHTTPServer
from cli_flow import BINARY, tool, response
from runtime_flow import Provider
from team_flow import wait

class ResearchProvider(Provider):
    def do_POST(self):
        if "/aiserver.v1.AiService/" in self.path:
            args = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
            self.server.searches.append(args["searchTerm"])
            body = json.dumps({"answer":"Offline fixture evidence",
                "documents":[{"url":"https://example.com/official","title":"Official fixture"}]}).encode()
            self.send_response(200)
            self.send_header("Content-Type","application/json")
            self.send_header("Content-Length",str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        else:
            super().do_POST()

def main():
    with tempfile.TemporaryDirectory(prefix="research-flow-") as root, ThreadingHTTPServer(("127.0.0.1",0),ResearchProvider) as server:
        root=Path(root)
        server.errors=[]
        server.searches=[]
        def callback(req):
            messages=req["messages"]
            system=messages[0]["content"]
            if system.startswith("Extract memory"):
                return response({"role":"assistant","content":json.dumps({"expertise":"","memories":[]})})
            if "persistent LazyBoy main agent" in system:
                calls=[c for m in messages for c in m.get("tool_calls",[])]
                if not calls:
                    return response(tool("spawn_agent",{"goal":"研究蝦皮分潤：離線測試攻略","task_type":"research"}))
                return response({"role":"assistant","content":"已開始研究。"})
            results=[json.loads(m["content"]) for m in messages if m["role"]=="tool"]
            n=len(results)
            if n==0: return response(tool("send_message",{"content":"先查來源。"}))
            if n==1: return response(tool("web_search",{"searchTerm":"official"}))
            if n==2: return response(tool("web_search",{"searchTerm":"official"}))
            if n==3:
                assert results[-1]["cached"]
                return response(tool("web_search",{"searchTerm":"practice"}))
            if n==4: return response(tool("web_search",{"searchTerm":"must be blocked"}))
            if n==5:
                assert "closed" in results[-1]["error"]
                return response(tool("publish_research",{"summary":"離線測試：先小規模驗證。",
                    "steps":["選定受眾","測試內容"],"sources":["https://example.com/official"],
                    "unknowns":["實際轉換率待測"],"gaps":["計算方式"]}))
            if n==6:
                assert results[-1]["research_complete"] is False
                return response(tool("web_search",{"searchTerm":"gap evidence","gap":0}))
            if n==7:
                return response(tool("publish_research",{"summary":"離線測試補充：計算方式仍待官方確認。",
                    "steps":[],"sources":["https://example.com/official"],
                    "unknowns":["計算方式"],"gaps":[]}))
            raise AssertionError("unexpected extra model round")
        server.callback=callback
        threading.Thread(target=server.serve_forever,daemon=True).start()
        env={**os.environ,"LAZYBOY_DATA_DIR":str(root/"data"),"LAZYBOY_API_KEY":"offline",
            "LAZYBOY_MODEL":"mock","LAZYBOY_BASE_URL":f"http://127.0.0.1:{server.server_port}/v1",
            "LAZYBOY_WEB_BACKEND_URL":f"http://127.0.0.1:{server.server_port}",
            "LAZYBOY_WEB_PORT":"0"}
        logs=[]
        process=subprocess.Popen([str(BINARY),"serve"],cwd=root,env=env,
            stdout=subprocess.DEVNULL,stderr=subprocess.PIPE,text=True)
        threading.Thread(target=lambda:logs.extend(process.stderr),daemon=True).start()
        def rpc(op,**kwargs):
            with socket.socket(socket.AF_UNIX) as sock:
                sock.settimeout(5)
                sock.connect(str(root/"data"/"service.sock"))
                sock.sendall((json.dumps({"op":op,**kwargs})+"\n").encode())
                result=json.loads(sock.makefile().readline())
                assert "error" not in result,result
                return result
        try:
            wait(lambda:(root/"data"/"service.sock").exists())
            rpc("create",name="owner",cwd=str(root))
            rpc("chat",agent="owner",message="我要你研究蝦皮分潤，整理攻略")
            def finished():
                tasks=rpc("tasks",agent="owner")
                return tasks[0] if tasks and tasks[0]["state"]=="terminal" else None
            task=wait(finished,20)
            assert task["verdict"]=="answer",task
            assert server.searches==["official","practice","gap evidence"],server.searches
            history=rpc("get",agent="owner")["transcript"]
            text=[row["content"] for row in history]
            assert sum("第一版攻略" in t for t in text)==1,text
            assert sum("補充結果" in t for t in text)==1,text
            assert rpc("get",agent="owner")["transcript"]==history
            events=rpc("events",agent=task["agent_id"],after=0)["events"]
            stages={e["payload"].get("stage") for e in events if e["kind"]=="timing"}
            assert {"worker_model_queue","model_response","research_first_delivery","research_final_delivery","tool_execution"} <= stages,stages
            assert not server.errors,server.errors
            print("PASS bounded research, cache reuse, first/final publication, durable transcript and correlated timing")
        finally:
            process.terminate()
            try: process.wait(timeout=5)
            except subprocess.TimeoutExpired: process.kill();process.wait()
            server.shutdown()
            if process.returncode not in (0,-15):
                print("".join(logs)[-2000:])

if __name__=="__main__":
    main()