feat: Add immediate graph deprovisioning with subscription cancellation

[jfrench9]

Summary

Implements immediate graph deprovisioning by adding subscription cancellation logic that triggers resource cleanup when a graph is deleted. This replaces or augments the previous end-of-period cancellation behavior with an immediate cancellation flow, ensuring that billing and infrastructure are promptly reconciled when users delete a graph.

Key Accomplishments

Billing & Subscription Cancellation

• Added a new CancelSubscriptionRequest API model to formally represent cancellation requests
• Extended the subscription billing router with enhanced cancellation logic supporting immediate cancellation semantics
• Updated the core subscription model to support new cancellation states/fields (backed by a new database migration)
• Introduced a dedicated graph subscriptions router (routers/graphs/subscriptions.py) to manage the relationship between graph lifecycle and subscription state

Graph Delete Operation

• Enhanced the delete-graph operation in routers/graphs/operations.py with comprehensive logic that coordinates subscription cancellation as part of graph deletion
• Added detailed operation documentation/models in models/api/graphs/operations.py

Infrastructure & Lifecycle Sensors

• Updated the Dagster graph lifecycle sensor to account for immediate deprovisioning scenarios, ensuring background jobs respond correctly when a graph is cancelled mid-cycle

Observability

• Extended the audit logger with new event types to capture subscription cancellation and graph deletion actions

Database Migration

• Added migration ac32cd15e221_immediate_cancel to support the new immediate cancellation fields on the subscription model

Breaking Changes

• Subscription cancellation behavior now defaults to immediate rather than end-of-period for graph deletions. Consumers relying on the previous end-of-period grace period should verify integration behavior.
• New fields on the subscription model require running the database migration before deploying this change.

Testing

• New test suite: tests/routers/graphs/test_delete_graph_op.py — comprehensive coverage of the delete-graph operation including edge cases around subscription state
• Expanded tests: tests/routers/billing/test_subscriptions.py (+339 lines) and tests/routers/graphs/test_subscriptions_router.py (+261 lines) significantly expanded to cover new cancellation flows
• Sensor tests: Added tests in tests/dagster/test_graph_lifecycle.py for the updated lifecycle sensor behavior
• Model tests: Updated tests/models/core/billing/test_subscription.py for new subscription model fields

Infrastructure Considerations

• A database migration must be applied prior to deployment to add the immediate cancellation columns
• The Dagster sensor changes should be deployed in coordination with the API changes to avoid inconsistent state handling
• Verify that any external billing provider webhooks or integrations are compatible with the new immediate cancellation flow

---

🤖 Generated with Claude Code

Branch Info:

• Source: feature/graph-immediate-deprovision
• Target: main
• Type: feature

Co-Authored-By: Claude noreply@anthropic.com

[jfrench9]

@claude please review this PR

[claude]

Claude finished @jfrench9's task in 2m 54s —— View job

---

PR Review: feat: Add immediate graph deprovisioning with subscription cancellation

• Read all changed files
• Review migration and model changes
• Review router and business logic changes
• Review test coverage
• Post final review

---

Overall this is well-structured work — the two-mode design (immediate vs at_period_end), the confirm-token guard, and the sensor branching on cancellation_type are all sound. I have a few bugs and one architecture note to flag before merging.

---

Bugs

1. Wrong SecurityEventType on Stripe failure [routers/graphs/operations.py:396]

