from statebase import StateBase

sb = StateBase(api_key="your-key")

def orchestrator(session_id, task):
    # 1. Decompose
    subtasks = llm.decompose(task)

    # 2. Record the plan in shared state
    sb.sessions.update_state(
        session_id=session_id,
        state={"plan": subtasks, "results": {}},
        reasoning=f"Decomposed task into {len(subtasks)} subtasks"
    )

    # 3. Dispatch workers (each writes to the same session)
    for i, subtask in enumerate(subtasks):
        worker_id = dispatch_worker(session_id, i, subtask)
        sb.sessions.update_state(
            session_id=session_id,
            state={f"worker_{i}": worker_id},
            reasoning=f"Dispatched worker {i}"
        )

    # 4. Wait for all workers, merge results
    while not all_workers_done(session_id):
        time.sleep(5)

    results = sb.sessions.get(session_id).state["results"]
    return merge(results)

def worker(session_id, worker_id, subtask):
    sb.sessions.update_state(
        session_id=session_id,
        state={f"worker_{worker_id}": "in_progress"},
        reasoning=f"Worker {worker_id} started"
    )

    result = execute(subtask)

    # Write result back to the shared session
    state = sb.sessions.get(session_id).state
    sb.sessions.update_state(
        session_id=session_id,
        state={"results": {**state.get("results", {}), worker_id: result}},
        reasoning=f"Worker {worker_id} completed"
    )

def supervisor(session_id, goal):
    while not goal_achieved(session_id):
        plan = llm.plan(goal)

        # Checkpoint the plan before acting
        sb.sessions.update_state(
            session_id=session_id,
            state={"plan": plan, "iteration": plan_count + 1},
            reasoning="New plan"
        )

        for action in plan["actions"]:
            result = execute(action)

        # Evaluate — rollback if the loop is off track
        if not evaluate(session_id):
            sb.sessions.rollback(session_id=session_id, version=-1)

sb.sessions.add_turn(
    session_id=session_id,
    input={"from": "planner", "action": "handoff"},
    output={"to": "executor", "context": handoff_context},
    metadata={"handoff": True, "depth": 1},
    reasoning="Planner -> Executor handoff"
)

# Any agent can persist a durable fact
sb.memory.add(
    session_id=session_id,
    content="Production API returns 429 under 1000 RPS",
    tags=["infra", "learned"],
    type="observation"
)

# Any agent can retrieve it
memories = sb.memory.search(
    session_id=session_id,
    query="rate limit behavior"
)