Distributed Runtime Architecture¶

import restate
from ravi.distributed.activities import do_llm_call, do_tool_exec

agent_workflow = restate.Workflow("AgentWorkflow")

@agent_workflow.main()
async def run(ctx: restate.WorkflowContext, payload: dict) -> dict:
    """Durable ReAct agent loop.

    Each LLM call and tool execution is a separate ctx.run() — the minimum
    durable unit. If the process crashes, Restate replays completed steps
    from its journal and resumes at the exact point of failure.
    """
    thread_id = payload["thread_id"]
    user_content = payload["user_content"]
    claims = payload["claims"]  # user_id, role from JWT
    model = payload.get("model")
    max_iterations = payload.get("max_iterations", 50)

    # Restore conversation memory (Redis) — idempotent
    await ctx.run("restore_memory", restore_memory, thread_id=thread_id)

    for iteration in range(max_iterations):
        # 1. THINK — call LLM (durable step)
        #    Inside this ctx.run(), the LLM callback publishes text_delta
        #    tokens to NATS for live streaming to the frontend.
        llm_result = await ctx.run(
            f"llm_call_{iteration}",
            do_llm_call,
            thread_id=thread_id,
            model=model,
        )

        # No tool calls → agent is done
        if not llm_result.get("tool_calls"):
            await publish_nats(thread_id, {"type": "completion", **llm_result})
            break

        # 2. ACT — execute each tool call (durable step per tool)
        for tc in llm_result["tool_calls"]:
            tool_name = tc["name"]
            tool_policy = get_tool_policy(tool_name)

            if tool_policy.requires_approval:
                # HITL: suspend until human approves/rejects
                request_id = ctx.uuid()
                await publish_nats(thread_id, {
                    "type": "tool_approval_request",
                    "requestId": request_id,
                    "tool_name": tool_name,
                    "input": tc["arguments"],
                })
                decision = await ctx.promise(f"approval-{request_id}")
                if decision.get("action") == "reject":
                    # Record rejection and continue to next iteration
                    await ctx.run(f"reject_{request_id}", record_rejection, ...)
                    continue

            if tool_policy.is_hitl_input:
                # Human input: suspend until human responds
                request_id = ctx.uuid()
                await publish_nats(thread_id, {
                    "type": "human_input_request",
                    "requestId": request_id,
                    "prompt": tc["arguments"].get("prompt", ""),
                })
                human_response = await ctx.promise(f"human-input-{request_id}")
                tool_result = {"content": human_response.get("response", "")}
            else:
                # Regular tool execution (durable, with per-tool timeout)
                idempotency_key = ctx.uuid() if tool_policy.needs_idempotency else None
                tool_result = await ctx.run(
                    f"tool_{tool_name}_{iteration}",
                    do_tool_exec,
                    tool_name=tool_name,
                    arguments=tc["arguments"],
                    timeout_seconds=tool_policy.timeout,
                    idempotency_key=idempotency_key,
                )

            # Publish tool result to NATS for frontend
            await publish_nats(thread_id, {
                "type": "tool_result",
                "tool_name": tool_name,
                **tool_result,
            })

        # 3. OBSERVE — results are in memory, loop continues

    return {"status": "completed", "thread_id": thread_id}

class ToolPolicy:
    """Execution policy for distributed runtime."""
    timeout: float = 30.0          # seconds, 0 = no timeout
    needs_idempotency: bool = False # generate ctx.uuid() idempotency key
    requires_approval: bool = False # HITL tool approval gate
    is_hitl_input: bool = False     # HITL human input (ask_human)
    large_payload: bool = False     # use DataRef for results > 32KB

# Per-tool policy examples:
TOOL_POLICIES = {
    # HITL tools — suspend with zero resources, no timeout
    "ask_human":          ToolPolicy(is_hitl_input=True),

    # Critical tools — idempotency key, approval required
    "send_email":         ToolPolicy(timeout=30, needs_idempotency=True, requires_approval=True),
    "manage_tasks":       ToolPolicy(timeout=10),

    # Network tools — configurable timeout
    "web_surfer":         ToolPolicy(timeout=60, large_payload=True),
    "db_query":           ToolPolicy(timeout=30),

    # Compute tools — long timeout with heartbeat
    "code_interpreter":   ToolPolicy(timeout=300, large_payload=True),

    # In-process tools — direct call, no distribution overhead
    "capability_search":  ToolPolicy(timeout=5),
    "clock":              ToolPolicy(timeout=1),
}

import nats
from nats.js import JetStreamContext