SecurityAuditLogger.log_security_event(
    event_type=SecurityEventType.AUTH_TOKEN_INVALID,  # ← wrong type
    ...
    "action": "stripe_cancel_failed_during_delete_graph",

AUTH_TOKEN_INVALID is an auth event. The right type here is SecurityEventType.OPERATION_FAILED (it already exists in the enum). As-is, a Stripe cancellation failure would show up in auth dashboards/alerting.

Fix this →

---

2. cancel(immediate=False) silently sets ends_at = None when current_period_end is null [models/core/billing/subscription.py:312]

else:
    self.ends_at = self.current_period_end  # None if never activated

A subscription in pending, provisioning, paused, or past_due state may never have had current_period_end set. Calling cancel(immediate=False) on such a subscription sets ends_at = None, which the expired_graph_subscription_sensor filters out (ends_at.isnot(None)). Result: the graph is never suspended and its infrastructure is never cleaned up.

delete_graph_op only guards against canceled/canceling/upgrading — it will happily accept a pending subscription with at_period_end=True.

Suggested fix: add an explicit guard in cancel():

if not immediate:
    if self.current_period_end is None:
        raise ValueError(
            f"Cannot defer cancellation for subscription {self.id}: "
            "current_period_end is not set (subscription may not have been activated)."
        )
    self.ends_at = self.current_period_end

…and/or add a pre-flight check in delete_graph_op requiring the subscription be active before proceeding with at_period_end=True.

Fix this →

---

Design Notes

3. _get_plan_display_name duplicated across two routers

Identical function defined at routers/billing/subscriptions.py:28 and routers/graphs/subscriptions.py:42. The second could simply import from the first (or both could import from a shared utility). Stale copies tend to diverge.

---

4. cancel_subscription endpoint is now unreachable for all current resource types [routers/billing/subscriptions.py:271–289]

Both "graph" and "repository" resource types now raise 400 pointing callers to their canonical paths. That's correct routing, but it means POST /billing/subscriptions/{org_id}/subscription/{subscription_id}/cancel is currently a dead endpoint for any existing resource. If it's being reserved for future platform-level add-ons, a comment to that effect would help; otherwise it could be removed to avoid confusion in the OpenAPI docs.

---

5. Body parameter inconsistency in cancel_repository_subscription [routers/graphs/subscriptions.py:592]

body: CancelSubscriptionRequest = ...,          # graphs/subscriptions.py
body: CancelSubscriptionRequest = Body(default_factory=CancelSubscriptionRequest),  # billing/subscriptions.py

The = ... form makes the body technically required at the FastAPI level. Since all fields in CancelSubscriptionRequest have defaults, a caller omitting the body entirely (vs sending {}) may get a validation error. The billing router's Body(default_factory=...) form is more permissive and intentional — the graphs router should match it.

---

6. logging.getLogger instead of get_logger in graphs/subscriptions.py [routers/graphs/subscriptions.py:30]

import logging
logger = logging.getLogger(__name__)  # ← not the project's get_logger

Every other router in the codebase uses from ...logger import get_logger / logger = get_logger(__name__). Inconsistent logger initialization may mean structured fields or handlers from the project logger aren't applied here.

---

Test Coverage

The new test suite in test_delete_graph_op.py is thorough for the router-level contract (shared repo gate, confirm token, admin role, Stripe branching). Two gaps worth noting:

• No test for cancel(immediate=False) with current_period_end=None in tests/models/core/billing/test_subscription.py — this is the scenario from bug Fix MCP tools query results to include units and improve dimensions #2.
• test_query_includes_immediate_cancellation_clause verifies the SQLAlchemy filter expression is called but doesn't assert the actual OR predicate shape; the string-compile approach only weakly validates it. This is acceptable at unit-test granularity, but an integration test that exercises the full query path would catch query regressions that mocks miss.

---

Summary

Category	Finding
Bug	Wrong SecurityEventType on Stripe failure (AUTH_TOKEN_INVALID → OPERATION_FAILED)
Bug	cancel(immediate=False) sets ends_at=None for unactivated subscriptions → infrastructure leak
Design	Dead endpoint for all current resource types in billing cancel router
Design	_get_plan_display_name duplicated across two routers
Nit	logging.getLogger vs project get_logger
Nit	Body parameter inconsistency (= ... vs Body(default_factory=...))

The two bugs (#1 and #2) should be fixed before merging; the rest are improvements that can be tracked separately.

[jfrench9]

@claude have we addressed all issues worth fixing?

[claude]

Claude finished @jfrench9's task in 1m 28s —— View job

---

Review issue status — all bugs fixed, one design note outstanding

Here's the current state of each item from the previous review:

---

Bug #1 — Wrong SecurityEventType on Stripe failure ✅ Fixed

routers/graphs/operations.py:441 now correctly uses SecurityEventType.OPERATION_FAILED. The Stripe failure no longer pollutes auth dashboards.

---

Bug #2 — cancel(immediate=False) null current_period_end ✅ Fixed

Two-layer fix landed:

• models/core/billing/subscription.py:311–316 — cancel() now raises ValueError if current_period_end is None when immediate=False
• routers/graphs/operations.py:399–407 — delete_graph_op pre-flight check rejects at_period_end=True if subscription is not active, with a clear error message pointing callers to the immediate path

---

Nit #5 — Body parameter inconsistency ✅ Fixed

cancel_repository_subscription at graphs/subscriptions.py:584 now uses Body(default_factory=CancelSubscriptionRequest), consistent with the billing router.

---

Nit #6 — logging.getLogger vs get_logger ✅ Fixed

graphs/subscriptions.py:14,29 now imports and uses get_logger(__name__), matching the rest of the codebase.

---

Design #3 — _get_plan_display_name duplication ✅ Substantially fixed

Logic is centralized in BillingConfig.get_plan_display_name (config/billing/core.py:201). Both routers now have a thin one-liner wrapper that delegates to it (billing/subscriptions.py:28–32, graphs/subscriptions.py:41–43). The duplication risk is gone since divergence would only affect the stub, not the logic. The two wrappers could be removed and callers could import BillingConfig directly, but this is a minor cleanup.

---

Design #4 — cancel_subscription billing endpoint unreachable ⚠️ Still present (by design)

routers/billing/subscriptions.py:260–276 still redirects both "graph" (→ delete-graph) and "repository" (→ subscriptions/cancel) resource types with 400s. The endpoint remains reachable in principle but has no current resource type that flows through it. This appears intentional — it's a routing layer reserved for future platform-level subscription types. No code comment explains this, which could confuse readers of the OpenAPI docs. Not a blocking issue, but worth a # reserved for future add-on subscription types comment if this is the intent.

---

Summary

Issue	Status
Bug: Wrong SecurityEventType	✅ Fixed
Bug: ends_at=None infrastructure leak	✅ Fixed
Nit: Body(...) inconsistency	✅ Fixed
Nit: logging.getLogger vs get_logger	✅ Fixed
Design: _get_plan_display_name duplication	✅ Substantially fixed (logic centralized)
Design: Dead billing cancel endpoint	⚠️ Still present — intentional redirect, no blocking impact

All blocking items are resolved. The remaining design note (#4) is not a blocker — a short comment in the billing router would close it, but it can also be tracked separately. This is ready to merge.