ddd refactor: copilot streaming

apps/rowboat/app/actions/copilot.actions.ts

@@ -1,5 +1,5 @@
 'use server';
-import { 
+import {
     CopilotAPIRequest,
     CopilotChatContext, CopilotMessage,
     DataSourceSchemaForCopilot,
@@ -8,15 +8,18 @@ import {
     Workflow} from "../lib/types/workflow_types";
 import { z } from 'zod';
 import { projectAuthCheck } from "./project.actions";
-import { redisClient } from "../lib/redis";
 import { authorizeUserAction, logUsage } from "./billing.actions";
 import { USE_BILLING } from "../lib/feature_flags";
 import { getEditAgentInstructionsResponse } from "../../src/application/lib/copilot/copilot";
 import { container } from "@/di/container";
 import { IUsageQuotaPolicy } from "@/src/application/policies/usage-quota.policy.interface";
 import { UsageTracker } from "../lib/billing";
+import { authCheck } from "./auth.actions";
+import { ICreateCopilotCachedTurnController } from "@/src/interface-adapters/controllers/copilot/create-copilot-cached-turn.controller";
+import { BillingError } from "@/src/entities/errors/common";
 
 const usageQuotaPolicy = container.resolve<IUsageQuotaPolicy>('usageQuotaPolicy');
+const createCopilotCachedTurnController = container.resolve<ICreateCopilotCachedTurnController>('createCopilotCachedTurnController');
 
 export async function getCopilotResponseStream(
     projectId: string,
@@ -27,40 +30,29 @@ export async function getCopilotResponseStream(
 ): Promise<{
     streamId: string;
 } | { billingError: string }> {
-    await projectAuthCheck(projectId);
-    await usageQuotaPolicy.assertAndConsumeProjectAction(projectId);
+    const user = await authCheck();
 
-    // Check billing authorization
-    const authResponse = await authorizeUserAction({
-        type: 'use_credits',
-    });
-    if (!authResponse.success) {
-        return { billingError: authResponse.error || 'Billing error' };
+    try {
+        const { key } = await createCopilotCachedTurnController.execute({
+            caller: 'user',
+            userId: user.id,
+            data: {
+                projectId,
+                messages,
+                workflow: current_workflow_config,
+                context,
+                dataSources,
+            }
+        });
+        return {
+            streamId: key,
+        };
+    } catch (err) {
+        if (err instanceof BillingError) {
+            return { billingError: err.message };
+        }
+        throw err;
     }
-
-    await usageQuotaPolicy.assertAndConsumeProjectAction(projectId);
-    
-    // prepare request
-    const request: z.infer<typeof CopilotAPIRequest> = {
-        projectId,
-        messages,
-        workflow: current_workflow_config,
-        context,
-        dataSources: dataSources,
-    };
-
-    // serialize the request
-    const payload = JSON.stringify(request);
-
-    // create a uuid for the stream
-    const streamId = crypto.randomUUID();
-
-    // store payload in redis
-    await redisClient.set(`copilot-stream-${streamId}`, payload, 'EX', 60 * 10); // expire in 10 minutes
-
-    return {
-        streamId,
-    };
 }
 
 export async function getCopilotAgentInstructions(

apps/rowboat/app/api/copilot-stream-response/[streamId]/route.ts

@@ -1,70 +1,45 @@
-import { getCustomerIdForProject, logUsage, UsageTracker } from "@/app/lib/billing";
-import { USE_BILLING } from "@/app/lib/feature_flags";
-import { redisClient } from "@/app/lib/redis";
-import { CopilotAPIRequest } from "@/src/entities/models/copilot";
-import { streamMultiAgentResponse } from "@/src/application/lib/copilot/copilot";
+import { container } from "@/di/container";
+import { IRunCopilotCachedTurnController } from "@/src/interface-adapters/controllers/copilot/run-copilot-cached-turn.controller";
+import { requireAuth } from "@/app/lib/auth";
 
 export const maxDuration = 300;
 
 export async function GET(request: Request, props: { params: Promise<{ streamId: string }> }) {
   const params = await props.params;
-  // get the payload from redis
-  const payload = await redisClient.get(`copilot-stream-${params.streamId}`);
-  if (!payload) {
-    return new Response("Stream not found", { status: 404 });
-  }
 
-  // parse the payload
-  const { projectId, context, messages, workflow, dataSources } = CopilotAPIRequest.parse(JSON.parse(payload));
+  // get user data
+  const user = await requireAuth();
 
-  // fetch billing customer id
-  let billingCustomerId: string | null = null;
-  if (USE_BILLING) {
-    billingCustomerId = await getCustomerIdForProject(projectId);
-  }
+  const runCopilotCachedTurnController = container.resolve<IRunCopilotCachedTurnController>("runCopilotCachedTurnController");
 
-  const usageTracker = new UsageTracker();
   const encoder = new TextEncoder();
-  let messageCount = 0;
 
   const stream = new ReadableStream({
     async start(controller) {
       try {
         // Iterate over the copilot stream generator
-        for await (const event of streamMultiAgentResponse(
-          usageTracker,
-          projectId,
-          context,
-          messages,
-          workflow,
-          dataSources || [],
-        )) {
+        for await (const event of runCopilotCachedTurnController.execute({
+          caller: "user",
+          userId: user.id,
+          apiKey: request.headers.get("Authorization")?.split(" ")[1],
+          key: params.streamId,
+        })) {
           // Check if this is a content event
           if ('content' in event) {
-            messageCount++;
             controller.enqueue(encoder.encode(`event: message\ndata: ${JSON.stringify(event)}\n\n`));
           } else if ('type' in event && event.type === 'tool-call') {
             controller.enqueue(encoder.encode(`event: tool-call\ndata: ${JSON.stringify(event)}\n\n`));
           } else if ('type' in event && event.type === 'tool-result') {
             controller.enqueue(encoder.encode(`event: tool-result\ndata: ${JSON.stringify(event)}\n\n`));
-          } else {
-            controller.enqueue(encoder.encode(`event: done\ndata: ${JSON.stringify(event)}\n\n`));
           }
         }
       } catch (error) {
         console.error('Error processing copilot stream:', error);
         controller.error(new Error("Something went wrong. Please try again."));
       } finally {
-        // log copilot usage
-        if (USE_BILLING && billingCustomerId) {
-          try {
-            await logUsage(billingCustomerId, {
-              items: usageTracker.flush(),
-            });
-          } catch (error) {
-            console.error("Error logging usage", error);
-          }
-        }
+        console.log("closing stream");
+        controller.enqueue(encoder.encode(`event: done\ndata: ${JSON.stringify({ type: 'done' })}\n\n`));
+        controller.enqueue(encoder.encode("event: end\n\n"));
         controller.close();
       }
     },