Implement Subscription Plans

src/AccountDO.ts

@@ -1,4 +1,8 @@
-import type { StartupAPIEnv } from './StartupAPIEnv';
+import { DurableObject } from 'cloudflare:workers';
+import { initPlans } from './billing/plansConfig';
+import { Plan } from './billing/Plan';
+import { MockPaymentEngine } from './billing/PaymentEngine';
+import { StartupAPIEnv } from './StartupAPIEnv';
 
 /**
  * A Durable Object representing an Account (Tenant).
@@ -11,11 +15,16 @@ export class AccountDO implements DurableObject {
   state: DurableObjectState;
   env: StartupAPIEnv;
   sql: SqlStorage;
+  paymentEngine: MockPaymentEngine;
 
   constructor(state: DurableObjectState, env: StartupAPIEnv) {
     this.state = state;
     this.env = env;
     this.sql = state.storage.sql;
+    this.paymentEngine = new MockPaymentEngine();
+
+    // Initialize plans
+    initPlans();
 
     // Initialize database schema
     this.sql.exec(`
@@ -48,6 +57,12 @@ export class AccountDO implements DurableObject {
     } else if (path.startsWith('/members/') && method === 'DELETE') {
       const userId = path.replace('/members/', '');
       return this.removeMember(userId);
+    } else if (path === '/billing' && method === 'GET') {
+      return this.getBillingInfo();
+    } else if (path === '/billing/subscribe' && method === 'POST') {
+      return this.subscribe(request);
+    } else if (path === '/billing/cancel' && method === 'POST') {
+      return this.cancelSubscription();
     }
 
     return new Response('Not Found', { status: 404 });
@@ -129,4 +144,100 @@ export class AccountDO implements DurableObject {
 
     return Response.json({ success: true });
   }
+
+  // Billing Implementation
+
+  private getBillingState(): any {
+    const result = this.sql.exec("SELECT value FROM account_info WHERE key = 'billing'");
+    for (const row of result) {
+      // @ts-ignore
+      return JSON.parse(row.value as string);
+    }
+    return {
+      plan_slug: 'free',
+      status: 'active',
+    };
+  }
+
+  private setBillingState(state: any) {
+    this.state.storage.transactionSync(() => {
+        this.sql.exec("INSERT OR REPLACE INTO account_info (key, value) VALUES ('billing', ?)", JSON.stringify(state));
+    });
+  }
+
+  async getBillingInfo(): Promise<Response> {
+    const state = this.getBillingState();
+    const plan = Plan.get(state.plan_slug);
+    return Response.json({
+      state,
+      plan_details: plan,
+    });
+  }
+
+  async subscribe(request: Request): Promise<Response> {
+    const { plan_slug, schedule_idx = 0 } = (await request.json()) as { plan_slug: string; schedule_idx?: number };
+    const plan = Plan.get(plan_slug);
+
+    if (!plan) {
+      return new Response('Plan not found', { status: 400 });
+    }
+
+    const currentState = this.getBillingState();
+
+    // Call hook if changing plans (simplification)
+    if (currentState.plan_slug !== plan_slug) {
+      if (currentState.plan_slug) {
+         const oldPlan = Plan.get(currentState.plan_slug);
+         if (oldPlan?.account_deactivate_hook) {
+             await oldPlan.account_deactivate_hook(this.state.id.toString());
+         }
+      }
+      if (plan.account_activate_hook) {
+          await plan.account_activate_hook(this.state.id.toString());
+      }
+    }
+
+    // Setup recurring payment
+    try {
+      await this.paymentEngine.setupRecurring(this.state.id.toString(), plan_slug, schedule_idx);
+    } catch (e: any) {
+      return new Response(`Payment setup failed: ${e.message}`, { status: 500 });
+    }
+
+    const newState = {
+      ...currentState,
+      plan_slug,
+      status: 'active',
+      schedule_idx,
+      next_billing_date: Date.now() + (plan.schedules[schedule_idx]?.charge_period || 30) * 24 * 60 * 60 * 1000
+    };
+
+    this.setBillingState(newState);
+
+    return Response.json({ success: true, state: newState });
+  }
+
+  async cancelSubscription(): Promise<Response> {
+    const currentState = this.getBillingState();
+    const currentPlan = Plan.get(currentState.plan_slug);
+
+    if (!currentPlan) {
+        return new Response('No active plan', { status: 400 });
+    }
+
+    await this.paymentEngine.cancelRecurring(this.state.id.toString());
+
+    // Downgrade logic (immediate or scheduled - simplification: scheduled if downgrade_to_slug exists)
+    // For this prototype, we'll mark it as canceled and set the next plan if applicable.
+
+    const newState = {
+        ...currentState,
+        status: 'canceled',
+        next_plan_slug: currentPlan.downgrade_to_slug
+    };
+
+    this.setBillingState(newState);
+
+    return Response.json({ success: true, state: newState });
+  }
 }

src/billing/PaymentEngine.ts

@@ -0,0 +1,20 @@
+export abstract class PaymentEngine {
+  abstract charge(accountId: string, amount: number, currency: string): Promise<boolean>;
+  abstract setupRecurring(accountId: string, planSlug: string, scheduleIdx: number): Promise<void>;
+  abstract cancelRecurring(accountId: string): Promise<void>;
+}
+
+export class MockPaymentEngine extends PaymentEngine {
+  async charge(accountId: string, amount: number, currency: string): Promise<boolean> {
+    console.log(`[MockPaymentEngine] Charging ${accountId} ${amount} ${currency}`);
+    return true;
+  }
+
+  async setupRecurring(accountId: string, planSlug: string, scheduleIdx: number): Promise<void> {
+    console.log(`[MockPaymentEngine] Setup recurring for ${accountId} on plan ${planSlug} schedule ${scheduleIdx}`);
+  }
+
+  async cancelRecurring(accountId: string): Promise<void> {
+    console.log(`[MockPaymentEngine] Cancel recurring for ${accountId}`);
+  }
+}

src/billing/Plan.ts

@@ -0,0 +1,72 @@
+export interface PaymentScheduleConfig {
+  charge_amount: number;
+  charge_period: number; // in days
+  is_default?: boolean;
+}
+
+export class PaymentSchedule {
+  charge_amount: number;
+  charge_period: number;
+  is_default: boolean;
+
+  constructor(config: PaymentScheduleConfig) {
+    this.charge_amount = config.charge_amount;
+    this.charge_period = config.charge_period;
+    this.is_default = config.is_default || false;
+  }
+}
+
+export interface PlanConfig {
+  slug: string;
+  name: string;
+  capabilities?: Record<string, boolean>;
+  downgrade_to_slug?: string;
+  grace_period?: number;
+  schedules?: PaymentScheduleConfig[];
+  account_activate_hook?: (accountId: string) => Promise<void>;
+  account_deactivate_hook?: (accountId: string) => Promise<void>;
+}
+
+export class Plan {
+  slug: string;
+  name: string;
+  capabilities: Record<string, boolean>;
+  downgrade_to_slug?: string;
+  grace_period: number;
+  schedules: PaymentSchedule[];
+  account_activate_hook?: (accountId: string) => Promise<void>;
+  account_deactivate_hook?: (accountId: string) => Promise<void>;
+
+  constructor(config: PlanConfig) {
+    this.slug = config.slug;
+    this.name = config.name;
+    this.capabilities = config.capabilities || {};
+    this.downgrade_to_slug = config.downgrade_to_slug;
+    this.grace_period = config.grace_period || 0;
+    this.schedules = (config.schedules || []).map(s => new PaymentSchedule(s));
+    this.account_activate_hook = config.account_activate_hook;
+    this.account_deactivate_hook = config.account_deactivate_hook;
+  }
+
+  // Registry of plans
+  private static plans: Map<string, Plan> = new Map();
+
+  static init(plans: PlanConfig[]) {
+    Plan.plans.clear();
+    for (const p of plans) {
+      Plan.plans.set(p.slug, new Plan(p));
+    }
+  }
+
+  static get(slug: string): Plan | undefined {
+    return Plan.plans.get(slug);
+  }
+
+  static getAll(): Plan[] {
+    return Array.from(Plan.plans.values());
+  }
+
+  getDefaultSchedule(): PaymentSchedule | undefined {
+    return this.schedules.find(s => s.is_default) || this.schedules[0];
+  }
+}

src/billing/plansConfig.ts

@@ -0,0 +1,33 @@
+import { Plan, PlanConfig } from './Plan';
+
+const plans: PlanConfig[] = [
+  {
+    slug: 'free',
+    name: 'Free',
+    capabilities: {
+      can_access_basic: true,
+      can_access_pro: false,
+    },
+    schedules: [
+      { charge_amount: 0, charge_period: 30, is_default: true }
+    ]
+  },
+  {
+    slug: 'pro',
+    name: 'Pro',
+    capabilities: {
+      can_access_basic: true,
+      can_access_pro: true,
+    },
+    downgrade_to_slug: 'free',
+    grace_period: 7,
+    schedules: [
+      { charge_amount: 2900, charge_period: 30, is_default: true }, // $29.00 / month
+      { charge_amount: 29000, charge_period: 365 } // $290.00 / year
+    ]
+  }
+];
+
+export function initPlans() {
+  Plan.init(plans);
+}

test/billing.spec.ts

@@ -0,0 +1,70 @@
+import { env } from 'cloudflare:test';
+import { describe, it, expect } from 'vitest';
+
+describe('Billing Logic in AccountDO', () => {
+  it('should start with default free plan', async () => {
+    const id = env.ACCOUNT.newUniqueId();
+    const stub = env.ACCOUNT.get(id);
+
+    const res = await stub.fetch('http://do/billing');
+    expect(res.status).toBe(200);
+    const data: any = await res.json();
+
+    expect(data.state.plan_slug).toBe('free');
+    expect(data.state.status).toBe('active');
+    expect(data.plan_details.slug).toBe('free');
+  });
+
+  it('should subscribe to a new plan', async () => {
+    const id = env.ACCOUNT.newUniqueId();
+    const stub = env.ACCOUNT.get(id);
+
+    // Subscribe to Pro
+    const res = await stub.fetch('http://do/billing/subscribe', {
+        method: 'POST',
+        body: JSON.stringify({ plan_slug: 'pro', schedule_idx: 0 })
+    });
+    expect(res.status).toBe(200);
+    const result: any = await res.json();
+    expect(result.success).toBe(true);
+    expect(result.state.plan_slug).toBe('pro');
+    expect(result.state.status).toBe('active');
+    expect(result.state.next_billing_date).toBeDefined();
+
+    // Verify persistence
+    const infoRes = await stub.fetch('http://do/billing');
+    const info: any = await infoRes.json();
+    expect(info.state.plan_slug).toBe('pro');
+  });
+
+  it('should fail to subscribe to invalid plan', async () => {
+    const id = env.ACCOUNT.newUniqueId();
+    const stub = env.ACCOUNT.get(id);
+
+    const res = await stub.fetch('http://do/billing/subscribe', {
+        method: 'POST',
+        body: JSON.stringify({ plan_slug: 'invalid-plan' })
+    });
+    expect(res.status).toBe(400);
+  });
+
+  it('should cancel subscription', async () => {
+    const id = env.ACCOUNT.newUniqueId();
+    const stub = env.ACCOUNT.get(id);
+
+    // Subscribe first
+    await stub.fetch('http://do/billing/subscribe', {
+        method: 'POST',
+        body: JSON.stringify({ plan_slug: 'pro' })
+    });
+
+    // Cancel
+    const res = await stub.fetch('http://do/billing/cancel', {
+        method: 'POST'
+    });
+    expect(res.status).toBe(200);
+    const result: any = await res.json();
+    expect(result.state.status).toBe('canceled');
+    expect(result.state.next_plan_slug).toBe('free'); // Based on plansConfig.ts
+  });
+});