Hamilton Stripe Integration: Restricted API Keys, Spend Caps, and Agent Governance

Hamilton is a Python micro-framework from DAGWorks that turns ordinary functions into a computation DAG — function names become node names, parameter names reference upstream nodes or inputs, and a Driver resolves the graph and executes it in topological order. Teams use it for ML feature pipelines, data transformation, and increasingly for AI agent workflows where the same function-as-node model maps cleanly onto agent steps. Its execution model — driver-level retry on the caller, Parallelizable[T] for concurrent fan-out, and lifecycle adapters (NodeExecutionHook) for cross-cutting behavior like automatic retry — introduces three specific Stripe billing failure modes: a driver retry loop re-executes every node in the graph including the billing function; Parallelizable[T] dispatches N concurrent billing tasks all sharing one unrestricted stripe.api_key with no per-task spend cap; and a NodeExecutionHook adapter that retries transient errors fires a fresh stripe.charges.create() call without an idempotency key each time the hook decides to retry.

This post covers all three failure modes with Python code, and the two-layer governance pattern — content-hash idempotency keys plus per-execution vault keys via a spend-cap proxy — that eliminates all three without restructuring your Hamilton dataflow.

Failure mode 1: driver retry loop re-executes the billing node from the start after a downstream failure

Hamilton's Driver.execute() is synchronous and atomic from the caller's perspective: if any node raises an exception during execution, the entire call raises and no partial result is returned. When a downstream node fails after the billing node has already succeeded — for example, when write_charge_to_db raises a database timeout after charge_customer called stripe.charges.create() and received a charge ID — the driver propagates the exception to the caller. The caller, typically a scheduled job or a service endpoint, catches the exception and retries the entire driver.execute() call. Hamilton has no built-in node-level checkpointing or resume-from-failure: every retry starts from the graph inputs and re-executes every node in dependency order, including charge_customer. Without an idempotency key, Stripe treats the second call as a new charge and creates a second charge object for the same customer.

# billing_module.py — UNSAFE: no idempotency key, driver retry re-charges customer
import stripe
import os

stripe.api_key = os.environ["STRIPE_SECRET_KEY"]  # unrestricted live key

def charge_customer(
    customer_id: str,
    amount_cents: int,
    billing_period: str,
) -> dict:
    # Driver retry starts here — no idempotency key, Stripe creates ch_B on retry
    charge = stripe.charges.create(
        amount=amount_cents,
        currency="usd",
        customer=customer_id,
        description=f"Subscription {billing_period}",
    )
    return {"charge_id": charge.id, "customer_id": customer_id}

def write_charge_to_db(charge_customer: dict, db_conn: object) -> bool:
    # If this raises (e.g. DB timeout), caller retries driver.execute() from the top
    db_conn.execute(
        "INSERT INTO charges VALUES (?, ?, ?)",
        (charge_customer["customer_id"], charge_customer["charge_id"], True),
    )
    return True

# run.py — UNSAFE: retry loop re-runs billing_module.charge_customer on db failure
import hamilton
from hamilton import driver
import billing_module, db_module

dr = driver.Driver({}, billing_module, db_module)

for attempt in range(3):
    try:
        result = dr.execute(
            final_vars=["write_charge_to_db"],
            inputs={
                "customer_id": "cus_abc123",
                "amount_cents": 4999,
                "billing_period": "2026-07",
                "db_conn": get_db_connection(),
            },
        )
        break
    except Exception:
        if attempt == 2:
            raise
        time.sleep(2 ** attempt)
# On DB timeout: ch_A created, write_charge_to_db raises, retry → ch_B created

The fix is to derive a content-hash idempotency key from the billing inputs inside charge_customer and pass it with every stripe.charges.create() call. Because customer_id, amount_cents, and billing_period are Hamilton inputs that remain constant across all retries of the same logical billing job, the key is stable across every retry. Stripe's idempotency layer recognizes the key on the second call and returns the original charge object without creating a new charge.

# billing_module.py — SAFE: content-hash idempotency key survives driver retries
import stripe
import hashlib
import os

stripe.api_key = os.environ["STRIPE_SECRET_KEY"]

