feat: implement organization-based billing and user management

[jfrench9]

Summary

This PR implements a comprehensive organization-based billing and user management system, transitioning from individual user billing to organization-level billing structures. The changes introduce multi-tenant capabilities with proper organization hierarchy and billing controls.

Key Accomplishments

🏢 Organization Management

• New Organization Model: Added complete organization entity with hierarchical user management
• Organization Limits: Implemented org-level resource limits and quota management
• Member Management: Added organization member invitation, role management, and access controls
• Organization API: New REST endpoints for organization CRUD operations, member management, and usage tracking

💳 Billing System Overhaul

• Org-Level Billing: Migrated from user-based to organization-based billing and subscriptions
• Enhanced Invoicing: Improved invoice models with better organization association and audit trails
• Subscription Management: Updated subscription handling to support organization-level plans
• Payment Provider Integration: Restructured payment provider interfaces for organization billing

🔄 Data Model Migration

• Database Schema: Major migration consolidating billing tables and adding organization structures
• Model Refactoring: Renamed and restructured multiple models for consistency:
  • user_graph → graph_user
  • graph_usage_tracking → graph_usage
  • tier_config → graph_tier
• Legacy Cleanup: Removed deprecated user limits and billing tables

🛡️ Access Control Updates

• Organization-Based Auth: Updated authentication middleware to support org-level permissions
• Graph Access Control: Modified graph operations to respect organization boundaries
• Resource Isolation: Enhanced resource isolation between organizations

📊 Usage & Analytics

• Organization Usage Tracking: New usage collection and reporting at organization level
• Credit Management: Updated credit allocation system for organization-based distribution
• Storage Analytics: Enhanced storage breakdown reporting for organizational insights

Breaking Changes

• API Structure: User-centric endpoints now require organization context
• Authentication Flow: Auth tokens now include organization scope
• Billing Endpoints: Legacy user billing endpoints deprecated in favor of organization billing
• Database Schema: Significant changes to billing and user relationship tables

Testing Notes

• All existing user workflows should be tested with organization context
• Billing integration tests need org-level validation
• Member invitation and role management flows require verification
• Graph access controls should be validated across organization boundaries
• Migration scripts should be tested against production-like datasets

Infrastructure Considerations

• Database migration requires careful planning due to extensive schema changes
• Payment provider configurations need updates for organization-level billing
• Existing user data requires migration to organization structures
• API clients will need updates to accommodate new organization-scoped endpoints
• Monitoring and logging should account for organization-level metrics

---

🤖 Generated with Claude Code

Branch Info:

• Source: feature/org-billing
• Target: main
• Type: feature

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

[jfrench9]

@claude please review this PR

[claude]

Claude finished @jfrench9's task —— View job

---

PR Review Complete ✅

I've completed a comprehensive review of this organization-based billing implementation. This is a significant architectural change with 100 files modified. Here's my detailed feedback:

---

Critical Issues 🚨

1. Data Migration Risk: org_id is Nullable on Graphs

Location: alembic/versions/e40df94401d2_org_billing_subscriptions_invoices.py:340

The migration adds org_id to the graphs table as nullable:

op.add_column("graphs", sa.Column("org_id", sa.String(), nullable=True))

Problem: This creates a data integrity issue. All existing graphs will have org_id=NULL after migration, but there's no data migration step to populate this field. This means:

• Existing graphs won't be associated with any organization
• Billing queries that filter by org_id will miss these graphs
• Users may lose access to their existing graphs

Recommendation: Add a data migration step that:

1. For each existing graph, look up the owner via user_graphs table
2. Get that user's organization from org_users
3. Set the graph.org_id to match
4. Make org_id non-nullable after backfill (or keep nullable only for shared repositories)

---

2. Webhook Security: Missing Signature Verification Logging

Location: robosystems/routers/admin/webhooks.py:50-54

try:
    event = provider.verify_webhook(payload, signature)
except ValueError as e:
    logger.error(f"Invalid webhook signature: {e}")
    raise HTTPException(status_code=400, detail="Invalid webhook signature")

Problem: Failed webhook signature attempts are logged but not tracked in the security audit system. This could hide webhook replay attacks or misconfigured webhooks.

