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

Mage AI is an open-source data pipeline framework where work is structured as typed Python blocks — data_loader, transformer, data_exporter — arranged in a DAG and executed as a pipeline run. Each run is independent, has its own execution environment, and can be triggered by a cron schedule, an API call, or a UI backfill. When the work being executed is a Stripe charge, that independence introduces three distinct failure modes that idempotency-free billing blocks cannot survive.

This post covers three Mage-specific failure modes — block retry re-execution, concurrent trigger runs, and backfill parallel runs — each of which causes a second stripe.charges.create() call on an already-billed customer, and the two-layer governance pattern that closes all three: content-hash idempotency keys plus per-run vault keys via a spend-cap proxy, without restructuring your pipeline DAG.

Failure mode 1: Block retry re-invokes the billing transformer from line 1 after stripe.charges.create() has already succeeded

Mage blocks support retry configuration via a retry_config dict in the block's metadata, or at the pipeline level. When a block raises an unhandled Python exception, Mage waits the configured delay and re-executes the entire block function from line 1. This is the correct behavior for idempotent operations like database reads or API GETs — but it is dangerous for billing blocks where the Stripe call is not the last line.

The failure window is between stripe.charges.create() returning successfully and the final statement in your block completing. Consider a transformer that charges a customer and then writes the charge ID to a database. If the database write raises a connection exception — or if any subsequent code in the block fails — Mage catches the exception, waits the retry delay, and re-executes the block. The Stripe call fires again. Without an idempotency key, Stripe creates ch_B. The customer has now been charged twice, and both charge objects are valid in your Stripe dashboard.

# transformers/bill_customer.py — UNSAFE: no idempotency key
import stripe
import os

stripe.api_key = os.environ["STRIPE_SECRET_KEY"]  # shared, unrestricted

@transformer
def transform(customer: dict, *args, **kwargs):
    # If this succeeds but the DB write below fails → block retried → ch_B
    charge = stripe.charges.create(
        amount=customer["amount_cents"],
        currency="usd",
        customer=customer["stripe_customer_id"],
        description=f"Subscription {customer['billing_period']}",
    )
    # Connection exception here → Mage re-executes block from line 1
    db.execute(
        "INSERT INTO charges VALUES (?, ?, ?)",
        [charge.id, customer["id"], customer["billing_period"]],
    )
    return {"charge_id": charge.id, "status": "charged"}

The fix is a content-hash idempotency key computed from inputs that are stable across all block re-executions. Mage passes the same block input data on every retry — the customer dict is identical — so a key derived from customer_id + amount_cents + billing_period is identical on the retried execution, and Stripe returns the original ch_A rather than creating ch_B. Permanent Stripe errors (card declined, invalid customer) should be returned as a result dict rather than raised; a raised exception triggers another retry cycle that will re-attempt a charge that should not be retried.

# transformers/bill_customer.py — SAFE: stable idempotency key + permanent-error handling
import hashlib
import stripe
import os

@transformer
def transform(customer: dict, vault_key: str, *args, **kwargs):
    # vault_key issued by upstream load_vault_key block, scoped to POST /v1/charges
    idempotency_key = hashlib.sha256(
        f"{customer['id']}:{customer['amount_cents']}:{customer['billing_period']}:mage-billing".encode()
    ).hexdigest()[:32]

    stripe_client = stripe.Stripe(
        api_key=vault_key,
        base_url="https://proxy.keybrake.com/stripe/",
    )

    try:
        charge = stripe_client.charges.create(
            amount=customer["amount_cents"],
            currency="usd",
            customer=customer["stripe_customer_id"],
            description=f"Subscription {customer['billing_period']}",
            idempotency_key=idempotency_key,
        )
        db.execute(
            "INSERT INTO charges VALUES (?, ?, ?)",
            [charge.id, customer["id"], customer["billing_period"]],
        )
        return {"charge_id": charge.id, "status": "charged"}
    except stripe.error.CardError as e:
        # Return as result — do NOT raise; raising triggers block retry on a declined card
        return {"status": "card_declined", "error_code": e.code, "charge_id": None}
    except stripe.error.InvalidRequestError as e:
        return {"status": "invalid_request", "error": str(e), "charge_id": None}