class NATSStreamingBridge:
    """Publishes SSE events from workers → NATS → Gateway → Frontend.

    Subject pattern: agent.events.{thread_id}
    Stream: AGENT_EVENTS (retention: 1 hour, per-subject)
    """

    def __init__(self, nats_url: str = "nats://localhost:4222"):
        self._nc: nats.NATS = None
        self._js: JetStreamContext = None

    async def connect(self):
        self._nc = await nats.connect(self._nats_url)
        self._js = self._nc.jetstream()
        # Create stream if not exists
        await self._js.add_stream(
            name="AGENT_EVENTS",
            subjects=["agent.events.*"],
            retention="limits",
            max_age=3600,  # 1 hour retention
        )

    async def publish(self, thread_id: str, event: dict):
        """Publish an SSE event for a specific thread."""
        subject = f"agent.events.{thread_id}"
        await self._js.publish(subject, json.dumps(event).encode())

    async def subscribe(self, thread_id: str):
        """Subscribe to SSE events for a specific thread. Used by gateway."""
        subject = f"agent.events.{thread_id}"
        sub = await self._js.subscribe(subject, deliver_policy="new")
        async for msg in sub.messages:
            yield json.loads(msg.data.decode())
            await msg.ack()

@router.post("/chat")
async def chat(body: ChatRequest, request: Request):
    # 1. Validate thread, extract JWT claims
    claims = get_current_user(request)

    # 2. Start durable workflow via Restate
    restate_client = request.app.state.restate_client
    workflow_id = f"agent-{body.thread_id}-{uuid4().hex[:8]}"
    await restate_client.workflow("AgentWorkflow").send(
        workflow_id,
        payload={
            "thread_id": str(body.thread_id),
            "user_content": body.messages[-1].content,
            "claims": {"user_id": claims.user_id, "role": claims.role},
            "model": body.model,
        },
    )

    # 3. Stream SSE events from NATS
    nats_bridge = request.app.state.nats_bridge
    async def sse_generator():
        async for event in nats_bridge.subscribe(str(body.thread_id)):
            yield f"data: {json.dumps(event)}\n\n"
            if event.get("type") in ("completion", "error", "cancelled"):
                if not event.get("has_tool_calls"):
                    break

    return StreamingResponse(sse_generator(), media_type="text/event-stream")

@router.post("/chat/respond/{request_id}")
async def respond_hitl(request_id: str, body: HITLResponse, request: Request):
    """Resolve a pending HITL promise in Restate.

    When the user approves/rejects a tool or provides human input,
    this resolves the ctx.promise() in the running workflow.
    """
    restate_client = request.app.state.restate_client

    # Resolve the durable promise — Restate routes to the correct workflow
    await restate_client.workflow("AgentWorkflow").resolve_promise(
        workflow_id=body.workflow_id,
        promise_name=f"{body.promise_type}-{request_id}",
        value=body.payload,
    )
    return {"status": "resolved"}

@router.post("/chat/{thread_id}/cancel")
async def cancel_chat(thread_id: str, request: Request):
    """Cancel a running agent workflow via Restate's cancel API."""
    restate_client = request.app.state.restate_client

    # Find active workflow for this thread
    workflow_id = await get_active_workflow(thread_id)
    if workflow_id:
        await restate_client.workflow("AgentWorkflow").cancel(workflow_id)

    return {"status": "cancelled"}

async def do_llm_call(*, thread_id: str, model: str, **kwargs) -> dict:
    """Activity: call LLM with streaming tokens published to NATS."""
    nats_bridge = get_nats_bridge()  # global singleton
    client = get_model_client(model)

    full_response = ""
    async for chunk in client.generate_stream(messages=..., tools=...):
        if chunk.text:
            full_response += chunk.text
            # Publish token to NATS (ephemeral, fire-and-forget)
            await nats_bridge.publish(thread_id, {
                "type": "text_delta",
                "content": chunk.text,
                "partial": True,
            })

    # Return full response for Restate journal (durable)
    return {
        "content": full_response,
        "tool_calls": chunk.tool_calls,
        "usage": chunk.usage,
    }

async def do_tool_exec(*, tool_name: str, arguments: dict, **kwargs) -> dict:
    tool = get_tool(tool_name)
    result = await tool.run(**arguments)

    policy = TOOL_POLICIES.get(tool_name, ToolPolicy())

    if policy.large_payload and len(str(result.content)) > 32_768:
        # Store in S3/file store, return reference
        data_store = get_data_store()
        ref = await data_store.store(
            data=str(result.content).encode(),
            content_type="text/plain",
        )
        return {
            "content": f"[Result stored as DataRef: {ref.ref_id}]",
            "data_ref_id": str(ref.ref_id),
            "success": True,
        }

    return {
        "content": result.content,
        "app_data": result.app_data,
        "success": True,
    }

@agent_workflow.main()
async def run(ctx: restate.WorkflowContext, payload: dict):
    # Orchestrator decides to hand off to a sub-agent
    sub_result = await ctx.service_call(
        "AgentWorkflow",
        "run",
        payload={
            "thread_id": payload["thread_id"],
            "user_content": "Research this topic...",
            "claims": payload["claims"],
            "model": "gpt-4o-mini",  # sub-agent can use different model
        },
    )
    # Sub-agent result is durable — survives orchestrator crash