Recommendation: Add security event logging:

except ValueError as e:
    SecurityAuditLogger.log_security_event(
        event_type=SecurityEventType.AUTHORIZATION_DENIED,
        ip_address=client_ip,
        endpoint="/admin/v1/webhooks/stripe",
        details={"reason": "invalid_webhook_signature", "error": str(e)},
        risk_level="high",
    )

---

3. Member Invitation Creates Inactive Users with Predictable Temp Passwords

Location: robosystems/routers/orgs/members.py:128-136

invited_user = User(
    email=request.email,
    name=request.email.split("@")[0],
    password_hash=f"pending_{secrets.token_urlsafe(32)}",  # Temporary password
    is_active=False,  # Inactive until they set password
)

Problems:

1. The password isn't hashed - it's a plaintext string starting with "pending_". Should use hash_password().
2. Invited users are created as inactive, but there's no email verification flow implemented (TODO comment at line 152).
3. No cleanup mechanism for abandoned invitations (inactive users that never activate).

Recommendations:

1. Hash the temporary password: password_hash=hash_password(f"pending_{secrets.token_urlsafe(32)}")
2. Implement the invitation email system with a secure token
3. Add a periodic cleanup job for abandoned invitations (e.g., delete inactive invited users after 7 days)

---

4. Graph Provisioning Task Error Handling Has Race Condition

Location: robosystems/tasks/graph_operations/provision_graph.py:239-248

try:
    if "subscription" not in locals():
        subscription = (
            session.query(BillingSubscription)
            .filter(BillingSubscription.id == subscription_id)
            .first()
        )

Problem: The error handler attempts to check if subscription variable exists in locals(), but this is unreliable and could fail if the exception occurred during subscription fetch. The session might also be in a bad state.

Recommendation: Restructure error handling:

except Exception as e:
    logger.error(f"Graph provisioning failed: {str(e)}", exc_info=True)
    
    try:
        # Always fetch fresh subscription in exception handler
        session.rollback()  # Clean up first
        subscription = session.query(BillingSubscription).filter(
            BillingSubscription.id == subscription_id
        ).first()
        
        if subscription:
            subscription.status = "failed"
            # ... rest of error handling
    except Exception as update_error:
        logger.error(f"Failed to update subscription status: {update_error}")

---

High Priority Issues ⚠️

5. Checkout Flow: Only First Org Used

Location: robosystems/routers/billing/checkout.py:69-77

user_orgs = OrgUser.get_user_orgs(current_user.id, db)
if not user_orgs:
    raise HTTPException(...)

membership = user_orgs[0]  # Always uses first org!
org_id = membership.org_id

Problem: If a user belongs to multiple organizations, this always uses the first one. There's no way for the user to specify which organization they're creating resources for.

Recommendation: Add org_id parameter to CreateCheckoutRequest model and validate the user is a member:

if request.org_id:
    membership = OrgUser.get_by_org_and_user(request.org_id, current_user.id, db)
    if not membership:
        raise HTTPException(403, "Not a member of specified organization")
else:
    # Fallback to default org
    membership = user_orgs[0]

---

6. Registration Creates "Phantom Org" Without User Awareness

Location: robosystems/routers/auth/register.py:277-282

org = Org.create_phantom_org_for_user(
    user_id=user.id,
    user_name=sanitized_name,
    session=session,
)
logger.info(f"Created personal org {org.id} for user {sanitized_email}")

Problem: Every user registration silently creates a "phantom" personal organization. This is a significant architectural decision that:

• Isn't documented in the API response
• May confuse users who think they're individual users, not organizations
• Creates hidden database records users don't know about

Recommendation:

1. Document this behavior in the API documentation
2. Include org information in the registration response
3. Consider making it explicit: "Your Personal Workspace" in the UI

---

7. Missing Transaction Boundaries in Webhook Handlers

Location: robosystems/routers/admin/webhooks.py:109-172

The webhook handler functions (handle_checkout_completed, handle_payment_succeeded, etc.) make multiple database changes but don't use explicit transaction boundaries. If an exception occurs mid-processing, the database could be left in an inconsistent state.

Recommendation: Wrap each handler in a transaction:

async def handle_checkout_completed(session_data: dict, db: Session):
    try:
        # All database operations
        db.commit()
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to process checkout completed: {e}")
        raise

---

8. Graph Tier Configuration Not Validated Against Billing Plan

Location: robosystems/tasks/graph_operations/provision_graph.py:82

tier = graph_config.get("tier", "kuzu-standard")

Problem: The tier is taken directly from the webhook metadata without validating it matches the subscription plan. A malicious or buggy client could request a higher tier than they paid for.

Recommendation: Validate tier against subscription:

tier = subscription.plan_name  # Use tier from paid subscription
if graph_config.get("tier") != tier:
    logger.warning(f"Tier mismatch: requested {graph_config.get('tier')}, using paid tier {tier}")

---

Medium Priority Issues 📋

9. Audit Log Missing Foreign Key Validation

Location: robosystems/models/billing/audit_log.py

The audit log has foreign keys to org_id, subscription_id, invoice_id, and actor_user_id, but these are all nullable and there's no validation that at least one is provided. This could result in audit logs that can't be traced back to any entity.

---

10. Invoice Number Generation Not Atomic

Location: robosystems/models/billing/invoice.py:136-144

count = (
    session.query(cls)
    .filter(cls.invoice_number.like(f"INV-{year}-{month:02d}-%"))
    .count()
    + 1
)
return f"INV-{year}-{month:02d}-{count:04d}"

Problem: If two invoices are created concurrently in the same month, they could get the same invoice number. The unique constraint will catch this, but it will cause a confusing error.

Recommendation: Use PostgreSQL sequences or add retry logic with exponential backoff.

---

11. Organization Deletion Would Cascade Delete Everything

Location: robosystems/models/iam/org.py:42-43

users = relationship("OrgUser", back_populates="org", cascade="all, delete-orphan")
graphs = relationship("Graph", back_populates="org", cascade="all, delete-orphan")

Problem: Deleting an organization would cascade delete all graphs and memberships. While this might be intentional, there's no soft-delete mechanism or safety check. This could lead to accidental data loss.

Recommendation:

1. Add a deleted_at field for soft deletes
2. Require explicit confirmation for org deletion
3. Add a grace period before hard deletion

---

12. No Rate Limiting on Webhook Endpoint

Location: robosystems/routers/admin/webhooks.py:16-39

The webhook endpoint doesn't have rate limiting. While Stripe has their own retry logic, a misconfigured webhook or attack could overwhelm the system.

Recommendation: Add webhook-specific rate limiting (separate from API rate limits).

---

Code Quality Issues 🔧

13. Inconsistent Error Messages in Org Members Endpoint

Location: robosystems/routers/orgs/members.py:118-124

The error message when trying to invite an existing user is very verbose (4 lines). While helpful, it's inconsistent with other error messages in the codebase which are typically 1-2 lines.

---

14. Type Ignore Comments in Webhook Handler