Failure mode 2: Scheduled trigger and concurrent API trigger create two independent pipeline runs for the same billing period

Mage separates pipeline definition from pipeline execution: each trigger — whether a cron schedule or an API call to POST /api/pipeline_runs — creates an independent pipeline run with its own pipeline_run_id. Mage does not enforce at-most-once execution across trigger types for the same execution window. This means a scheduled midnight trigger and an operator-initiated API trigger can both create pipeline runs for the same billing period and execute concurrently.

The scenario is common in practice: an on-call engineer sees what looks like a failed pipeline run in the Mage UI and triggers a manual API run to retry — but the scheduled run was actually mid-execution, not failed. Both runs proceed through all blocks including the billing transformer. Each run independently calls stripe.charges.create() for every customer in the dataset. Without an idempotency key, Stripe creates ch_A from the scheduled run and ch_B from the API-triggered run, milliseconds apart. The customer is charged twice. Both pipeline runs show as successful in the Mage UI because Stripe accepted both requests.

# API trigger creates a second concurrent run — UNSAFE without idempotency key
import requests

# Run 1 already executing from scheduled trigger at midnight
# Run 2 created by operator via API — believes Run 1 failed
resp = requests.post(
    "http://mage-server:6789/api/pipeline_runs",
    headers={"x-api-key": MAGE_API_KEY},
    json={
        "pipeline_run": {
            "pipeline_uuid": "monthly_billing",
            "variables": {
                "billing_period": "2026-07",
                "execution_date": "2026-07-01",
            }
        }
    }
)
# Both runs now execute bill_customer block for all customers
# Without idempotency keys: ch_A (from Run 1) + ch_B (from Run 2) for every customer

The content-hash idempotency key from Failure Mode 1 handles concurrent runs transparently: because both pipeline runs receive the same billing_period and customer_id inputs, they compute the same idempotency key, and Stripe deduplicates the second request. The vault key issued by the upstream block also provides a secondary control: configure the vault key with a per-key dollar cap equal to the customer's expected charge amount. Even if the idempotency key is accidentally varied (e.g., because a developer added a timestamp to it), the vault key spend cap prevents the second charge from exceeding the customer's billing amount.

For pipelines where you need to enforce at-most-one execution at the Mage level — not just at the Stripe level — set a concurrency limit on the billing pipeline trigger. In Mage's trigger configuration, set pipeline_concurrency to 1 so that a new trigger attempt is queued or skipped if a run for that pipeline is already active. Combine this with the Stripe-layer idempotency key for defense in depth.

Failure mode 3: Backfill launches N parallel pipeline runs that charge customers simultaneously with no cross-run deduplication

Mage's backfill feature creates one pipeline run per execution interval between the backfill start and end dates, all running in parallel. A thirty-day backfill from July 1 to July 30 creates thirty independent pipeline runs, each executing the full pipeline DAG — including the billing transformer — for its assigned date. All thirty runs execute concurrently, subject only to the parallelism setting on the Mage instance.

The exposure is: every run charges customers independently, with no coordination between runs. If your billing query uses a date range that overlaps between intervals — for example, a monthly aggregate where July 1 and July 2 both pull the same customer's month-to-date charges — two runs will fire stripe.charges.create() for the same customer with identical billing amounts. Without an idempotency key, Stripe creates two charges. With 30 parallel runs and 500 customers per run, a data error in the billing amount query propagates to all 15,000 Stripe calls before any individual run returns a result — there is no early-exit mechanism to halt the remaining runs once the first error surfaces.

# Mage backfill via UI or CLI — UNSAFE without per-date idempotency key
# mage run monthly_billing --start-time 2026-07-01 --end-time 2026-07-30
# Creates 30 parallel pipeline runs:
#   Run 1: execution_date=2026-07-01 → bill_customer block fires for all customers
#   Run 2: execution_date=2026-07-02 → bill_customer block fires for all customers
#   ...
#   Run 30: execution_date=2026-07-30 → bill_customer block fires for all customers
#
# If billing query overlap bug causes cus_ABC to appear in both date 01 and date 02:
#   Run 1 fires stripe.charges.create(customer="cus_ABC", amount=2999) → ch_A
#   Run 2 fires stripe.charges.create(customer="cus_ABC", amount=2999) → ch_B
#   Without idempotency key: customer billed twice, both charges valid