def billing_idempotency_key(
    customer_id: str, amount_cents: int, billing_period: str
) -> str:
    raw = f"{customer_id}:{amount_cents}:{billing_period}:hamilton-billing"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]

def charge_customer(
    customer_id: str,
    amount_cents: int,
    billing_period: str,
    billing_idempotency_key: str,  # Hamilton resolves this from the node above
) -> dict:
    charge = stripe.charges.create(
        amount=amount_cents,
        currency="usd",
        customer=customer_id,
        description=f"Subscription {billing_period}",
        idempotency_key=billing_idempotency_key,  # stable across all retries
    )
    return {"charge_id": charge.id, "customer_id": customer_id}

def write_charge_to_db(charge_customer: dict, db_conn: object) -> bool:
    db_conn.execute(
        "INSERT INTO charges VALUES (?, ?, ?)",
        (charge_customer["customer_id"], charge_customer["charge_id"], True),
    )
    return True

Note that billing_idempotency_key is itself a Hamilton node — a pure function of the billing inputs. Hamilton resolves it automatically before calling charge_customer, which takes it as a parameter named billing_idempotency_key. There is no manual wiring required. This pattern also makes the key auditable as a first-class node in your Hamilton graph, not a hidden implementation detail inside a function body.

Failure mode 2: Parallelizable[T] fan-out runs N concurrent billing tasks sharing one unrestricted key

Hamilton 0.83+ supports parallel execution through Parallelizable[T] and Collect[T]. A function that returns Parallelizable[Customer] acts as a generator: Hamilton expands it into N independent execution branches, one per yielded value, and runs all branches concurrently using a ThreadPoolExecutor or ProcessPoolExecutor via the h.execution.executors adapter. Each branch executes the downstream subgraph — including any billing node — with the per-item input. All N branches read the same process-level stripe.api_key: there is no per-branch key isolation. If the data pipeline that generates the customer list has a unit bug — passing amount_dollars where amount_cents is expected, or computing a monthly total from daily data without multiplying by 30 — all N branches call stripe.charges.create() with the wrong amount simultaneously. By the time the first error surfaces from Stripe's response validation, every customer in the batch has been charged the wrong amount with no circuit breaker to halt the remaining branches.

# batch_billing_module.py — UNSAFE: Parallelizable fan-out shares one unrestricted key
from hamilton.htypes import Parallelizable, Collect
import stripe
import os
from dataclasses import dataclass

stripe.api_key = os.environ["STRIPE_SECRET_KEY"]  # shared across all N branches

@dataclass
class Customer:
    customer_id: str
    amount_cents: int
    billing_period: str

def customers_to_bill(customers_json: list) -> Parallelizable[Customer]:
    for c in customers_json:
        yield Customer(**c)  # N branches, all sharing stripe.api_key above

def charge_customer(customers_to_bill: Customer) -> dict:
    # Runs N times concurrently — same unrestricted key, no per-customer spend cap
    charge = stripe.charges.create(
        amount=customers_to_bill.amount_cents,  # unit bug → 100x charges all N customers
        currency="usd",
        customer=customers_to_bill.customer_id,
    )
    return {"charge_id": charge.id, "customer_id": customers_to_bill.customer_id}

def billing_results(charge_customer: Collect[dict]) -> list:
    return list(charge_customer)

The fix is to issue a per-customer vault key before the fan-out and pass it through the Customer dataclass so each branch uses its own key. Each vault key is scoped to POST /v1/charges and capped at the individual customer's expected charge amount plus a 10% buffer. A unit bug that passes amount_dollars instead of amount_cents — producing a $499.99 charge where $4.99 was expected — exhausts the vault key's cap on the very first Stripe call. The proxy blocks the oversized charge and returns a 402 error, preventing the charge for that customer. Because each branch has its own vault key, the cap failure for one branch does not affect others: each branch independently attempts its charge and fails fast against its own per-customer cap. A consistent unit bug that affects every customer will fail every branch — which is the correct behavior, since you want to detect the bug before any customers are charged, not after all of them have been overcharged.

# batch_billing_module.py — SAFE: per-customer vault keys issued before fan-out
from hamilton.htypes import Parallelizable, Collect
import stripe
import httpx
import hashlib
import os
from dataclasses import dataclass

