Environments

.github/workflows/e2e-tests.yml

@@ -15,7 +15,7 @@ jobs:
   e2e-tests:
     name: 🧪 Run E2E Tests
     runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 25
 
     steps:
       - name: 📥 Checkout code

app/api/[resource]/[resourceId]/members/route.ts

@@ -1,7 +1,7 @@
 import { NextRequest, NextResponse } from 'next/server';
 import { z } from 'zod';
-import { userService, type UserProfile } from '~/lib/services/userService';
-import { getDataSource } from '~/lib/services/datasourceService';
+import { type UserProfile, userService } from '~/lib/services/userService';
+import { getDataSource } from '~/lib/services/dataSourceService';
 import { getEnvironment } from '~/lib/services/environmentService';
 import { getWebsite } from '~/lib/services/websiteService';
 import { getPermissionLevelDetails, type PermissionLevel } from '~/lib/services/permissionService';

app/api/chat/route.ts

@@ -10,11 +10,9 @@ import { messageService } from '~/lib/services/messageService';
 import { MESSAGE_ROLE } from '~/types/database';
 import { createId } from '@paralleldrive/cuid2';
 import { conversationService } from '~/lib/services/conversationService';
-import type { StarterPluginId } from '~/lib/plugins/types';
 import type { FileMap } from '~/lib/stores/files';
 import { getTelemetry } from '~/lib/telemetry/telemetry-manager';
 import { TelemetryEventType } from '~/lib/telemetry/telemetry-types';
-import { prisma } from '~/lib/prisma';
 import { LLMManager } from '~/lib/modules/llm/manager';
 import { DataSourcePluginManager } from '~/lib/plugins/data-access/data-access-plugin-manager';
 import { type UserProfile, userService } from '~/lib/services/userService';
@@ -23,6 +21,7 @@ import { getDatabaseSchema } from '~/lib/schema';
 import { requireUserId } from '~/auth/session';
 import { formatDbSchemaForLLM } from '~/lib/.server/llm/database-source';
 import { AI_SDK_INVALID_KEY_ERROR } from '~/utils/constants';
+import { getDatabaseUrl } from '~/lib/services/dataSourceService';
 
 const WORK_DIR = '/home/project';
 
@@ -201,8 +200,13 @@ async function chatAction(request: NextRequest) {
             };
             dataStream.writeData(currentProgressAnnotation);
 