The idempotency key must include the execution date to be unique per backfill run while remaining stable across retries of the same run. The key derivation should be SHA-256(customer_id + ":" + amount_cents + ":" + billing_period + ":" + execution_date + ":" + "mage-billing")[:32]. This key is different for July 1 and July 2 even for the same customer — preventing double-charging if the billing query overlap bug causes the customer to appear in multiple backfill intervals. It is identical on a retry of Run 1 — preventing double-charging if Run 1 fails and is restarted.

# data_loaders/load_vault_key.py — issue per-run vault key before billing block
import requests
import os
from mage_ai.data_preparation.shared.secrets import get_secret_value

@data_loader
def load_vault_key(*args, **kwargs):
    execution_date = kwargs.get("execution_date", "")
    billing_period = kwargs.get("billing_period", "")

    KEYBRAKE_API_KEY = get_secret_value("KEYBRAKE_API_KEY")

    # Estimate max spend for this run: customer_count × max_amount × 1.10
    # In practice, fetch from billing config or pipeline variable
    max_daily_cap = float(kwargs.get("max_run_spend_usd", "5000"))

    resp = requests.post(
        "https://proxy.keybrake.com/keys",
        headers={"Authorization": f"Bearer {KEYBRAKE_API_KEY}"},
        json={
            "vendor": "stripe",
            "allowed_endpoints": ["POST /v1/charges"],
            "daily_usd_cap": max_daily_cap,
            "expires_in_seconds": 3600,
            "label": f"mage-billing-{billing_period}-{execution_date}",
        },
        timeout=10,
    )
    resp.raise_for_status()
    return {"vault_key": resp.json()["vault_key"]}

# transformers/bill_customer.py — SAFE: execution-date-scoped idempotency key
import hashlib
import stripe

@transformer
def transform(customers: list, vault_data: dict, *args, **kwargs):
    vault_key = vault_data["vault_key"]
    billing_period = kwargs.get("billing_period", "")
    execution_date = kwargs.get("execution_date", "")

    stripe_client = stripe.Stripe(
        api_key=vault_key,
        base_url="https://proxy.keybrake.com/stripe/",
    )

    results = []
    for customer in customers:
        idempotency_key = hashlib.sha256(
            f"{customer['id']}:{customer['amount_cents']}:{billing_period}:{execution_date}:mage-billing".encode()
        ).hexdigest()[:32]

        try:
            charge = stripe_client.charges.create(
                amount=customer["amount_cents"],
                currency="usd",
                customer=customer["stripe_customer_id"],
                description=f"Subscription {billing_period}",
                idempotency_key=idempotency_key,
            )
            results.append({"customer_id": customer["id"], "charge_id": charge.id, "status": "charged"})
        except stripe.error.CardError as e:
            results.append({"customer_id": customer["id"], "status": "card_declined", "error_code": e.code})
        except stripe.error.InvalidRequestError as e:
            results.append({"customer_id": customer["id"], "status": "invalid_request", "error": str(e)})

    return results

Approach comparison

Gap analysis

1. data_exporter blocks making direct Stripe calls bypass upstream vault key governance

Mage's block types are conventions, not enforcement boundaries — a data_exporter block can call stripe.charges.create() directly just as a transformer can. Teams occasionally put billing calls in data_exporter blocks because they think of it as "writing output to Stripe." If this exporter block reads STRIPE_SECRET_KEY from the environment rather than accepting a vault key as input from an upstream block, it bypasses the spend cap. All Stripe-calling blocks in the pipeline must receive their vault key from the load_vault_key upstream block and use it via the proxy — not from the environment directly.

2. Mage webhook triggers deliver at least once and can fire duplicate pipeline runs