KEYBRAKE_ADMIN_KEY = os.environ["KEYBRAKE_ADMIN_KEY"]
KEYBRAKE_PROXY_URL = os.environ.get("KEYBRAKE_PROXY_URL", "https://proxy.keybrake.com")

@dataclass
class Customer:
    customer_id: str
    amount_cents: int
    billing_period: str
    vault_key: str  # issued before fan-out, capped at amount × 1.10

def issue_vault_key(customer_id: str, amount_cents: int, billing_period: str) -> str:
    resp = httpx.post(
        f"{KEYBRAKE_PROXY_URL}/keys",
        headers={"Authorization": f"Bearer {KEYBRAKE_ADMIN_KEY}"},
        json={
            "label": f"hamilton-{customer_id}-{billing_period}",
            "vendor": "stripe",
            "daily_usd_cap": round(amount_cents * 1.10 / 100, 2),
            "allowed_endpoints": ["POST /v1/charges"],
            "expires_in_seconds": 3600,
        },
    )
    resp.raise_for_status()
    return resp.json()["vault_key"]

def customers_to_bill(customers_json: list) -> Parallelizable[Customer]:
    for c in customers_json:
        vault_key = issue_vault_key(
            customer_id=c["customer_id"],
            amount_cents=c["amount_cents"],
            billing_period=c["billing_period"],
        )
        yield Customer(**c, vault_key=vault_key)  # each branch gets its own capped key

def billing_idempotency_key_for_customer(customers_to_bill: Customer) -> str:
    c = customers_to_bill
    raw = f"{c.customer_id}:{c.amount_cents}:{c.billing_period}:hamilton-billing"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]

def charge_customer(
    customers_to_bill: Customer,
    billing_idempotency_key_for_customer: str,
) -> dict:
    # Each branch uses its own vault key — capped at this customer's amount × 1.10
    stripe.api_key = customers_to_bill.vault_key
    stripe.api_base = f"{KEYBRAKE_PROXY_URL}/stripe"
    charge = stripe.charges.create(
        amount=customers_to_bill.amount_cents,
        currency="usd",
        customer=customers_to_bill.customer_id,
        description=f"Subscription {customers_to_bill.billing_period}",
        idempotency_key=billing_idempotency_key_for_customer,
    )
    return {"charge_id": charge.id, "customer_id": customers_to_bill.customer_id}

def billing_results(charge_customer: Collect[dict]) -> list:
    return list(charge_customer)

One implementation note: setting stripe.api_key and stripe.api_base inside charge_customer rather than at module level is deliberate. With concurrent branches running in a ThreadPoolExecutor, module-level mutation of the stripe module attributes is not thread-safe — one branch's write can race with another's read. Passing the key and base URL as local overrides to each stripe.charges.create() call via the Stripe client instance is the safe pattern for concurrent Hamilton execution.

Failure mode 3: NodeExecutionHook lifecycle adapter retries the billing node without an idempotency key

Hamilton's lifecycle adapter system allows cross-cutting behavior to be attached to graph execution without modifying individual functions. The NodeExecutionHook interface receives callbacks before and after each node executes, and on node failure. A common production pattern is a retry adapter that catches transient errors — APIConnectionError, RateLimitError, or any IOError — on any node and re-executes the node callable directly. When this adapter wraps the billing node and a network timeout occurs after stripe.charges.create() has already submitted the charge but before the HTTP response arrives at the client, the adapter classifies the Timeout as a transient error and re-calls the node. The second call reaches Stripe without an idempotency key — because the adapter calls the raw node function directly, bypassing any key the function would normally construct from its inputs — and Stripe creates a second charge.

# retry_adapter.py — UNSAFE: retries billing node after timeout without idempotency key
from hamilton import lifecycle
from hamilton.lifecycle import base

class RetryOnTransientError(base.BaseDoNodeExecuteHook):
    """Retry any node up to 3 times on transient network errors."""

    def do_node_execute(
        self,
        *,
        run_id: str,
        graph_position: "NodeGroupPurpose",
        node_: "Node",
        kwargs: dict,
        task_id: str | None,
    ):
        last_exc = None
        for attempt in range(3):
            try:
                return node_.callable(**kwargs)  # re-calls charge_customer on retry
            except (stripe.error.APIConnectionError, TimeoutError) as e:
                last_exc = e
                time.sleep(2 ** attempt)
        raise last_exc

