from failproofai.costcap import SessionCostCap

cost_cap = SessionCostCap(max_usd=5.00, max_parallel=32)

def guarded_agent_infer(session_id, agent_func, *args):
    if cost_cap.should_kill(session_id):
        raise RuntimeError(f"Session {session_id}: cost/parallel cap hit")
    result = agent_func(*args)
    cost_cap.record(session_id, agent_func, result)
    return result