-            const dataSource = await conversationService.getConversationDataSource(conversationId);
-            const databaseSchema = await getDatabaseSchema(dataSource.id, userId);
+            const environmentDataSource =
+              await conversationService.getConversationEnvironmentDataSource(conversationId);
+            const databaseSchema = await getDatabaseSchema(
+              environmentDataSource.dataSourceId,
+              environmentDataSource.environmentId,
+              userId,
+            );
 
             implementationPlan = await createImplementationPlan({
               isFirstUserMessage: !!userMessageProperties.isFirstUserMessage,
@@ -280,6 +284,7 @@ async function chatAction(request: NextRequest) {
 
                 const userId = await requireUserId(request);
                 const user = await userService.getUser(userId);
+
                 await trackChatPrompt(conversationId, currentModel, user, userMessageProperties.content);
               } catch (error) {
                 logger.error('Failed to save prompt', error);
@@ -325,7 +330,7 @@ async function chatAction(request: NextRequest) {
             options,
             files,
             promptId,
-            starterId: conversation?.starterId as StarterPluginId,
+            conversation: conversation!,
             contextOptimization,
             contextFiles: filteredFiles,
             summary,
@@ -430,36 +435,33 @@ async function trackChatPrompt(
   userMessage: string,
 ): Promise<void> {
   try {
-    const conversationWithDataSource = await prisma.conversation.findUnique({
-      where: { id: conversationId },
-      include: {
-        dataSource: {
-          select: {
-            connectionString: true,
-          },
-        },
-      },
-    });
+    const environmentDataSource = await conversationService.getConversationEnvironmentDataSource(conversationId);
+    const dataSourceUrl = await getDatabaseUrl(
+      user.id,
+      environmentDataSource.dataSourceId,
+      environmentDataSource.environmentId,
+    );
 
-    if (conversationWithDataSource?.dataSource.connectionString) {
-      const pluginId = DataSourcePluginManager.getAccessorPluginId(
-        conversationWithDataSource.dataSource.connectionString,
-      );
-
-      const telemetry = await getTelemetry();
-      await telemetry.trackTelemetryEvent(
-        {
-          eventType: TelemetryEventType.USER_CHAT_PROMPT,
-          properties: {
-            conversationId,
-            dataSourceType: pluginId,
-            llmModel,
-            userMessage,
-          },
-        },
-        user,
-      );
+    if (!dataSourceUrl) {
+      logger.warn('No data source URL found for telemetry tracking');
+      return;
     }
+
+    const pluginId = DataSourcePluginManager.getAccessorPluginId(dataSourceUrl);
+
+    const telemetry = await getTelemetry();
+    await telemetry.trackTelemetryEvent(
+      {
+        eventType: TelemetryEventType.USER_CHAT_PROMPT,
+        properties: {
+          conversationId,
+          dataSourceType: pluginId,
+          llmModel,
+          userMessage,
+        },
+      },
+      user,
+    );
   } catch (telemetryError) {
     logger.error('Failed to track telemetry event', telemetryError);
   }

app/api/conversations/[conversationId]/route.ts

@@ -108,18 +108,16 @@ async function handleDelete(conversationId: string, userId: string) {
 
 const UPDATE_CONVERSATION_SCHEMA = z.object({
   description: z.string().nullable().optional(),
+  environmentId: z.string().optional(),
+  dataSourceId: z.string().optional(),
 });
 
 async function handlePatch(conversationId: string, userId: string, request: NextRequest) {
   try {
     const body = await request.json();
     const updateData = UPDATE_CONVERSATION_SCHEMA.parse(body);
 
-    const updatedConversation = await conversationService.updateConversationDescription(
-      conversationId,
-      userId,
-      updateData,
-    );
+    const updatedConversation = await conversationService.updateConversation(conversationId, userId, updateData);
 
     if (!updatedConversation) {
       return NextResponse.json({ error: 'Conversation not found' }, { status: 404 });

app/api/conversations/[conversationId]/snapshots/message/[messageId]/route.ts

@@ -0,0 +1,94 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { z } from 'zod';
+import { conversationService } from '~/lib/services/conversationService';
+import { StorageServiceFactory } from '~/lib/services/storage/storage-service-factory';
+import { snapshotService } from '~/lib/services/snapshotService';
+import { logger } from '~/utils/logger';
+import { requireUserId } from '~/auth/session';
+
+// Zod schema for the request body
+const UPDATE_SNAPSHOT_SCHEMA = z.object({
+  fileMap: z.record(
+    z.string(),
+    z
+      .union([
+        z.object({
+          type: z.literal('file'),
+          content: z.string(),
+          isBinary: z.boolean(),
+        }),
+        z.object({
+          type: z.literal('folder'),
+        }),
+      ])
+      .optional(),
+  ),
+});
+
+export async function PUT(
+  request: NextRequest,
+  { params }: { params: Promise<{ conversationId: string; messageId: string }> },
+) {
+  const { conversationId, messageId } = await params;
+
+  const userId = await requireUserId(request);
+
+  const conversation = await conversationService.getConversation(conversationId);
+
+  if (!conversation) {
+    return NextResponse.json({ error: 'Conversation not found' }, { status: 404 });
+  }
+
+  // Check if the conversation belongs to the authenticated user
+  if (conversation.userId !== userId) {
+    return NextResponse.json({ error: 'Conversation not found' }, { status: 404 });
+  }
+
+  try {
+    // Find the snapshot associated with this message ID
+    const snapshot = await snapshotService.getSnapshotByMessageId(messageId);
+
+    if (!snapshot) {
+      logger.error(`Snapshot not found for message ${messageId}`);
+      return NextResponse.json({ error: 'Snapshot not found' }, { status: 404 });
+    }
+
+    if (snapshot.conversationId !== conversationId) {
+      logger.error(`Snapshot ${snapshot.id} does not belong to conversation ${conversationId}`);
+      return NextResponse.json({ error: 'Snapshot does not belong to this conversation' }, { status: 403 });
+    }
+
+    const storageService = StorageServiceFactory.get();
+    const data = await storageService.get(snapshot.storageKey);
+    const fileMap = JSON.parse(data.toString());
+
+    if (!fileMap) {
+      logger.error(`No snapshot files found for snapshot ${snapshot.id}`);
+      return NextResponse.json({ error: 'No snapshot files found' }, { status: 404 });
+    }
+
+    // Validate and parse the request body
+    const body = await request.json();
+    const requestData = UPDATE_SNAPSHOT_SCHEMA.parse(body);
+    const updatedFileMap = { ...fileMap, ...requestData.fileMap };
+
+    const serializedData = Buffer.from(JSON.stringify(updatedFileMap, null, 2));
+    await storageService.save(snapshot.storageKey, serializedData);
+
+    return NextResponse.json({
+      success: true,
+      snapshot: {
+        ...snapshot,
+        fileMap: updatedFileMap,
+      },
+    });
+  } catch (error) {
+    if (error instanceof z.ZodError) {
+      return NextResponse.json({ error: 'Invalid request body', details: error.errors }, { status: 400 });
+    }
+
+    logger.error('Failed to update snapshot:', error);
+
+    return NextResponse.json({ error: 'Failed to update snapshot' }, { status: 500 });
+  }
+}

app/api/conversations/route.ts

@@ -16,6 +16,7 @@ export async function POST(request: NextRequest) {
       dataSourceId: string;
       description?: string;
       messages?: Message[];
+      environmentId: string;
     };
 
     if (!body?.dataSourceId) {
@@ -27,6 +28,7 @@ export async function POST(request: NextRequest) {
 
       const conversation = await conversationService.createConversation(
         body.dataSourceId,
+        body.environmentId,
         userId,
         body.description,
         tx,

app/api/data-sources/[id]/route.ts

@@ -2,35 +2,47 @@ import { NextRequest, NextResponse } from 'next/server';
 import {
   deleteDataSource,
   getConversationCount,
-  getDataSource,
+  getEnvironmentDataSource,
   updateDataSource,
-} from '~/lib/services/datasourceService';
+} from '~/lib/services/dataSourceService';
 import { requireUserId } from '~/auth/session';
 
 export async function GET(request: NextRequest, { params }: { params: Promise<{ id: string }> }) {
   const userId = await requireUserId(request);
 
   const { id } = await params;
+  const { searchParams } = new URL(request.url);
+  const environmentId = searchParams.get('environmentId');
 
-  const dataSource = await getDataSource(id);
+  if (!environmentId) {
+    return NextResponse.json({ success: false, error: 'Environment ID is required' }, { status: 400 });
+  }
+
+  const environmentDataSource = await getEnvironmentDataSource(id, userId, environmentId);
 
-  if (!dataSource) {
+  if (!environmentDataSource) {
     return NextResponse.json({ success: false, error: 'Data source not found' }, { status: 404 });
   }
 
   const conversationCount = await getConversationCount(id, userId);
 
-  return NextResponse.json({ success: true, dataSource, conversationCount });
+  return NextResponse.json({ success: true, environmentDataSource, conversationCount });
 }
 
 export async function PUT(request: NextRequest, { params }: { params: Promise<{ id: string }> }) {
   const userId = await requireUserId(request);
 
   const { id } = await params;
+  const { searchParams } = new URL(request.url);
+  const environmentId = searchParams.get('environmentId');
+
+  if (!environmentId) {
+    return NextResponse.json({ success: false, error: 'Environment ID is required' }, { status: 400 });
+  }
 
-  const dataSource = await getDataSource(id);
+  const environmentDataSource = await getEnvironmentDataSource(id, userId, environmentId);
 
-  if (!dataSource) {
+  if (!environmentDataSource) {
     return NextResponse.json({ success: false, error: 'Data source not found' }, { status: 404 });
   }
 
@@ -53,10 +65,16 @@ export async function PUT(request: NextRequest, { params }: { params: Promise<{
 export async function DELETE(request: NextRequest, { params }: { params: Promise<{ id: string }> }) {
   const userId = await requireUserId(request);
   const { id } = await params;
+  const { searchParams } = new URL(request.url);
+  const environmentId = searchParams.get('environmentId');
+
+  if (!environmentId) {
+    return NextResponse.json({ success: false, error: 'Environment ID is required' }, { status: 400 });
+  }
 
-  const dataSource = await getDataSource(id);
+  const environmentDataSource = await getEnvironmentDataSource(id, userId, environmentId);
 
-  if (!dataSource) {
+  if (!environmentDataSource) {
     return NextResponse.json({ success: false, error: 'Data source not found' }, { status: 404 });
   }
 

app/api/data-sources/[id]/url/route.ts

@@ -1,11 +1,24 @@
 import { NextRequest, NextResponse } from 'next/server';
-import { getDatabaseUrl } from '~/lib/services/datasourceService';
+import { getDatabaseUrl } from '~/lib/services/dataSourceService';
 import { requireUserId } from '~/auth/session';
 
 export async function GET(request: NextRequest, { params }: { params: Promise<{ id: string }> }) {
   const userId = await requireUserId(request);
   const { id } = await params;
-  const url = await getDatabaseUrl(userId, id);
+  const { searchParams } = new URL(request.url);
+  const environmentId = searchParams.get('environmentId');
 
-  return NextResponse.json({ url, success: true });
+  if (!environmentId) {
+    return NextResponse.json({ success: false, error: 'Environment ID is required' }, { status: 400 });
+  }
+
+  try {
+    const url = await getDatabaseUrl(userId, id, environmentId);
+    return NextResponse.json({ url, success: true });
+  } catch (error) {
+    return NextResponse.json(
+      { success: false, error: error instanceof Error ? error.message : 'Failed to get database URL' },
+      { status: 404 },
+    );
+  }
 }