· AI Agents  · 5 min read

Kitaru by ZenML (Part 2): Human-in-the-Loop Wait Gates and Framework Adapters

Enterprise agents cannot run unchecked, nor can containers sit idle for hours awaiting approval. Learn how to use kitaru.wait(), design clean checkpoint boundaries, and wrap PydanticAI or LangGraph agents.

Enterprise agents cannot run unchecked, nor can containers sit idle for hours awaiting approval. Learn how to use kitaru.wait(), design clean checkpoint boundaries, and wrap PydanticAI or LangGraph agents.

This is Part 2 of our 3-part technical deep dive into Kitaru by ZenML. Read Part 1: Architecture & Durable Execution first if you missed it.


In Part 1, we saw how @flow and @checkpoint protect AI agents from mid-run crashes and redundant LLM token spend.

In real enterprise workflows, however, agents do not just fail due to network timeouts. They also need to stop and ask for permission. Whether an agent is about to terminate an idle GPU cluster, issue a £2,000 customer credit, or merge a database migration PR, compliance teams require a Human-in-the-Loop (HITL) approval gate.

If a Python worker blocks on input() or polls a database in a while loop for four hours waiting for a manager to click “Approve” in Slack, you waste container compute and risk losing the process to an autoscaler scale-in event.


1. Zero-Compute Suspension with kitaru.wait()

Kitaru’s wait() primitive turns human approval into a durable state transition:

  1. When the @flow reaches kitaru.wait(prompt=...), Kitaru persists the current execution state and transitions the run to a paused status.
  2. The worker process exits cleanly, releasing CPU and memory resources entirely.
  3. Hours or days later, when an operator submits an approval payload (via CLI, UI, or webhook), Kitaru wakes the flow, replays completed upstream @checkpoint functions from cache in milliseconds, injects the human decision at the wait() boundary, and resumes execution.

Step-by-Step Implementation: A Governed Cloud FinOps Remediation Agent

Let’s build a complete FinOps agent that scans billing anomalies, drafts a Terraform patch, pauses via kitaru.wait() if the monthly cost impact exceeds £500, and applies the change only upon explicit approval:

import json
from typing import Dict, Any
import kitaru
from kitaru import flow, checkpoint


@checkpoint
def fetch_billing_anomalies(business_unit: str) -> Dict[str, Any]:
    """
    Boundary 1: Expensive BigQuery / Snowflake billing scan.
    """
    print(f"[Checkpoint] Scanning billing telemetry for {business_unit}...")
    return {
        "business_unit": business_unit,
        "resource_id": "gke-ml-training-pool-a100",
        "idle_hours": 96,
        "estimated_monthly_waste_gbp": 2450.00,
    }


@checkpoint
def generate_remediation_plan(anomaly: Dict[str, Any]) -> Dict[str, Any]:
    """
    Boundary 2: Frontier LLM reasoning step.
    Cached so resuming after human approval never re-bills LLM tokens.
    """
    print(f"[Checkpoint] Generating remediation plan for {anomaly['resource_id']}...")
    return {
        "target_resource": anomaly["resource_id"],
        "action": "scale_node_pool_to_zero",
        "monthly_savings_gbp": anomaly["estimated_monthly_waste_gbp"],
        "terraform_patch": (
            'resource "google_container_node_pool" "ml_pool" { node_count = 0 }'
        ),
    }


@checkpoint
def apply_remediation(plan: Dict[str, Any], approved_by: str) -> str:
    """
    Boundary 3: Side-effecting infrastructure mutation.
    """
    print(f"[Checkpoint] Applying patch to {plan['target_resource']} (approved by {approved_by})")
    return json.dumps(
        {
            "status": "APPLIED",
            "resource": plan["target_resource"],
            "approved_by": approved_by,
            "savings_gbp": plan["monthly_savings_gbp"],
        }
    )


@flow
def finops_remediation_agent(business_unit: str) -> str:
    anomaly = fetch_billing_anomalies(business_unit)
    plan = generate_remediation_plan(anomaly)

    if plan["monthly_savings_gbp"] > 500.00:
        approval = kitaru.wait(
            prompt=(
                f"Approve scaling '{plan['target_resource']}' to zero "
                f"(saving £{plan['monthly_savings_gbp']:.2f}/month)?"
            )
        )
        if not approval or not approval.get("approved", False):
            return json.dumps({"status": "REJECTED", "resource": plan["target_resource"]})

        operator = approval.get("operator", "finops-lead")
    else:
        operator = "auto-policy"

    return apply_remediation(plan, approved_by=operator)


if __name__ == "__main__":
    execution = finops_remediation_agent.run("EMEA-AI-LABS")
    print(execution.wait())

2. Three Rules for Designing Clean Checkpoint Boundaries

Because Kitaru replays the @flow function from the top up to the last completed @checkpoint or wait() boundary on resume, follow these three rules to guarantee deterministic recovery:

  1. Checkpoint at Non-Deterministic or Costly Boundaries: Put LLM completions, embedding calls, vector queries, and external HTTP/database requests inside @checkpoint. Keep lightweight deterministic control flow (if/else, list comprehensions, string formatting) in the @flow body.
  2. Never Generate Randomness Between Checkpoints: Avoid calling uuid.uuid4(), random.choice(), or datetime.now() directly inside the @flow body if those values are passed into subsequent checkpoints. Generate timestamps and IDs inside a @checkpoint so that replaying the flow yields identical arguments.
  3. Keep Return Types Serialisable: Return JSON-compatible primitives (dict, list, str, int, float, bool) or Pydantic models from @checkpoint functions. Avoid returning open file handles, database connections, or unserialisable SDK client objects.

3. Integrating Kitaru with Existing Agent Frameworks

You do not have to abandon your preferred agent SDK to adopt Kitaru. Kitaru is framework-agnostic and ships native adapters for:

  • Python: PydanticAI, LangGraph, and the OpenAI Agents SDK.
  • TypeScript: Mastra and the Vercel AI SDK.

The integration pattern is straightforward: keep your agent’s internal tool definitions and prompt templates in PydanticAI or LangGraph, and wrap each logical agent turn or multi-step delegation in a Kitaru @checkpoint inside an outer @flow.

from kitaru import flow, checkpoint


@checkpoint
def run_research_subagent(query: str) -> str:
    # Invoke your existing PydanticAI or LangGraph agent here
    # Kitaru records the full trace and caches the final structured output
    return f"Research findings for: {query}"


@checkpoint
def run_reviewer_subagent(draft: str) -> dict:
    # Invoke a second specialist agent to validate policy compliance
    return {"approved": True, "notes": "All citations verified."}


@flow
def multi_agent_compliance_flow(query: str) -> dict:
    findings = run_research_subagent(query)
    review = run_reviewer_subagent(findings)
    return {"findings": findings, "review": review}

By separating agent cognition (handled by your framework of choice) from durable execution and state persistence (handled by Kitaru), you can upgrade agent frameworks or swap underlying models without rewriting your infrastructure layer.


Next in This Series

In Part 3, we will explore how Kitaru turns production traces into automated evaluations: importing traces from Langfuse or LangSmith, extracting cross-session insights, integrating Kitaru’s MCP server with coding assistants, and running deterministic session replays in CI/CD.

Why Alps Agility?

At Alps Agility, we combine deep AI engineering expertise with production MLOps and cloud architecture to help enterprises build durable, observable, and cost-effective agentic systems.

Contact us today to architect your production AI agent runtime.

Struggling to move AI from prototype to production? We help enterprises build robust, scalable AI architectures. Book a Generative AI Readiness Assessment.

Back to Knowledge Hub

Related Posts

View All Posts »