Location: `robosystems/routers/admin/webhooks.py:326, 364, 420, etc.**

Multiple type: ignore comments throughout the webhook handlers indicate the type system isn't properly modeling the data structures. This could hide bugs.

Recommendation: Define proper TypedDict classes for Stripe event structures.

---

15. Hardcoded Retry Values in Provision Task

Location: robosystems/tasks/graph_operations/provision_graph.py:21-24

autoretry_for=(ConnectionError, TimeoutError, OperationalError),
retry_kwargs={"max_retries": 3, "countdown": 60},
retry_backoff=True,
retry_backoff_max=600,

These should be configuration values, not hardcoded.

---

Positive Observations ✨

1. Excellent Security Audit Logging: The registration flow has comprehensive security event logging (robosystems/routers/auth/register.py).

2. Good Separation of Concerns: Billing models are properly separated from IAM models with clear boundaries.

3. Idempotent Webhook Processing: The webhook handler correctly checks for duplicate events using BillingAuditLog.is_webhook_processed().

4. Strong Access Control: Organization member management has proper role-based authorization with checks for owner/admin permissions.

5. Comprehensive Documentation: API endpoints have detailed docstrings explaining the flow and requirements.

6. Proper Index Usage: Database migrations include appropriate indexes for query performance (billing_invoices, billing_subscriptions, etc.).

---

Testing Recommendations 🧪

Before merging, ensure these scenarios are tested:

1. Migration Testing: Test the migration on a copy of production data to ensure existing graphs get properly associated with orgs.

2. Webhook Replay: Test webhook idempotency by sending the same webhook event multiple times.

3. Concurrent Checkout: Test multiple users from the same org trying to checkout simultaneously.

4. Failed Provisioning: Test the error recovery path when graph provisioning fails mid-process.

5. Multi-Org Users: Test users who belong to multiple organizations to ensure they can manage resources for the correct org.

6. Abandoned Invitations: Test the invitation flow with users who never accept.

---

Overall Assessment

This is a substantial and well-architected change that properly separates billing concerns and establishes a solid foundation for organization-based multi-tenancy. The code quality is generally high, with good documentation and security considerations.

However, the data migration risk (#1) is critical and must be addressed before merging to production. The other critical issues (#2-4) should also be resolved to avoid security and data integrity problems.

The high-priority issues (#5-8) should be addressed soon after merge if not before, as they affect the user experience and security posture.

Recommendation: 🟡 Approve with required changes

• Fix critical issue Inital open source release of robosystems #1 (data migration) before merging
• Consider addressing critical issues Fix MCP tools query results to include units and improve dimensions #2-4 before or immediately after merge
• Create follow-up tickets for high and medium priority issues

---

[chatgpt-codex-connector]

💡 Codex Review

robosystems/robosystems/middleware/billing/enforcement.py

Lines 14 to 32 in 56964fc

	def check_can_provision_graph(
	user_id: str,
	requested_tier: GraphTier,
	session: Session,
	) -> Tuple[bool, Optional[str]]:
	"""Check if a user can provision a new graph.
	Args:
	user_id: The user ID
	requested_tier: The requested graph tier
	session: Database session
	Returns:
	Tuple of (can_provision, error_message)
	"""
	billing_customer = BillingCustomer.get_or_create(user_id, session)
	can_provision, error_message = billing_customer.can_provision_resources(
	environment=env.ENVIRONMENT, billing_enabled=env.BILLING_ENABLED

Use org_id when loading billing customer for provisioning checks

After moving billing customers to be keyed by organization, check_can_provision_graph still calls BillingCustomer.get_or_create(user_id, session) with the caller’s user id. The model now declares org_id as the foreign key, so this call attempts to create a billing customer row for a non‑existent org and will raise an integrity error (or associate provisioning decisions with an invalid org) as soon as a user without a matching org id invokes the middleware. To enforce org‑level limits correctly, the function needs to resolve the user’s organization and pass the org id into get_or_create.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[jfrench9]

@claude can you review the latest changes? A graph may not have an org id because a shared repository is not owned by an org

[claude]

Claude finished @jfrench9's task —— View job

---

Updated Review - Latest Changes ✅

I've reviewed the latest commits (51622ed, 18e92e7, 56964fc) addressing several critical issues from my previous review. Great progress!

---

Summary of Changes Since Last Review

The latest commit (51622ed) made significant improvements:

1. ✅ Fixed Critical Issue Fix MCP tools query results to include units and improve dimensions #2 - Billing enforcement now uses org_id instead of user_id
2. ✅ Fixed High Priority Issue Remove unused Account and Disclosure transaction nodes from graph schema #5 - Checkout flow now retrieves user's organization properly
3. ✅ Fixed High Priority Issue Fix worker monitor Valkey connection with Secrets Manager auth #6 - Registration flow improved with better naming and documentation
4. ✅ Addressed Critical Issue Fix MCP tools query results to include units and improve dimensions #2 - Added security audit logging for webhook signature failures
5. ✅ Improved Critical Issue chore: improve documentation and suppress build output #4 - Better error handling in graph provisioning task with proper rollback

---

Re-evaluation: Nullable org_id for Shared Repositories ✅

Status: RESOLVED - Your clarification is correct and the implementation is sound.

Why This Is Safe

Graph Model Documentation (robosystems/models/iam/graph.py:65-69):

# Ownership - graph is owned by an organization
# Nullable for shared repositories which are system-wide
org_id = Column(
  String, ForeignKey("orgs.id"), nullable=True
)  # Organization that owns and pays for this graph (None for shared repositories)

Repository Creation (robosystems/models/iam/graph.py:428):

repository = cls.create(
  graph_id=graph_id,
  org_id=None,  # Shared repositories are not owned by any org
  ...
)

Proper Safeguards in Place

1. Clear Distinction: is_repository flag distinguishes shared repos from org-owned graphs
2. Access Control: Shared repositories use UserRepository table (subscription-based), not GraphUser (org-based)
3. Query Safety: All org-level queries properly filter by org_id, which excludes shared repositories (null != any org_id)
4. Billing Logic: Shared repositories are not included in org billing calculations

Example Query Patterns (All Safe)

• db.query(Graph).filter(Graph.org_id == org_id) - excludes shared repos ✓
• Graph.get_all_repositories() - explicitly queries is_repository=True ✓
• Graph.get_repository_by_type() - uses is_repository filter ✓

Previous concern #1 is RESOLVED - No data migration needed. The nullable org_id is intentional and properly implemented.

---

Remaining Issues to Address

🔴 Critical: Clear-Text Logging of Sensitive Data

GitHub Advanced Security flagged this twice in robosystems/operations/graph/subscription_service.py:

Line 200-204 (creates Stripe customer):

logger.info(
  f"Found existing billing customer for user {user_id}",
  extra={
    "stripe_customer_id": f"{customer.stripe_customer_id[:8]}***{customer.stripe_customer_id[-4:]}"
    if customer.stripe_customer_id
    else None,
    "plan_name": plan_name,
  },
)

Line 236-242 (creates Stripe subscription):

logger.info(
  f"Created Stripe subscription {redacted_stripe_id} for graph {graph_id}",
  extra={
    "user_id": user_id,
    "subscription_id": subscription.id,
    "stripe_subscription_id": redacted_stripe_id,  # ← Already redacted in string, but...
  },
)

Problem: While the IDs are redacted in the log message, the security scanner may be flagging that customer.stripe_customer_id and other sensitive fields are being accessed. This is likely a false positive since you're already redacting them properly.

Recommendation:

1. If these are false positives, mark them as such in GitHub Security
2. If you want to be extra safe, ensure redaction happens before the logger call:

   redacted_customer_id = f"{customer.stripe_customer_id[:8]}***{customer.stripe_customer_id[-4:]}" if customer.stripe_customer_id else None
   logger.info(
     f"Found existing billing customer for user {user_id}",
     extra={"stripe_customer_id": redacted_customer_id, "plan_name": plan_name}
   )

---

⚠️ Medium Priority Issues Still Outstanding

From my previous review, these medium-priority issues remain:

Issue #9: Audit Log Missing Foreign Key Validation

• Location: robosystems/models/billing/audit_log.py
• All foreign keys (org_id, subscription_id, invoice_id, actor_user_id) are nullable
• No validation that at least one is provided
• Could result in untraceable audit logs

Issue #10: Invoice Number Generation Not Atomic

• Location: robosystems/models/billing/invoice.py:136-144
• Concurrent invoice creation could get duplicate numbers
• Unique constraint catches it but causes confusing errors
• Recommendation: Add retry logic or use PostgreSQL sequences

Issue #11: Organization Deletion Cascade

• Location: robosystems/models/iam/org.py:42-43
• Deleting an org cascades to all graphs and memberships
• No soft-delete or safety mechanism
• Consider adding deleted_at field and grace period

Issue #12: No Rate Limiting on Webhook Endpoint

• Location: robosystems/routers/admin/webhooks.py
• Webhook endpoint lacks rate limiting
• Misconfigured webhook could overwhelm system
• Recommendation: Add webhook-specific rate limits

Issue #13-15: Code Quality Issues

• Inconsistent error messages
• Multiple type: ignore comments in webhook handlers
• Hardcoded retry values in provision task

---

What's Been Fixed ✨

1. Billing Enforcement Now Uses Organization ✅

File: robosystems/middleware/billing/enforcement.py:30-41

# Get user's organization - billing is org-level, not user-level
org_user = session.query(OrgUser).filter(OrgUser.user_id == user_id).first()

if not org_user:
  logger.error(f"User {user_id} has no organization - cannot check billing")
  return (False, "User is not a member of any organization")

# Get or create billing customer for the user's organization
billing_customer = BillingCustomer.get_or_create(org_user.org_id, session)

Previously: Used BillingCustomer.get_or_create(user_id, session)
Now: Correctly uses organization ID for billing checks

---

2. Webhook Security Audit Logging ✅

File: robosystems/routers/admin/webhooks.py:65-86

try:
  event = provider.verify_webhook(payload, signature)
except ValueError as e:
  logger.error(f"Invalid webhook signature: {e}")
  
  # Log security event for failed signature verification
  client_ip = request.client.host if request.client else "unknown"
  SecurityAuditLogger.log_security_event(
    event_type=SecurityEventType.AUTHORIZATION_DENIED,
    user_id=None,
    ip_address=client_ip,
    endpoint="/admin/v1/webhooks/stripe",
    details={"reason": "invalid_webhook_signature", "error": str(e)},
    risk_level="high",
  )

Previously: Only logged errors
Now: Full security audit trail for webhook auth failures

---

3. Transaction Boundaries in Webhook Handlers ✅

File: robosystems/routers/admin/webhooks.py:206-213

except Exception as e:
  db.rollback()
  logger.error(
    f"Failed to process checkout completed webhook: {e}",
    extra={"session_id": session_data.get("id")},
    exc_info=True,
  )
  raise

Previously: No explicit transaction management
Now: Proper rollback on errors, commit on success

---

4. Improved Graph Provisioning Error Handling ✅

File: robosystems/tasks/graph_operations/provision_graph.py:241-278

try:
  session.rollback()
  
  subscription = (
    session.query(BillingSubscription)
    .filter(BillingSubscription.id == subscription_id)
    .first()
  )
  
  if subscription:
    subscription.status = "failed"
    # ... error metadata handling
    session.commit()

Previously: Used unreliable locals() check
Now: Always fetches fresh subscription after rollback, proper nested error handling

---

5. Personal Organization Creation Documented ✅

File: robosystems/routers/auth/register.py:285-295

# Create personal organization for the user
# RoboSystems is org-centric: all resources belong to orgs, not users
org = Org.create_personal_org_for_user(
  user_id=user.id,
  user_name=sanitized_name,
  session=session,
)

Previously: Called create_phantom_org_for_user with no explanation
Now: Clear naming, inline documentation, and org info in response

AuthResponse Model now includes:

org: dict[str, object] | None = Field(
  default=None,
  description="Organization information (personal org created automatically on registration)",
)

---

Testing Recommendations 🧪

Before merging, test these specific scenarios:

1. Shared Repository Access

   • Create a shared repository (e.g., SEC data)
   • Verify it has org_id=None and is_repository=True
   • Confirm users can access it via UserRepository subscription
   • Ensure it's excluded from org billing/usage queries

2. Organization Billing

   • Create a user and verify personal org is created
   • Provision a graph and confirm billing uses org_id
   • Verify multiple users in same org share billing

3. Webhook Error Recovery

   • Test webhook signature failures get logged to security audit
   • Verify database rollback on webhook processing errors
   • Confirm idempotency (duplicate webhooks handled correctly)

4. Graph Provisioning Failures

   • Trigger a provisioning failure (e.g., invalid tier)
   • Verify subscription status changes to "failed"
   • Confirm error metadata is stored
   • Check rollback leaves database consistent

---

Overall Assessment

Status: 🟢 Significant Improvement - Recommend Approval with Minor Follow-ups

What Changed

• 5 critical/high-priority issues from previous review have been addressed
• Security audit logging properly implemented
• Organization-based billing is now correct
• Error handling is much more robust
• Documentation and naming improved

Remaining Work

• 1 critical issue (likely false positive security scan)
• 5 medium-priority issues that can be addressed post-merge
• Code quality improvements (non-blocking)

Recommendation

✅ Approve for merge with follow-up tasks for:

1. Review security scanning alerts (likely false positives)
2. Add validation to audit log model
3. Consider atomic invoice number generation
4. Add soft-delete for organizations
5. Implement webhook rate limiting

The nullable org_id implementation is correct and well-documented. Shared repositories are properly isolated from organization billing and access control.

---

Branch: feature/org-billing