Mage supports event-triggered pipeline runs via webhooks configured on a trigger. Upstream systems (payment processors, CRMs, event buses) that send Mage webhook events use at-least-once delivery semantics: if the Mage server does not acknowledge within the timeout, the upstream retries the event. If the Mage server is mid-restart or under load when the first event arrives, the acknowledgment is delayed, and the upstream fires a second webhook. Two pipeline runs are created for the same event. The content-hash idempotency key closes this: if the event payload includes a stable customer identifier, billing period, and amount, both runs compute the same Stripe idempotency key and only one charge is created. Include a unique event ID from the webhook payload as an additional field in the idempotency key hash if the upstream system guarantees event ID stability across retries.

3. Mage global pipeline variables are shared across all concurrent runs with no per-run isolation

Mage's get_global_variable(pipeline_uuid, key) function reads a value from the pipeline's global variable store, which is shared across all concurrent pipeline runs. If your billing configuration (pricing tiers, billing multipliers, fee structures) is stored as a global variable and updated mid-backfill to fix a pricing error, all thirty concurrent backfill runs will immediately start reading the updated value. Runs that already completed billing blocks used the old value; runs that reach the billing block after the update use the new value. The same customers may be charged different amounts by different runs in the same backfill. Pass billing configuration as pipeline-run-level variables via the trigger payload, not as global variables, so each run is sealed to its configuration at trigger time.

4. Mage pipeline re-runs from the UI re-execute all blocks including completed billing blocks

When you click "Re-run" on a pipeline run in the Mage UI, Mage creates a new pipeline run that executes all blocks from the beginning — including the load_vault_key block (which issues a new vault key) and the bill_customer block (which charges customers). The new run has a different pipeline_run_id, but if the billing inputs are identical to the original run, the content-hash idempotency key is also identical. Stripe deduplicates the second charge request and returns the original charge object. This is the correct behavior — the idempotency key absorbs the re-run safely. However, the new vault key issued by load_vault_key has a fresh spend cap, so the proxy's per-run accounting resets. Add a pre-flight check in load_vault_key that queries your database for an existing vault key for the billing period before issuing a new one, so re-runs consume the same tracked key.

Enforcement tests

# test_mage_billing.py
import hashlib
import pytest
from unittest.mock import patch, MagicMock

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

def test_idempotency_key_stable_across_block_retries():
    key1 = make_idempotency_key("cus_ABC", 2999, "2026-07")
    key2 = make_idempotency_key("cus_ABC", 2999, "2026-07")
    assert key1 == key2, "Key must be identical across block re-executions"

def test_idempotency_key_differs_by_execution_date():
    key_day1 = make_idempotency_key("cus_ABC", 2999, "2026-07", "2026-07-01")
    key_day2 = make_idempotency_key("cus_ABC", 2999, "2026-07", "2026-07-02")
    assert key_day1 != key_day2, "Backfill runs on different dates must use different keys"

def test_idempotency_key_same_across_concurrent_trigger_runs():
    # Scheduled run and API-triggered run receive same inputs → same key → Stripe deduplicates
    key_scheduled = make_idempotency_key("cus_ABC", 2999, "2026-07")
    key_api_run   = make_idempotency_key("cus_ABC", 2999, "2026-07")
    assert key_scheduled == key_api_run, "Concurrent trigger runs must produce the same idempotency key"

def test_card_declined_returns_struct_not_raises():
    from transformers.bill_customer import transform
    with patch("transformers.bill_customer.stripe.Stripe") as mock_stripe:
        mock_stripe.return_value.charges.create.side_effect = (
            __import__("stripe").error.CardError("Declined", None, "card_declined", http_status=402)
        )
        result = transform(
            customers=[{"id": "cus_A", "stripe_customer_id": "cus_A", "amount_cents": 2999}],
            vault_data={"vault_key": "vk_test"},
            billing_period="2026-07",
            execution_date="2026-07-01",
        )
    assert result[0]["status"] == "card_declined"
    assert result[0].get("charge_id") is None

def test_vault_key_issued_before_billing_block():
    from data_loaders.load_vault_key import load_vault_key
    with patch("data_loaders.load_vault_key.requests.post") as mock_post:
        mock_post.return_value.json.return_value = {"vault_key": "vk_test_abc"}
        mock_post.return_value.raise_for_status = lambda: None
        result = load_vault_key(
            billing_period="2026-07",
            execution_date="2026-07-01",
            max_run_spend_usd="5000",
        )
    assert "vault_key" in result
    call_json = mock_post.call_args.kwargs["json"]
    assert call_json["allowed_endpoints"] == ["POST /v1/charges"]
    assert call_json["daily_usd_cap"] == 5000.0