# run.py — adapter attached to driver wraps ALL nodes including charge_customer
from hamilton import driver
import billing_module
from retry_adapter import RetryOnTransientError

dr = (
    driver.Builder()
    .with_modules(billing_module)
    .with_adapters(RetryOnTransientError())  # wraps every node, including billing
    .build()
)

result = dr.execute(
    final_vars=["write_charge_to_db"],
    inputs={
        "customer_id": "cus_abc123",
        "amount_cents": 4999,
        "billing_period": "2026-07",
        "db_conn": get_db_connection(),
    },
)
# Network timeout after charge_customer submits charge → adapter retries node →
# charge_customer runs again with no idempotency key → Stripe creates ch_B

The failure is subtle because the adapter is a generic utility shared across all nodes. A developer adding it to improve reliability of I/O-bound nodes — database writes, API reads, external enrichment calls — does not intend for it to apply to billing. But Hamilton's adapter system applies hooks uniformly: every node receives the same hooks unless you explicitly filter by node name or tags.

The fix is two-part. First, apply the retry adapter only to non-billing nodes by filtering on node tags or names. Second, ensure the billing node itself always carries a stable idempotency key so that even if the adapter fires it multiple times, Stripe's idempotency cache returns the same charge object on every call after the first.

# retry_adapter.py — SAFE: skips billing nodes by tag; billing node handles its own safety
from hamilton import lifecycle
from hamilton.lifecycle import base

BILLING_NODE_NAMES = {"charge_customer"}  # explicitly excluded from retry

class RetryOnTransientError(base.BaseDoNodeExecuteHook):
    """Retry non-billing nodes on transient errors. Billing nodes excluded."""

    def do_node_execute(self, *, node_: "Node", kwargs: dict, **_):
        if node_.name in BILLING_NODE_NAMES:
            # Billing node runs exactly once — idempotency key handles Stripe safety
            return node_.callable(**kwargs)

        last_exc = None
        for attempt in range(3):
            try:
                return node_.callable(**kwargs)
            except (stripe.error.APIConnectionError, TimeoutError) as e:
                last_exc = e
                time.sleep(2 ** attempt)
        raise last_exc

With the billing node excluded from automatic retry and the content-hash idempotency key present in the node's function body, the failure path becomes safe: a timeout after the charge is submitted causes the driver to raise, the caller's retry loop re-executes the driver, and the stable idempotency key causes Stripe to return the original charge on the second attempt — whether the timeout was at the network level, in Hamilton's node execution layer, or in the lifecycle adapter. All three paths converge on the same safe outcome.

Governance approach comparison

Gap analysis: four Hamilton billing edge cases the main patterns don't cover

1. driver.execute(overrides={"charge_customer": stale_result}) returning a refunded charge ID. Hamilton's overrides dict lets callers inject pre-computed values for any node, bypassing its computation. A caching layer that stores billing results and injects them via overrides to avoid re-charging customers on graph re-runs is a valid optimization — but if the injected charge ID belongs to a charge that was subsequently refunded or disputed, downstream nodes that treat charge_customer["charge_id"] as confirmation of a successful payment will produce incorrect accounting. The fix: always verify a cached charge ID against Stripe's GET /v1/charges/{id} endpoint (using an audit vault key scoped to GET /v1/charges) before injecting it as an override. The verification node should itself be part of the Hamilton graph so it is executed as a dependency of any node that acts on the charge ID.

2. DAGWorks Hamilton UI "replay run" re-executing the production graph including billing nodes. The DAGWorks platform (the commercial layer on top of open-source Hamilton) stores run artifacts and supports replaying historical runs for debugging. A replay run re-executes the full graph with the original inputs. If the graph includes a billing node, a replay run against a production run's inputs will call stripe.charges.create() again. The content-hash idempotency key built from the original inputs will match the original charge and Stripe will return it safely — but only if the same billing_period input is passed. A replay run that modifies the billing_period input to distinguish it from the original run generates a different idempotency key and creates a new charge. The fix: tag billing nodes with a safe_to_replay=False metadata tag and configure the DAGWorks platform to skip or mock nodes with that tag during replay execution.

