# Logging pass-through proxy to api.x.ai: records token usage of every request (incl. failed ones).
import os,json,time,aiohttp
from aiohttp import web
UP="https://api.x.ai"; KEY=os.environ["XAI_API_KEY"]; LOG="/workspace/tests/usage.jsonl"
async def h(req):
    path=req.path_qs; tag=""
    if path.startswith("/t/"): _,_,tag,rest=path.split("/",3); path="/"+rest
    body=await req.read(); j=None
    try:
        j=json.loads(body)
        if j.get("stream"): j.setdefault("stream_options",{})["include_usage"]=True
        for k in ("frequency_penalty","presence_penalty"):  # rejected by some xAI models
            if j.get(k) in (0,0.0): j.pop(k)
        body=json.dumps(j).encode()
    except Exception: pass
    hd={"Authorization":f"Bearer {KEY}","Content-Type":"application/json"}
    async with aiohttp.ClientSession() as s:
        async with s.request(req.method, UP+path, data=body if body else None, headers=hd) as r:
            resp=web.StreamResponse(status=r.status, headers={"Content-Type":r.headers.get("Content-Type","application/json")})
            await resp.prepare(req); buf=b""
            async for chunk in r.content.iter_any():
                buf+=chunk; await resp.write(chunk)
            await resp.write_eof()
    usage=None; txt=buf.decode("utf8","replace")
    try: usage=json.loads(txt).get("usage")
    except Exception:
        for line in txt.splitlines():
            if line.startswith("data:") and '"usage"' in line:
                try: u=json.loads(line[5:]).get("usage"); usage=u or usage
                except Exception: pass
    rec={"ts":time.time(),"tag":tag,"path":path,"status":r.status,"model":(j or {}).get("model"),"usage":usage}
    if r.status>=400: rec["error"]=txt[:500]
    open(LOG,"a").write(json.dumps(rec)+"\n"); return resp
app=web.Application(client_max_size=50*2**20); app.router.add_route("*","/{p:.*}",h)
web.run_app(app,host="127.0.0.1",port=8787,print=None)