# Billing State Machine Documentation

## Table of Contents
1. [Overview](#overview)
2. [Customer States](#customer-states)
3. [State Diagram](#state-diagram)
4. [Stripe Events](#stripe-events)
5. [State Transitions](#state-transitions)
6. [Decision Table](#decision-table)
7. [Sequence Diagrams](#sequence-diagrams)
8. [Edge Cases and Error Handling](#edge-cases-and-error-handling)
9. [Implementation Reference](#implementation-reference)

## Overview

The billing system implements a state machine that manages customer transitions between free and metered (paid) plans based on Stripe payment events. The system follows a **zero-ops philosophy** where all billing interactions happen through Stripe's hosted services, and we only react to webhooks to maintain state synchronization.

### Key Principles
- **Reactive System**: We never initiate payment flows, only react to Stripe events
- **Event-Driven**: All state changes are triggered by Stripe webhooks
- **Idempotent**: All operations are safe to retry
- **Stateless Service**: All state is persisted in DynamoDB

## Customer States

The system recognizes the following customer states:

| State | Description | API Usage Plan | Stripe Subscription | Payment Method |
|-------|-------------|----------------|---------------------|----------------|
| **FREE_NO_PM** | New customer without payment method | FREE_10K | None | None |
| **FREE_WITH_PM** | Customer with payment method but no active subscription (e.g., after payment failure) | FREE_10K | Canceled/None | Exists |
| **METERED_ACTIVE** | Customer with active metered subscription | METERED | Active | Exists (default) |
| **METERED_PAST_DUE** | Customer with overdue payment (grace period) | METERED → FREE_10K | Active but failing | Exists |

## State Diagram

```mermaid
stateDiagram-v2
    [*] --> FREE_NO_PM: New Customer Created

    FREE_NO_PM --> METERED_ACTIVE: payment_method.attached\n(with default PM)
    FREE_NO_PM --> METERED_ACTIVE: setup_intent.succeeded\n(with default PM)

    METERED_ACTIVE --> FREE_WITH_PM: invoice.payment_failed
    METERED_ACTIVE --> METERED_ACTIVE: customer.subscription.updated\n(maintains PM)
    METERED_ACTIVE --> FREE_NO_PM: payment_method.detached\n(last PM removed)

    FREE_WITH_PM --> METERED_ACTIVE: payment_method.attached\n(new default PM)
    FREE_WITH_PM --> METERED_ACTIVE: setup_intent.succeeded\n(payment retry success)
    FREE_WITH_PM --> METERED_ACTIVE: customer.subscription.created\n(resubscribe)
    FREE_WITH_PM --> FREE_NO_PM: payment_method.detached\n(all PMs removed)

    METERED_PAST_DUE --> FREE_WITH_PM: invoice.payment_failed\n(after grace period)
    METERED_PAST_DUE --> METERED_ACTIVE: invoice.payment_succeeded\n(payment recovered)

    note right of FREE_NO_PM
        Limited to 10K API calls/month
        No usage reporting to Stripe
    end note

    note right of METERED_ACTIVE
        Unlimited API calls
        Hourly usage reporting to Stripe
        Billed monthly based on usage
    end note

    note right of FREE_WITH_PM
        Customer can fix payment
        and auto-promote back
    end note
```

## Stripe Events

### Events That Trigger State Changes

| Event Type | Description | Typical Trigger | Target State Logic |
|------------|-------------|-----------------|-------------------|
| **payment_method.attached** | Payment method added to customer | Customer adds card via Portal | Check for default PM → METERED if exists |
| **setup_intent.succeeded** | Payment setup completed | Customer completes checkout | Check for default PM → METERED if exists |
| **customer.subscription.created** | New subscription created | System creates metered subscription | Check for default PM → METERED if exists |
| **customer.subscription.updated** | Subscription modified | Plan change, status update | Check for default PM → maintain or change |
| **invoice.payment_failed** | Payment attempt failed | Card declined, insufficient funds | Always → FREE |
| **invoice.payment_succeeded** | Payment successful | Retry succeeded, new payment | Check subscription status → METERED |

### Events Currently Ignored

| Event Type | Why Ignored | Potential Future Use |
|------------|-------------|---------------------|
| **payment_method.detached** | Not implemented | Could trigger demotion if last PM |
| **customer.deleted** | Not implemented | Could clean up DynamoDB records |
| **customer.subscription.deleted** | Handled via payment_failed | Could be explicit demotion trigger |

## State Transitions

### Comprehensive State Transition Matrix

```mermaid
graph TB
    subgraph "From FREE_NO_PM"
        A1[payment_method.attached] -->|has_default_pm| B1[METERED_ACTIVE]
        A2[setup_intent.succeeded] -->|has_default_pm| B1
        A3[payment_method.attached] -->|no_default_pm| C1[FREE_NO_PM]
    end

    subgraph "From METERED_ACTIVE"
        D1[invoice.payment_failed] --> E1[FREE_WITH_PM]
        D2[customer.subscription.updated] -->|still has PM| F1[METERED_ACTIVE]
        D3[payment_method.detached] -->|was last PM| G1[FREE_NO_PM]
    end

    subgraph "From FREE_WITH_PM"
        H1[payment_method.attached] -->|becomes default| I1[METERED_ACTIVE]
        H2[setup_intent.succeeded] -->|payment fixed| I1
        H3[customer.subscription.created] -->|resubscribe| I1
        H4[payment_method.detached] -->|all removed| J1[FREE_NO_PM]
    end
```

## Decision Table

### State Transition Decision Logic

| Current State | Event | Condition | Next State | Action | Side Effects |
|---------------|-------|-----------|------------|--------|--------------|
| **FREE_NO_PM** | payment_method.attached | has_default_payment_method() | METERED_ACTIVE | create_subscription() + move_to_metered_plan() | Start usage reporting |
| **FREE_NO_PM** | payment_method.attached | !has_default_payment_method() | FREE_NO_PM | none | Log "no default PM" |
| **FREE_NO_PM** | setup_intent.succeeded | has_default_payment_method() | METERED_ACTIVE | ensure_metered_subscription() | Start usage reporting |
| **METERED_ACTIVE** | invoice.payment_failed | - | FREE_WITH_PM | cancel_subscription() + move_to_free_plan() | Stop usage reporting |
| **METERED_ACTIVE** | customer.subscription.updated | has_default_payment_method() | METERED_ACTIVE | update_subscription_item_id() | Continue reporting |
| **METERED_ACTIVE** | customer.subscription.updated | !has_default_payment_method() | FREE_WITH_PM | move_to_free_plan() | Stop usage reporting |
| **FREE_WITH_PM** | payment_method.attached | becomes_default_payment_method() | METERED_ACTIVE | ensure_metered_subscription() | Resume usage reporting |
| **FREE_WITH_PM** | setup_intent.succeeded | payment_retry_successful() | METERED_ACTIVE | ensure_metered_subscription() | Resume usage reporting |
| **FREE_WITH_PM** | customer.subscription.created | - | METERED_ACTIVE | move_to_metered_plan() | Start usage reporting |

## Sequence Diagrams

### 1. New Customer Adds Payment Method (Promotion Flow)

```mermaid
sequenceDiagram
    autonumber
    participant Customer
    participant StripePortal as Stripe Customer Portal
    participant Stripe
    participant EventBridge
    participant Lambda as Lambda Handler
    participant Service as BillingService
    participant DynamoDB
    participant APIGateway as API Gateway

    Customer->>StripePortal: Add payment method
    StripePortal->>Stripe: Create payment method
    Stripe->>Stripe: Set as default payment method
    Stripe-->>EventBridge: payment_method.attached event

    EventBridge-->>Lambda: Trigger handler
    Lambda->>Service: handle_stripe_event(event)

    Service->>Service: Extract customer_id from event
    Service->>Stripe: has_default_payment_method(customer_id)
    Stripe-->>Service: true

    Service->>Stripe: ensure_metered_subscription(customer_id)
    Stripe-->>Service: subscription_item_id

    Service-->>Lambda: {target_plan: "metered", subscription_item_id: "si_xxx"}

    Lambda->>APIGateway: move_key_to_plan(api_key, "METERED")
    Lambda->>DynamoDB: save(plan="metered", subscription_item_id="si_xxx")

    Note over Customer: Customer now on METERED plan
```

### 2. Payment Failure (Demotion Flow)

```mermaid
sequenceDiagram
    autonumber
    participant Stripe
    participant EventBridge
    participant Lambda as Lambda Handler
    participant Service as BillingService
    participant DynamoDB
    participant APIGateway as API Gateway

    Stripe->>Stripe: Payment attempt fails
    Stripe-->>EventBridge: invoice.payment_failed event

    EventBridge-->>Lambda: Trigger handler
    Lambda->>Service: handle_stripe_event(event)

    Service->>Service: Extract customer_id
    Service->>Service: Determine target_plan = "free"

    Service-->>Lambda: {target_plan: "free", subscription_item_id: null}

    Lambda->>APIGateway: move_key_to_plan(api_key, "FREE_10K")
    Lambda->>DynamoDB: save(plan="free", subscription_item_id=null)

    Note over Stripe: Customer demoted to FREE plan
    Note over Stripe: Customer can fix payment to re-promote
```

### 3. Hourly Usage Reporting Workflow

```mermaid
sequenceDiagram
    autonumber
    participant CloudWatch as CloudWatch Events
    participant Lambda as Usage Reporter Lambda
    participant Service as BillingService
    participant DynamoDB
    participant APIGateway as API Gateway
    participant Stripe
    participant Audit as Audit Trail

    CloudWatch-->>Lambda: Hourly schedule trigger
    Lambda->>Service: reconcile_usage_window(since, until)

    Service->>DynamoDB: scan_metered() - Get all metered customers
    DynamoDB-->>Service: List of Links with plan="metered"

    loop For each metered customer
        Service->>Service: Calculate window_start from last_usage_window_end

        alt Gap detected (window_start > since)
            Service->>Service: Log warning about gap
        end

        Service->>APIGateway: get_usage(usage_plan_id, api_key_id, window_start, until)
        APIGateway-->>Service: usage_count

        alt usage_count > 0
            Service->>Service: Generate idempotency_key
            Service->>Stripe: report_usage(subscription_item_id, usage_count, timestamp, idempotency_key)
            Stripe-->>Service: usage_record

            opt Audit enabled
                Service->>Audit: record_usage_report(report_details)
            end

            Service->>DynamoDB: Update Link with last_usage_window_end=until
        end
    end

    Service-->>Lambda: List of usage reports
    Lambda->>Lambda: Log metrics
```

### 4. Manual Promotion (Admin Override)

```mermaid
sequenceDiagram
    autonumber
    participant Admin
    participant CLI as CLI Tool
    participant Service as BillingService
    participant Stripe
    participant DynamoDB
    participant APIGateway as API Gateway

    Admin->>CLI: ai-billing promote --api-key KEY123
    CLI->>Service: promote(api_key_id)

    Service->>DynamoDB: get_by_api_key(api_key_id)
    DynamoDB-->>Service: Link object

    Service->>Stripe: ensure_metered_subscription(customer_id)
    Stripe-->>Service: subscription_item_id

    Service-->>CLI: {target_plan: "metered", subscription_item_id: "si_xxx"}

    CLI->>APIGateway: move_key_to_plan(api_key_id, "METERED")
    CLI->>DynamoDB: save(plan="metered", subscription_item_id="si_xxx")

    CLI-->>Admin: Success - API key promoted
```

## Edge Cases and Error Handling

### 1. Concurrent State Changes

**Scenario**: Multiple Stripe events arrive simultaneously for the same customer

**Handling**:
- DynamoDB operations are atomic at the item level
- Last write wins for plan transitions
- Idempotency keys prevent duplicate usage reporting

```python
# Implementation pattern
def apply_plan_move_for_customer(customer_id, target_plan):
    for link in repo.get_by_customer(customer_id):
        # Atomic update per API key
        link.plan = target_plan
        repo.save(link)  # DynamoDB PutItem with conditions
```

### 2. Usage Reporting Gaps

**Scenario**: Lambda fails or is delayed, creating gaps in usage windows

**Detection**:
```python
if link.last_usage_window_end and window_start > since:
    gap_seconds = (window_start - since).total_seconds()
    if gap_seconds > 300:  # More than 5 minutes
        log_event("warning", "usage_window_gap_detected", ...)
```

**Recovery**:
- System automatically fills gaps on next run
- Uses last_usage_window_end as new window start
- Maintains continuity in billing

### 3. Payment Method Edge Cases

#### Multiple Payment Methods
- System only checks for `default_payment_method`
- Non-default payment methods don't trigger promotion
- Customer must set default in Stripe Portal

#### Payment Method Attached but Not Default
```python
# Current behavior
if has_default_payment_method(customer_id):
    # Promote to metered
else:
    # Stay on free plan
    log_event("info", "stay_free_no_default_pm")
```

### 4. Subscription Lifecycle Issues

#### Orphaned Subscriptions
**Problem**: Subscription exists but no payment method
**Solution**: `invoice.payment_failed` will trigger demotion

#### Multiple Active Subscriptions
**Problem**: Customer somehow has multiple subscriptions
**Solution**: System cancels all and creates single new one

```python
def ensure_metered_subscription(customer_id):
    existing = self._list_subs(customer_id)
    if existing:
        return existing[0]["items"]["data"][0]["id"]
    # Create new subscription
```

### 5. Race Conditions

#### Promotion During Usage Reporting
**Scenario**: Customer promoted while usage is being reported
**Solution**: Usage reporting uses point-in-time snapshot from DynamoDB scan

#### Demotion During Active API Call
**Scenario**: Customer demoted while making API calls
**Solution**: API Gateway enforces limits immediately after plan move

### 6. Error Recovery Patterns

#### Failed Plan Move
```python
try:
    plan_admin.move_key_to_plan(api_key_id, target_plan_id)
    link.plan = new_plan
    repo.save(link)
except Exception as e:
    # Plan move failed - state remains consistent
    log_event("error", "plan_move_failed", error=str(e))
    # Lambda will retry entire operation
```

#### Failed Usage Report
```python
try:
    stripe.report_usage(subscription_item_id, usage, timestamp, idempotency_key)
    # Only update state after successful report
    link.last_usage_window_end = until
    repo.save(link)
except Exception as e:
    # Don't update state - will retry in next window
    log_event("error", "usage_reporting_failed", error=str(e))
    raise  # Trigger Lambda retry
```

### 7. Monitoring and Alerting

#### Critical Metrics to Monitor
1. **State Transition Failures**
   - Failed promotions/demotions
   - Plan move errors

2. **Usage Reporting Health**
   - Gap detection frequency
   - Report failure rate
   - Idempotency key collisions

3. **State Consistency**
   - Customers with plan != usage_plan_id mismatch
   - Metered customers without subscription_item_id

## Implementation Reference

### Core Service Implementation

The state machine logic is primarily implemented in:

#### `/src/augint_billing_lib/service.py`

```python
def handle_stripe_event(self, evt: dict[str, Any]) -> dict[str, Any]:
    """Main state machine entry point"""

    # Events that can trigger promotion
    if etype in (
        "payment_method.attached",
        "setup_intent.succeeded",
        "customer.subscription.created",
        "customer.subscription.updated"
    ):
        if self.stripe.has_default_payment_method(customer_id):
            # Transition to METERED_ACTIVE
            return {"target_plan": "metered", ...}
        else:
            # Stay in or move to FREE
            return {"target_plan": "free", ...}

    # Events that trigger demotion
    if etype == "invoice.payment_failed":
        # Always transition to FREE_WITH_PM
        return {"target_plan": "free", ...}
```

### Data Model

#### `/src/augint_billing_lib/models.py`

```python
@dataclass
class Link:
    """Core state tracking entity"""
    api_key_id: str                          # Primary key
    stripe_customer_id: str                  # Stripe customer
    plan: Literal["free", "metered"]        # Current state
    usage_plan_id: str                       # API Gateway plan
    metered_subscription_item_id: str | None # For usage reporting
    last_usage_window_end: datetime | None   # Prevent gaps/overlaps
```

### Lambda Integration

#### `/src/augint_billing_lib/bootstrap.py`

```python
def process_event_and_apply_plan_moves(event: dict) -> dict:
    """Lambda handler for Stripe events"""
    service = _get_service()
    result = service.handle_stripe_event(event)

    if result["target_plan"]:
        # Apply the state transition
        moved = service.apply_plan_move_for_customer(
            customer_id=result["stripe_customer_id"],
            target=result["target_plan"],
            ...
        )
    return result
```

### Testing Coverage

The state machine is thoroughly tested in:
- `/tests/unit/test_service.py` - Unit tests for all transitions
- `/tests/integration/test_billing_workflow_integration.py` - Integration tests
- `/tests/e2e/test_billing_workflow.py` - End-to-end workflow tests

## Appendix: Future Enhancements

### Potential State Machine Extensions

1. **Grace Period State (METERED_PAST_DUE)**
   - Allow continued access for X days after payment failure
   - Requires tracking payment_due_date

2. **Trial State (TRIAL_ACTIVE)**
   - Free trial with full metered access
   - Auto-convert to METERED_ACTIVE or FREE_NO_PM

3. **Suspended State**
   - Manual suspension for abuse/compliance
   - Blocks all API access regardless of payment

4. **Enterprise State**
   - Custom pricing/limits
   - Bypass standard state machine

### Implementation Recommendations

1. **Add State Enum**
```python
from enum import Enum

class CustomerState(Enum):
    FREE_NO_PM = "free_no_payment_method"
    FREE_WITH_PM = "free_with_payment_method"
    METERED_ACTIVE = "metered_active"
    METERED_PAST_DUE = "metered_past_due"
```

2. **Explicit Transition Validation**
```python
VALID_TRANSITIONS = {
    (CustomerState.FREE_NO_PM, "payment_method.attached"): CustomerState.METERED_ACTIVE,
    (CustomerState.METERED_ACTIVE, "invoice.payment_failed"): CustomerState.FREE_WITH_PM,
    # ... etc
}
```

3. **Transition Hooks**
```python
def on_transition(from_state, to_state, customer_id):
    """Hook for state transition side effects"""
    if to_state == CustomerState.METERED_ACTIVE:
        start_usage_reporting(customer_id)
    elif from_state == CustomerState.METERED_ACTIVE:
        stop_usage_reporting(customer_id)
```

---

*This documentation reflects the current implementation as of v2.7.2. For the latest updates, see CHANGELOG.md*
