[feat] referrals system (#17)

Author: brandonkachen

.env.example

@@ -18,13 +18,16 @@ NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY='your publishable key'
 DATABASE_URL='postgresql://postgres:mysecretpassword@localhost:5432/postgres'
 
 # Web
-APP_URL='https://localhost:3000'
+NEXT_PUBLIC_APP_URL='https://localhost:3000'
 GOOGLE_SITE_VERIFICATION_ID='your google verification id'
 GITHUB_ID='your github client ID'
 GITHUB_SECRET='your github secret ID'
 NEXTAUTH_SECRET='your next-auth secret'
 NEXTAUTH_URL='http://localhost:3000'
 
+# All
+NEXT_PUBLIC_SUPPORT_EMAIL='support@manicode.ai'
+
 #################################################################
 # PLEASE SEE JAMES OR BRANDON FOR THE REAL ENVIRONMENT VARIABLES!
 #################################################################

authentication.knowledge.md

@@ -16,7 +16,7 @@ Manicode implements a secure authentication flow that involves the npm-app (CLI)
      a. `fingerprintId` (to link the current session with user credentials)
      b. Timestamp (5 minutes in the future, for link expiration)
      c. Hash of the above + a secret value (for request verification)
-   - Backend appends this auth code to the login URL: `${APP_URL}/login?auth_code=<token-goes-here>`
+   - Backend appends this auth code to the login URL: `${NEXT_PUBLIC_APP_URL}/login?auth_code=<token-goes-here>`
 
 3. User Login:
 

backend/knowledge.md

@@ -70,6 +70,31 @@ The backend handles file operations for the Manicode project:
 - **Reading Files**: The `read_files` tool allows the AI to access project file contents.
 - **Applying Changes**: The `applyChanges` function in `prompts.ts` processes and applies file modifications suggested by the AI.
 
+## Code Organization and Best Practices
+
+1. **Centralize Shared Logic**: When implementing functionality that's used in multiple places (e.g., web API and backend), create shared functions to promote code reuse and consistency. This is particularly important for business logic like referral handling and usage calculations.
+
+2. **Shared Function Location**: Place shared functions in a common directory accessible to both the web API and backend. Consider creating a `common/src/utils` directory for such shared functionality.
+
+3. **DRY Principle**: Always look for opportunities to refactor repeated code into shared, reusable functions. This not only reduces code duplication but also makes maintenance and updates easier.
+
+4. **Consistent API**: When creating shared functions, ensure they have a consistent API that can be easily used by different parts of the application.
+
+5. **Testing Shared Functions**: Implement unit tests for shared functions to ensure they work correctly in all contexts where they are used.
+
+6. **Documentation**: Document shared functions clearly, including their purpose, inputs, outputs, and any side effects, so that other developers can use them effectively.
+
+7. **Version Control**: When making changes to shared functions, consider the impact on all parts of the application that use them, and test thoroughly.
+
+8. **Dependency Injection**: Prefer pulling common dependencies (like database connections or environment variables) from centralized locations rather than passing them as parameters. This reduces function complexity and improves maintainability.
+
+9. **Single Responsibility Principle**: Design functions to have a single, well-defined purpose. For example, separate the logic of determining eligibility for a referral code from the generation of the full referral link.
+
+10. **Abstraction Refinement**: Be prepared to refine initial implementations as the system's needs become clearer. This might involve changing function signatures, splitting functions, or adjusting their purposes to better fit the overall architecture.
+
+
+
+
 ## Development Guidelines
 
 1. **Type Safety**: Utilize TypeScript's type system to ensure code reliability and catch errors early.
@@ -145,6 +170,29 @@ export const logger = pino({
 4. Create a robust testing suite for backend components.
 5. Optimize the file diff generation process for better reliability and performance.
 
+## Referral System
+
+The referral system is an important feature of our application. Here are key points to remember:
+
+1. **Referral Limit**: Users are limited to a maximum number of successful referrals (currently set to 5).
+
+2. **Limit Enforcement**: The referral limit must be enforced during the redemption process (POST request), not just when displaying referral information (GET request).
+
+3. **Centralized Logic**: The `hasMaxedReferrals` function in `common/src/util/referral.ts` is used to check if a user has reached their referral limit. This function should be used consistently across the application to ensure uniform enforcement of the referral limit.
+
+4. **Redemption Process**: When redeeming a referral code (in the POST request handler), always check if the referrer has maxed out their referrals before processing the redemption. This ensures that users cannot exceed their referral limit even if they distribute their referral code widely.
+
+5. **Error Handling**: Provide clear error messages when a referral code cannot be redeemed due to the referrer reaching their limit. This helps maintain a good user experience.
+
+Remember to keep the referral system logic consistent between the backend API and the websocket server to ensure uniform behavior across different parts of the application.
+
+
+
+
+
+
+
+
 ## Debugging Docker Issues
 
 - When encountering "Cannot find module" errors in a Docker container, it's important to verify the contents of the container itself, not just the local build.
@@ -194,3 +242,5 @@ This project uses Bun for testing instead of Jest. When writing tests, keep the
 - Bun's test API is similar to Jest's, but there are some differences in implementation.
 - When mocking methods, use `mock(object.method)` instead of Jest's `jest.spyOn(object, 'method')`.
 - Bun's `mock` function expects 0-1 arguments, not 2 like Jest's `spyOn`.
+
+Remember to keep this knowledge file updated as the application evolves or new features are added.

backend/src/env.mjs

@@ -24,7 +24,7 @@ export const env = createEnv({
     STRIPE_SUBSCRIPTION_PRICE_ID: z.string().min(1),
     PORT: z.coerce.number().min(1000),
     ENVIRONMENT: z.string().min(1),
-    APP_URL: z.string().min(1),
+    NEXT_PUBLIC_APP_URL: z.string().min(1),
     NEXTAUTH_SECRET: z.string().min(1),
   },
   runtimeEnv: process.env,

backend/src/websockets/middleware.ts

@@ -82,9 +82,6 @@ protec.use(async (action, _clientSessionId, _) => {
   logger.debug(`Protecting action of type: '${action.type}'`)
 })
 protec.use(async (action, _clientSessionId, ws) => {
-  if (env.ENVIRONMENT === 'local') {
-    return
-  }
   return match(action)
     .with(
       {

backend/src/websockets/websocket-action.ts

@@ -11,13 +11,14 @@ import { getSearchSystemPrompt } from '../system-prompt'
 import { promptClaude } from '../claude'
 import { env } from '../env.mjs'
 import db from 'common/db'
-import * as schema from 'common/db/schema'
 import { genAuthCode } from 'common/util/credentials'
+import * as schema from 'common/db/schema'
 import { claudeModels } from 'common/constants'
 import { protec } from './middleware'
 import { getQuotaManager } from '@/billing/quota-manager'
 import { logger, withLoggerContext } from '@/util/logger'
 import { generateCommitMessage } from '@/generate-commit-message'
+import { hasMaxedReferrals } from 'common/util/server/referral'
 
 export const sendAction = (ws: WebSocket, action: ServerAction) => {
   sendMessage(ws, {
@@ -51,10 +52,7 @@ async function calculateUsage(fingerprintId: string, userId?: string) {
     userId ? 'authenticated' : 'anonymous',
     userId ?? fingerprintId
   )
-  const { creditsUsed, quota } = await quotaManager.checkQuota()
-  if (env.ENVIRONMENT === 'local') {
-    return { usage: creditsUsed, limit: Infinity }
-  }
+  const { creditsUsed, quota } = await quotaManager.updateQuota()
   return { usage: creditsUsed, limit: quota }
 }
 
@@ -180,7 +178,10 @@ const onClearAuthTokenRequest = async (
 }
 
 const onLoginCodeRequest = (
-  { fingerprintId }: Extract<ClientAction, { type: 'login-code-request' }>,
+  {
+    fingerprintId,
+    referralCode,
+  }: Extract<ClientAction, { type: 'login-code-request' }>,
   _clientSessionId: string,
   ws: WebSocket
 ): void => {
@@ -191,7 +192,7 @@ const onLoginCodeRequest = (
       expiresAt.toString(),
       env.NEXTAUTH_SECRET
     )
-    const loginUrl = `${env.APP_URL}/login?auth_code=${fingerprintId}.${expiresAt}.${fingerprintHash}`
+    const loginUrl = `${env.NEXT_PUBLIC_APP_URL}/login?auth_code=${fingerprintId}.${expiresAt}.${fingerprintHash}${referralCode ? `&referral_code=${referralCode}` : ''}`
 
     sendAction(ws, {
       type: 'login-code-response',
@@ -323,10 +324,20 @@ const onUsageRequest = async (
   const { usage, limit } = await calculateUsage(fingerprintId, userId)
   await withLoggerContext({ fingerprintId, userId, usage, limit }, async () => {
     logger.info('Sending usage info')
+
+    let referralLink: string | undefined = undefined
+    if (userId) {
+      const shouldGenerateReferralLink = await hasMaxedReferrals(userId)
+      if (shouldGenerateReferralLink.reason === undefined) {
+        referralLink = shouldGenerateReferralLink.referralLink
+      }
+    }
+
     sendAction(ws, {
       type: 'usage-response',
       usage,
       limit,
+      referralLink,
     })
   })
 }
@@ -414,7 +425,7 @@ export const onWebsocketAction = async (
 }
 
 subscribeToAction('user-input', protec.run(onUserInput))
-subscribeToAction('init', () => protec.run(onInit))
+subscribeToAction('init', protec.run(onInit))
 
 subscribeToAction('clear-auth-token', onClearAuthTokenRequest)
 subscribeToAction('login-code-request', onLoginCodeRequest)

bun.lockb

@@ -91,6 +91,7 @@ export const CLIENT_ACTION_SCHEMA = z.discriminatedUnion('type', [
   z.object({
     type: z.literal('login-code-request'),
     fingerprintId: z.string(),
+    referralCode: z.string().optional(),
   }),
   z.object({
     type: z.literal('login-status-request'),
@@ -167,6 +168,7 @@ export const SERVER_ACTION_SCHEMA = z.discriminatedUnion('type', [
     type: z.literal('usage-response'),
     usage: z.number(),
     limit: z.number(),
+    referralLink: z.string().optional(),
   }),
   z.object({
     type: z.literal('quota-exceeded'),

common/src/actions.ts

@@ -1,3 +1,5 @@
+// import { env } from './env.mjs'
+
 export const STOP_MARKER = '[' + 'END]'
 export const FIND_FILES_MARKER = '[' + 'FIND_FILES_PLEASE]'
 export const TOOL_RESULT_MARKER = '[' + 'TOOL_RESULT]'
@@ -39,11 +41,20 @@ export const SKIPPED_TERMINAL_COMMANDS = [
 
 export const MAX_DATE = new Date(86399999999999)
 
-export const CREDITS_USAGE_LIMITS = {
-  ANON: 1_000,
-  FREE: 2_500,
-  PAID: 50_000,
-}
+export type UsageLimits = 'ANON' | 'FREE' | 'PAID'
+export const CREDITS_USAGE_LIMITS: Record<UsageLimits, number> =
+  process.env.NEXT_PUBLIC_ENVIRONMENT === 'local'
+    ? {
+        ANON: 1_000_000,
+        FREE: 2_500_000,
+        PAID: 50_000_000,
+      }
+    : {
+        ANON: 1_000,
+        FREE: 2_500,
+        PAID: 50_000,
+      }
+export const CREDITS_REFERRAL_BONUS = 500
 
 export const claudeModels = {
   sonnet: 'claude-3-5-sonnet-20240620',
@@ -61,3 +72,5 @@ export const models = {
 }
 
 export const TEST_USER_ID = 'test-user-id'
+
+export const MAX_REFERRALS = 5

common/src/constants.ts

@@ -0,0 +1,30 @@
+DO $$ BEGIN
+ CREATE TYPE "public"."referral_status" AS ENUM('pending', 'completed');
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+CREATE TABLE IF NOT EXISTS "referral" (
+	"referrer_id" text NOT NULL,
+	"referred_id" text NOT NULL,
+	"status" "referral_status" DEFAULT 'pending' NOT NULL,
+	"credits" integer NOT NULL,
+	"created_at" timestamp DEFAULT now() NOT NULL,
+	"completed_at" timestamp,
+	CONSTRAINT "referral_referrer_id_referred_id_pk" PRIMARY KEY("referrer_id","referred_id")
+);
+--> statement-breakpoint
+ALTER TABLE "user" ADD COLUMN "referral_code" text DEFAULT 'ref-' || gen_random_uuid();--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "referral" ADD CONSTRAINT "referral_referrer_id_user_id_fk" FOREIGN KEY ("referrer_id") REFERENCES "public"."user"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "referral" ADD CONSTRAINT "referral_referred_id_user_id_fk" FOREIGN KEY ("referred_id") REFERENCES "public"."user"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+ALTER TABLE "user" ADD CONSTRAINT "user_referral_code_unique" UNIQUE("referral_code");