3. AsyncDriver concurrent execution of multiple inputs shares the stripe module's global state. Hamilton's h.async_driver.AsyncDriver can execute the same graph for multiple independent inputs concurrently using asyncio.gather(). If the billing module sets stripe.api_key at module level and two concurrent async executions both invoke charge_customer, they share the same module-level key. If one execution's vault key is set as a global override just before the other execution reads it, the second execution may charge using the first execution's vault key — charging to the wrong per-customer cap. The fix: pass the vault key as a Hamilton input (not a module-level global) and use stripe.Stripe(api_key=vault_key) client instances per execution rather than mutating stripe.api_key.

4. h.graph.FunctionGraph visualization including billing inputs leaking to observability tools. Hamilton integrates with several observability platforms (OpenTelemetry, MLflow, the DAGWorks UI) and exports node inputs and outputs as trace attributes. A billing node's inputs — customer_id, amount_cents, vault key — are emitted as span attributes unless explicitly filtered. Vault keys in trace data expose them to anyone with access to the observability backend. The fix: implement a NodeExecutionHook that redacts vault_key from node kwargs before they are emitted to any span or log, and verify that the Hamilton adapter you use for observability supports input filtering.

Enforcement: five pytest tests for Hamilton billing graphs

import hashlib
import pytest
from unittest.mock import MagicMock, patch
from hamilton import driver
import billing_module

@pytest.fixture
def dr():
    return driver.Driver(
        {},
        billing_module,
        adapter=None,  # no retry adapter in tests
    )

def make_key(customer_id, amount_cents, billing_period):
    raw = f"{customer_id}:{amount_cents}:{billing_period}:hamilton-billing"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]

def test_idempotency_key_stable_across_driver_retries(dr):
    """Same inputs always produce the same idempotency key node output."""
    result1 = dr.execute(
        ["billing_idempotency_key"],
        inputs={"customer_id": "cus_001", "amount_cents": 4999, "billing_period": "2026-07"},
    )
    result2 = dr.execute(
        ["billing_idempotency_key"],
        inputs={"customer_id": "cus_001", "amount_cents": 4999, "billing_period": "2026-07"},
    )
    assert result1["billing_idempotency_key"] == result2["billing_idempotency_key"]

def test_idempotency_key_distinct_per_billing_period(dr):
    """Different billing periods produce different idempotency keys."""
    key_july = dr.execute(
        ["billing_idempotency_key"],
        inputs={"customer_id": "cus_001", "amount_cents": 4999, "billing_period": "2026-07"},
    )["billing_idempotency_key"]
    key_aug = dr.execute(
        ["billing_idempotency_key"],
        inputs={"customer_id": "cus_001", "amount_cents": 4999, "billing_period": "2026-08"},
    )["billing_idempotency_key"]
    assert key_july != key_aug

def test_idempotency_key_distinct_per_customer(dr):
    """Different customers produce different idempotency keys."""
    key_a = dr.execute(
        ["billing_idempotency_key"],
        inputs={"customer_id": "cus_001", "amount_cents": 4999, "billing_period": "2026-07"},
    )["billing_idempotency_key"]
    key_b = dr.execute(
        ["billing_idempotency_key"],
        inputs={"customer_id": "cus_002", "amount_cents": 4999, "billing_period": "2026-07"},
    )["billing_idempotency_key"]
    assert key_a != key_b

def test_vault_key_blocks_oversized_charge():
    """Vault key capped at amount × 1.10 blocks unit-bug charge (100× amount)."""
    import httpx
    from unittest.mock import patch as mock_patch

    with mock_patch("httpx.post") as mock_post:
        mock_post.return_value = MagicMock(
            status_code=200,
            json=lambda: {"vault_key": "vk_test_abc123"},
        )
        from batch_billing_module import issue_vault_key
        vault_key = issue_vault_key("cus_001", 4999, "2026-07")
        call_kwargs = mock_post.call_args.kwargs["json"]
        assert call_kwargs["daily_usd_cap"] == pytest.approx(4999 * 1.10 / 100, rel=0.01)
        assert call_kwargs["allowed_endpoints"] == ["POST /v1/charges"]