Frequently asked questions

Can I use Mage's built-in retry_config on a billing block if I add idempotency keys?

Yes — with idempotency keys, block retries are safe for the Stripe call itself. When the retried block fires stripe.charges.create() with the same idempotency key, Stripe returns the original charge object and no second charge is created. The retry may still cause issues in other parts of the block: if a database write is not idempotent (e.g., a plain INSERT without ON CONFLICT DO NOTHING), a second execution of the DB write will raise a uniqueness violation. Make all operations in a billing block idempotent — use INSERT OR IGNORE or ON CONFLICT DO NOTHING for database writes, use Stripe's idempotency key for the charge call, and return error status dicts rather than raising for permanent failures.

What is the difference between retrying a Mage block and retrying a pipeline run?

A block retry re-executes just the failed block with the same input data, within the same pipeline run. The pipeline_run_id is unchanged; block-level retry counts are incremented. A pipeline run re-run (via the UI or API) creates a new pipeline run with a new pipeline_run_id and re-executes all blocks from the beginning. Block retry is the more dangerous case for billing blocks because it is automatic and silent — Mage retries without operator confirmation. Pipeline run re-runs are manual and visible. Both cases are handled by the content-hash idempotency key, which produces the same Stripe key regardless of whether the re-execution is a block retry or a full run re-run.

Is Mage's pipeline_run_id a safe idempotency key for Stripe charges?

No. The pipeline_run_id is different on every pipeline run, including re-runs triggered after the original run failed. If you use pipeline_run_id as the Stripe idempotency key, a pipeline re-run will produce a new key, and Stripe will create a second charge. The idempotency key must be derived from the billing inputs themselves — customer_id, amount_cents, billing_period, and execution_date — not from Mage execution metadata. The pipeline_run_id is useful for logging and audit trail correlation, but it must not be the sole component of the Stripe idempotency key.

How do I prevent a backfill from double-charging a customer who appears in multiple date intervals?

Include execution_date in the idempotency key: SHA-256(customer_id:amount_cents:billing_period:execution_date:mage-billing)[:32]. This makes the key unique per backfill run per customer, so two runs for different dates that both include the same customer produce different idempotency keys — which is correct, because they are intended to be different charges for different periods. If they should not be different charges (i.e., a billing query overlap bug causes the same customer to appear in both), fix the query to prevent overlap. Do not suppress the idempotency key uniqueness to paper over a data bug — that would cause deduplication to hide the actual billing error instead of surfacing it.

How long should vault key TTLs be for backfill pipeline runs?

Set the TTL to cover the longest expected backfill run duration plus a buffer. For a thirty-day backfill where each daily run takes up to ten minutes, the last run could still be executing thirty minutes after the backfill started (assuming high parallelism). Set the vault key TTL to at least 60 × 2 = 120 minutes for this case. If your Mage instance has low concurrency and runs backfill jobs sequentially, the total time is 30 × 10 = 300 minutes — set TTL accordingly. Expired vault keys cause 401 responses from the proxy; the block will raise an exception, Mage will retry, and the load_vault_key block's pre-flight check will return the existing key from the database if the billing period is already recorded. Issue vault keys with generous TTLs or implement the pre-flight database lookup to handle expiry transparently.

How do I handle Stripe card declines in a Mage block without triggering block retry?

Catch stripe.error.CardError and stripe.error.InvalidRequestError and return them as fields in the block's output dict rather than re-raising. When a Mage block raises an exception, Mage treats it as a retriable failure and re-executes the block up to the retry limit. A card decline is not retriable — the card will still be declined on the next attempt — so re-raising creates wasted retry cycles and unnecessary latency before the run reaches a failed terminal state. Return a structured result dict with a status field ("card_declined", "invalid_request") and a charge_id of None. Your downstream data_exporter block can inspect the status field and route declined customers into a dunning notification pipeline rather than treating them as errors.