Merge pull request #699 from Merit-Systems/master

zdql · web-flow · commit a4cd5e26a5e6 · 2025-11-25T18:29:25.000-05:00
Merge Referral Fixes
diff --git a/packages/app/control/src/app/(app)/app/[id]/(overview)/_components/header/buttons/index.tsx b/packages/app/control/src/app/(app)/app/[id]/(overview)/_components/header/buttons/index.tsx
@@ -12,7 +12,11 @@ interface Props {
 }
 
 export const HeaderButtons: React.FC<Props> = ({ appId }) => {
-  const [isOwner] = api.apps.app.isOwner.useSuspenseQuery(appId);
+  const { data: isOwner, isLoading } = api.apps.app.isOwner.useQuery(appId);
+
+  if (isLoading || isOwner === undefined) {
+    return <LoadingHeaderButtons />;
+  }
 
   return (
     <div className="flex items-center gap-2">
diff --git a/packages/app/control/src/app/(app)/app/[id]/(overview)/_components/referral-handler.tsx b/packages/app/control/src/app/(app)/app/[id]/(overview)/_components/referral-handler.tsx
@@ -0,0 +1,37 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import { useSearchParams } from 'next/navigation';
+import { api } from '@/trpc/client';
+
+interface Props {
+  appId: string;
+}
+
+export const ReferralHandler: React.FC<Props> = ({ appId }) => {
+  const searchParams = useSearchParams();
+  const referralCode = searchParams.get('referral_code');
+  const [processed, setProcessed] = useState(false);
+
+  const { mutateAsync: registerReferral } =
+    api.apps.app.registerReferral.useMutation();
+
+  useEffect(() => {
+    if (!referralCode || processed) return;
+
+    const processReferralCode = async () => {
+      await registerReferral({
+        appId,
+        code: referralCode,
+      }).catch(() => {
+        // Silently fail - referral code may be invalid, expired, or user may already have a referrer
+      });
+
+      setProcessed(true);
+    };
+
+    void processReferralCode();
+  }, [referralCode, appId, registerReferral, processed]);
+
+  return null;
+};
diff --git a/packages/app/control/src/app/(app)/app/[id]/(overview)/page.tsx b/packages/app/control/src/app/(app)/app/[id]/(overview)/page.tsx
@@ -9,6 +9,7 @@ import { api, HydrateClient } from '@/trpc/server';
 import { HeaderCard, LoadingHeaderCard } from './_components/header';
 import { Setup } from './_components/setup';
 import { Overview } from './_components/overview';
+import { ReferralHandler } from './_components/referral-handler';
 import { userOrRedirect } from '@/auth/user-or-redirect';
 
 export default async function AppPage(props: PageProps<'/app/[id]'>) {
@@ -26,6 +27,7 @@ export default async function AppPage(props: PageProps<'/app/[id]'>) {
 
   return (
     <HydrateClient>
+      <ReferralHandler appId={id} />
       <Body className="gap-0 pt-0">
         <Suspense fallback={<LoadingHeaderCard />}>
           <HeaderCard appId={id} />
diff --git a/packages/app/control/src/app/api/v1/user/referral/route.ts b/packages/app/control/src/app/api/v1/user/referral/route.ts
@@ -3,12 +3,54 @@ import { z } from 'zod';
 import { appIdSchema } from '@/services/db/apps/lib/schemas';
 import { authRoute } from '../../../../../lib/api/auth-route';
 import { setAppMembershipReferrer } from '@/services/db/apps/membership';
+import {
+  getUserAppReferralCode,
+  createAppReferralCode,
+} from '@/services/db/apps/referral-code';
+
+const getUserReferralCodeSchema = z.object({
+  echoAppId: appIdSchema,
+});
 
 const setUserReferrerForAppSchema = z.object({
   echoAppId: appIdSchema,
   code: z.string(),
 });
 
+export const GET = authRoute
+  .query(getUserReferralCodeSchema)
+  .handler(async (_, context) => {
+    const { echoAppId } = context.query;
+    const userId = context.ctx.userId;
+
+    let referralCode = await getUserAppReferralCode(userId, echoAppId);
+
+    if (!referralCode) {
+      referralCode = await createAppReferralCode(userId, {
+        appId: echoAppId,
+        expiresAt: new Date(Date.now() + 365 * 24 * 60 * 60 * 1000),
+      });
+
+      if (!referralCode) {
+        return NextResponse.json(
+          {
+            success: false,
+            message: 'Failed to create referral code',
+          },
+          { status: 500 }
+        );
+      }
+    }
+
+    return NextResponse.json({
+      success: true,
+      message: 'Referral code retrieved successfully',
+      code: referralCode.code,
+      referralLinkUrl: referralCode.referralLinkUrl,
+      expiresAt: referralCode.expiresAt,
+    });
+  });
+
 export const POST = authRoute
   .body(setUserReferrerForAppSchema)
   .handler(async (_, context) => {
diff --git a/packages/app/control/src/services/db/apps/membership.ts b/packages/app/control/src/services/db/apps/membership.ts
@@ -163,21 +163,30 @@ export async function setAppMembershipReferrer(
   echoAppId: string,
   code: string
 ): Promise<boolean> {
-  const appMembership = await db.appMembership.findUnique({
+  // Get or create the app membership
+  const appMembership = await db.appMembership.upsert({
     where: {
       userId_echoAppId: {
         userId,
         echoAppId,
       },
-      referrerId: null,
     },
+    create: {
+      userId,
+      echoAppId,
+      role: AppRole.CUSTOMER,
+      status: MembershipStatus.ACTIVE,
+      totalSpent: 0,
+    },
+    update: {},
   });
 
-  if (appMembership) {
-    // If the user already has a referrer, return false
+  // Check if user already has a referrer
+  if (appMembership.referrerId) {
     return false;
   }
 
+  // Validate the referral code exists
   const referralCode = await db.referralCode.findUnique({
     where: {
       code,
diff --git a/packages/app/control/src/trpc/routers/apps/index.ts b/packages/app/control/src/trpc/routers/apps/index.ts
@@ -23,6 +23,7 @@ import {
   createAppMembershipSchema,
   updateAppMembershipReferrer,
   updateAppMembershipReferrerSchema,
+  setAppMembershipReferrer,
 } from '@/services/db/apps/membership';
 import {
   listAppsSchema,
@@ -262,6 +263,26 @@ export const appsRouter = createTRPCRouter({
         }),
     },
 
+    registerReferral: protectedProcedure
+      .input(z.object({ appId: appIdSchema, code: z.string() }))
+      .mutation(async ({ input, ctx }) => {
+        const success = await setAppMembershipReferrer(
+          ctx.session.user.id,
+          input.appId,
+          input.code
+        );
+
+        if (!success) {
+          throw new TRPCError({
+            code: 'BAD_REQUEST',
+            message:
+              'Referral code could not be applied. It may be invalid, expired, or you may already have a referrer for this app.',
+          });
+        }
+
+        return { success: true };
+      }),
+
     transactions: {
       list: paginatedProcedure
         .concat(protectedProcedure)
diff --git a/packages/tests/integration/README.md b/packages/tests/integration/README.md
@@ -1 +1,82 @@
-# Integration test trigger
+# Integration Tests
+
+This package contains integration tests for the Echo platform.
+
+## Setup
+
+1. Set up the test environment:
+```bash
+pnpm env:setup
+```
+
+2. Seed the test database:
+```bash
+pnpm db:seed
+```
+
+## Running Tests
+
+Run all integration tests:
+```bash
+pnpm test:watch
+```
+
+Run specific test suites:
+```bash
+# Echo Data Server tests
+pnpm test:echo-data-server
+
+# OAuth Protocol tests
+pnpm test:oauth-protocol
+```
+
+## Test Suites
+
+### Echo Data Server Tests
+Located in `tests/echo-data-server/`:
+- `api-key.client.test.ts` - API key authentication and usage
+- `402-auth.client.test.ts` - Payment required (402) authentication flow
+- `echo-access-jwt.client.test.ts` - JWT token validation
+- `free-tier.client.test.ts` - Free tier functionality
+- `referral-code.client.test.ts` - Referral code creation and application
+- `in-flight-requests.test.ts` - Concurrent request handling
+
+### OAuth Protocol Tests
+Located in `tests/oauth-protocol/`:
+- OAuth authorization flow
+- Token refresh and lifecycle
+- PKCE security
+- CSRF vulnerability testing
+
+## Referral Code Tests
+
+The referral code integration tests (`referral-code.client.test.ts`) cover:
+
+1. **GET endpoint** - Retrieval of referral codes:
+   - Retrieving an existing referral code for a user
+   - Auto-creating a referral code for users who don't have one
+   - Ensuring consistency across multiple requests
+
+2. **POST endpoint** - Application of referral codes:
+   - Successfully applying another user's referral code
+   - Rejecting invalid referral codes
+   - Preventing users from applying codes when they already have a referrer
+
+### Test Data
+
+Referral code test data is defined in `config/test-data.ts`:
+- Primary user has referral code: `TEST-REFERRAL-CODE-PRIMARY`
+- Secondary user has referral code: `TEST-REFERRAL-CODE-SECONDARY`
+- Tertiary user has no referral code (created during tests)
+
+## Database Management
+
+Reset and reseed the database:
+```bash
+pnpm db:reset-and-seed
+```
+
+Reset only:
+```bash
+pnpm db:reset
+```
diff --git a/packages/tests/integration/config/test-data.ts b/packages/tests/integration/config/test-data.ts
@@ -131,6 +131,26 @@ export const TEST_DATA = {
     },
   },
 
+  // Referral code configurations
+  referralCodes: {
+    primaryUserCode: {
+      id: '88888888-8888-4888-8888-888888888888',
+      code: 'TEST-REFERRAL-CODE-PRIMARY',
+      userId: '11111111-1111-4111-8111-111111111111', // Primary test user
+      expiresAt: new Date(Date.now() + 365 * 24 * 60 * 60 * 1000), // 1 year from now
+      isArchived: false,
+      usedAt: null,
+    },
+    secondaryUserCode: {
+      id: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa',
+      code: 'TEST-REFERRAL-CODE-SECONDARY',
+      userId: '33333333-3333-4333-8333-333333333333', // Secondary test user
+      expiresAt: new Date(Date.now() + 365 * 24 * 60 * 60 * 1000), // 1 year from now
+      isArchived: false,
+      usedAt: null,
+    },
+  },
+
   // Test timeouts and delays
   timeouts: {
     default: 30000,
@@ -234,6 +254,16 @@ export const TEST_SPEND_POOL_IDS = {
   primary: TEST_DATA.spendPools.primary.id,
 };
 
+export const TEST_REFERRAL_CODE_IDS = {
+  primary: TEST_DATA.referralCodes.primaryUserCode.id,
+  secondary: TEST_DATA.referralCodes.secondaryUserCode.id,
+};
+
+export const TEST_REFERRAL_CODES = {
+  primary: TEST_DATA.referralCodes.primaryUserCode.code,
+  secondary: TEST_DATA.referralCodes.secondaryUserCode.code,
+};
+
 // Type definitions for test data
 export type TestData = typeof TEST_DATA;
 export type TestUser = typeof TEST_DATA.users.primary;
@@ -242,3 +272,4 @@ export type TestApiKey = typeof TEST_DATA.apiKeys.primary;
 export type TestSpendPool = typeof TEST_DATA.spendPools.primary;
 export type TestUserSpendPoolUsage =
   typeof TEST_DATA.userSpendPoolUsage.tertiaryUserPrimaryPool;
+export type TestReferralCode = typeof TEST_DATA.referralCodes.primaryUserCode;
diff --git a/packages/tests/integration/scripts/seed-integration-db.ts b/packages/tests/integration/scripts/seed-integration-db.ts
@@ -25,6 +25,7 @@ export async function seedIntegrationDatabase() {
     await prisma.spendPool.deleteMany();
     await prisma.apiKey.deleteMany();
     await prisma.appMembership.deleteMany();
+    await prisma.referralCode.deleteMany();
     await prisma.echoApp.deleteMany();
     await prisma.user.deleteMany();
 
@@ -171,6 +172,17 @@ export async function seedIntegrationDatabase() {
 
     console.log('🤖 Created test LLM transaction');
 
+    // Create test referral codes
+    await prisma.referralCode.create({
+      data: TEST_DATA.referralCodes.primaryUserCode,
+    });
+
+    await prisma.referralCode.create({
+      data: TEST_DATA.referralCodes.secondaryUserCode,
+    });
+
+    console.log('🎟️ Created test referral codes');
+
     console.log('✅ Integration test database seeded successfully');
     console.log('\n📊 Summary:');
     console.log(`  - Users: 3`);
@@ -181,6 +193,7 @@ export async function seedIntegrationDatabase() {
     console.log(`  - User Spend Pool Usage: 1`);
     console.log(`  - Payments: 1`);
     console.log(`  - LLM Transactions: 1`);
+    console.log(`  - Referral Codes: 2`);
   } catch (error) {
     console.error('❌ Error seeding integration test database:', error);
     throw error;
diff --git a/packages/tests/integration/tests/echo-data-server/referral-code.client.test.ts b/packages/tests/integration/tests/echo-data-server/referral-code.client.test.ts
diff --git a/packages/tests/integration/utils/api-client.ts b/packages/tests/integration/utils/api-client.ts