def test_retry_adapter_excludes_billing_node():
    """RetryOnTransientError adapter does not retry charge_customer on timeout."""
    from retry_adapter import RetryOnTransientError
    adapter = RetryOnTransientError()

    call_count = 0
    def failing_billing_fn(**kwargs):
        nonlocal call_count
        call_count += 1
        raise TimeoutError("network timeout")

    node_mock = MagicMock()
    node_mock.name = "charge_customer"
    node_mock.callable = failing_billing_fn

    with pytest.raises(TimeoutError):
        adapter.do_node_execute(node_=node_mock, kwargs={}, run_id="r1", task_id=None, graph_position=None)

    assert call_count == 1  # called exactly once, not retried

Frequently asked questions

Can I use Hamilton's overrides dict to skip the billing node on a driver retry?

Yes — injecting a pre-existing charge ID via overrides={"charge_customer": cached_result} prevents the billing node from running and is a valid optimization. The risk is injecting a stale charge ID (refunded, disputed, or from the wrong billing period). Always validate cached charge IDs against Stripe's API before injecting them, and include the billing period in the cache key so a new period always triggers a fresh charge. Use a content-hash idempotency key as a second layer even when using overrides, in case the override is bypassed or the cache key is constructed incorrectly.

How do vault keys interact with Hamilton's functional parameter injection?

Vault keys work naturally as Hamilton inputs. Pass them as part of the inputs dict to driver.execute() or embed them inside a dataclass that Hamilton injects as a typed input. For Parallelizable[T] patterns, issue each vault key before the fan-out in a preparatory function and attach it to the per-item dataclass so each branch carries its own key. Vault key issuance should be a separate, non-retried step: if issuance fails for any customer, fail immediately rather than falling back to a raw Stripe key.

What happens if the Keybrake proxy is unreachable when a Hamilton billing node executes?

The stripe.charges.create() call will fail with a connection error because the proxy is the configured stripe.api_base. This is the correct behavior: the Hamilton driver raises, the caller's retry logic fires, and the billing node re-runs with the same idempotency key. No charge is created until the proxy is reachable. Do not configure a fallback to api.stripe.com directly — the fallback bypasses the spend cap and defeats the entire governance layer. If the proxy is down, billing should wait, not proceed unguarded.

Should I set stripe.api_key at module level or pass it through Hamilton inputs?

Pass it through Hamilton inputs or create a stripe.Stripe(api_key=...) client instance inside the billing function and use that client for the API call. Module-level stripe.api_key mutation is not thread-safe in concurrent Hamilton execution (ThreadPoolExecutor, AsyncDriver with asyncio.gather()). If you set stripe.api_key inside charge_customer, every concurrent branch immediately races on the global. Using a per-call client instance — client = stripe.Stripe(api_key=vault_key); client.charges.create(...) — avoids the race entirely and makes key scoping explicit in the function signature.

How do I test Hamilton billing graphs without calling Stripe's live API?

Use driver.execute(overrides={"charge_customer": mock_result}) to inject a mock charge result and test the downstream nodes (database write, confirmation email) independently from Stripe. For testing the billing node itself, patch stripe.charges.create at the function level and verify the idempotency key sent to Stripe matches the expected content-hash. Test the billing_idempotency_key node in isolation — it is a pure function with no external dependencies and should be fully deterministic. The five pytest tests above cover all three failure modes without any live Stripe calls.

Does Hamilton's AsyncDriver change how idempotency keys should be derived?

No — the key derivation logic is the same. What changes is where you place the stripe.Stripe client instantiation: in async Hamilton graphs, instantiate the client inside the async billing coroutine (not at module level) so each concurrent invocation gets its own client object. For vault key issuance before a parallel fan-out in an async graph, use asyncio.gather() to issue all keys concurrently before dispatching the billing subgraph, rather than issuing them serially in the fan-out generator. The idempotency key itself — derived from customer_id, amount_cents, and billing_period — is stable across both sync and async drivers.