use own db

ai-service/package.json

@@ -4,7 +4,7 @@
   "description": "",
   "main": "index.js",
   "scripts": {
-    "dev": "tsx watch src/server.ts",
+    "dev": "tsx watch src/server.ts --include \"./src/**/*\"",
     "test": "echo \"Error: no test specified\" && exit 1",
     "build": "rimraf dist && tsup",
     "start": "node dist/server.mjs"

ai-service/prisma/schema.prisma

@@ -16,11 +16,24 @@ datasource db {
 
 model Conversation {
   artifactId String
-  projectId String
-  id        String   @id @default(uuid())
+  id        String   @id
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
   title String
-  uiMessages  Json
-  modelMessages Json
+  messages Message[]
+  sandboxId String?
+}
+
+enum MessageRole {
+  user
+  assistant
+}
+
+model Message {
+  id String @id
+  role MessageRole
+  parts Json
+  conversationId String
+  createdAt DateTime @default(now())
+  conversation Conversation @relation(fields: [conversationId], references: [id])
 }

ai-service/src/handlers/chat/route copy.ts

@@ -1,226 +0,0 @@
-import { Context } from 'hono';
-import { chatRequestBodySchema, ChatRequestBodyType } from './types';
-import { z } from 'zod';
-import {
-    convertToModelMessages,
-    createUIMessageStream,
-    createUIMessageStreamResponse,
-    streamText,
-    TextPart
-} from 'ai';
-import { createRuntimeContext, WriterType } from '../../lib/ai/mastra/utils/runtime-context';
-import { mastra } from '../../lib/ai/mastra';
-import { createImagineClient } from '@/lib/imagine/create-artifact-client';
-import { anthropic } from '@ai-sdk/anthropic';
-import { Workspace } from '@/lib/imagine/workspaces-api-client';
-import { AppwriteException } from '@appwrite.io/console';
-import { createSynapseClient } from '@/lib/synapse-http-client';
-
-export const handleChatRequest = async (c: Context) => {
-    c.res.headers.set('x-vercel-ai-ui-message-stream', 'v1');
-    const signal = c.req.raw.signal;
-
-    const token = c.req.header('X-Imagine-Token');
-
-    console.log('token', token);
-
-    if (!token) {
-        return c.json({ error: 'Unauthorized' }, 401);
-    }
-
-    let body: ChatRequestBodyType;
-
-    /** Create workspace */
-
-    // Parse request body
-    try {
-        const json = await c.req.json();
-        body = chatRequestBodySchema.parse(json);
-    } catch (error) {
-        if (error instanceof z.ZodError) {
-            return c.json({ errors: error.issues }, 400);
-        }
-
-        console.error('Error', error);
-        return c.json('An unknown error occurred', 500);
-    }
-
-    const { id: conversationId, messages, trigger, artifactId, projectId } = body;
-
-    if (trigger === 'submit-tool-result') {
-        // save to file
-        // return c.body(null, 200);
-        return streamText({
-            model: anthropic('claude-3-7-sonnet-20250219'),
-            messages: convertToModelMessages(messages)
-        }).toUIMessageStreamResponse();
-    }
-
-    const { imagineClient, workspacesClient } = await createImagineClient({
-        projectId,
-        token // TODO: use the token from the request
-    });
-
-    const imagineConvo = await imagineClient.getConversation(artifactId, conversationId);
-
-    if (!imagineConvo) {
-        throw new Error('Conversation not found');
-    }
-
-    // Create workspace
-    let workspace: Workspace;
-    const workspaceUrl = `${process.env.WORKSPACE_URL_PROTOCOL}://${artifactId}.${process.env.WORKSPACE_URL_DOMAIN}:${process.env.WORKSPACE_URL_PORT}`;
-    console.log('workspaceUrl', workspaceUrl);
-
-    const synapseClient = createSynapseClient({
-        artifactId
-    });
-    
-
-    try {
-        console.log('Getting workspace');
-        workspace = await workspacesClient.get(artifactId);
-        console.log('Found existing workspace', workspace);
-    } catch (error) {
-        if ((error as AppwriteException).type === 'workspace_not_found') {
-            console.log('Workspace not found, creating...');
-            workspace = await workspacesClient.create(artifactId, artifactId, "s-1vcpu-1gb");
-            console.log('Created new workspace', workspace);
-            console.log('Creating proxy rule');
-            console.time("createWorkspaceProxyRule");
-            const proxyRule = await workspacesClient.createWorkspaceProxyRule(
-                `${artifactId}.${process.env.WORKSPACE_URL_DOMAIN}`,
-                workspace.$id
-            );
-
-            await new Promise((resolve) => setTimeout(resolve, 3000));
-            console.timeEnd("createWorkspaceProxyRule");
-            console.log('Created proxy rule', proxyRule);
-
-            console.log("Creating artifact directory");
-            await synapseClient.executeCommand({
-                command: "mkdir -p artifact",
-                cwd: "/usr/local"
-            });
-            
-            console.log("Cloning template")
-            await synapseClient.executeCommand({
-                command: "bunx giget@latest gh:appwrite/templates-for-frameworks/base-vite-template .",
-                cwd: "/usr/local/artifact",
-                timeout: 60000,
-            });
-        
-            console.log("Installing dependencies")
-            await synapseClient.executeCommand({
-                command: "bun install",
-                cwd: "/usr/local/artifact",
-                timeout: 60000,
-            });
-        
-            console.log("Running bun dev in the background")
-            await synapseClient.startBackgroundProcess({
-                command: "bun",
-                args: ["run", "dev"],
-                cwd: "/usr/local/artifact", 
-            });
-        
-        } else {
-            throw error;
-        }
-    }
-
-
-    const convertedMessages = convertToModelMessages(messages);
-
-    const latestMessage = convertedMessages[convertedMessages.length - 1];
-
-    const latestMessageTextPart = latestMessage.content[0] as TextPart;
-    const restMessages = convertedMessages.slice(0, -1);
-    const isNewConversation = restMessages.length === 0;
-
-    let didError = false;
-
-    const stream = createUIMessageStream({
-        originalMessages: messages,
-        execute: async (params) => {
-            const writer = params.writer as WriterType;
-            const runtimeContext = createRuntimeContext({
-                writer,
-                artifactId,
-                restMessages,
-                isFirstMessage: isNewConversation,
-                signal
-            });
-
-            writer.write({
-               type: "data-workspace-state",
-               data: {
-                state: "ready",
-                workspaceUrl,
-               },
-               transient: true
-            })
-
-            c.set('runtimeContext', runtimeContext);
-            const run = await mastra.getWorkflow('codeWorkflow').createRunAsync();
-
-            const result = run.stream({
-                inputData: {
-                    userPrompt: latestMessageTextPart.text
-                },
-                runtimeContext
-            });
-
-            writer.write({
-                type: 'start'
-            });
-
-            writer.write({
-                type: 'start-step'
-            });
-
-            for await (const chunk of result.stream) {
-                // We must await the stream
-            }
-
-            writer.write({
-                type: 'finish-step'
-            });
-
-            writer.write({
-                type: 'finish'
-            });
-        },
-        onFinish: async (event) => {
-            console.log('onFinish', { didError });
-            if (didError) {
-                return;
-            }
-
-            const { messages } = event;
-
-            // The last message should NEVER be by the user. If that's the case, remove it.
-            const lastMessage = messages[messages.length - 1];
-
-            console.log('lastMessage', JSON.stringify(lastMessage, null, 2));
-
-            console.log('Saving messages to imagine');
-            await imagineClient.updateConversation(
-                artifactId,
-                imagineConvo.$id,
-                'Test Conversation',
-                messages
-            );
-            console.log('Messages saved to imagine');
-        },
-        onError: (error) => {
-            didError = true;
-            console.error('onError', error);
-            return (error as Error).message ?? 'An error occurred';
-        }
-    });
-
-    return createUIMessageStreamResponse({
-        stream
-    });
-};

ai-service/src/handlers/chat/route.ts

@@ -18,7 +18,22 @@ import { createSynapseClient } from '@/lib/synapse-http-client';
 import { daytona } from '@/lib/daytona-client';
 import { Sandbox } from '@daytonaio/sdk';
 import { getOrCreateArtifactSandbox, startDevServer } from '@/lib/daytona-utils';
-import { OnStepUpdateFn, WorkspaceStepId, workspaceStepSchema } from '@/lib/ai/custom-parts/workspace-state';
+import {
+    OnStepUpdateFn,
+    workspaceStepSchema
+} from '@/lib/ai/custom-parts/workspace-state';
+import {
+    getOrCreateConversation,
+    updateConversationHistory
+} from '@/lib/message-history';
+
+export const routeHandler = async (c: Context) => {
+    try {
+        return handleChatRequest(c);
+    } catch (e) {
+        console.error('Error', e);
+    }
+};
 
 export const handleChatRequest = async (c: Context) => {
     c.res.headers.set('x-vercel-ai-ui-message-stream', 'v1');
@@ -63,9 +78,12 @@ export const handleChatRequest = async (c: Context) => {
         token // TODO: use the token from the request
     });
 
-    const imagineConvo = await imagineClient.getConversation(artifactId, conversationId);
+    const convo = await getOrCreateConversation({
+        conversationId,
+        artifactId: conversationId,
+    });
 
-    if (!imagineConvo) {
+    if (!convo) {
         throw new Error('Conversation not found');
     }
 
@@ -100,7 +118,7 @@ export const handleChatRequest = async (c: Context) => {
                 id,
                 status,
                 text
-                }: Parameters<OnStepUpdateFn>[0]) => {
+            }: Parameters<OnStepUpdateFn>[0]) => {
                 const found = steps.find((step) => step.id === id);
 
                 if (found) {
@@ -111,13 +129,13 @@ export const handleChatRequest = async (c: Context) => {
                 }
 
                 writer.write({
-                    type: "data-workspace-state",
+                    type: 'data-workspace-state',
                     data: {
-                        state: "in-progress",
+                        state: 'in-progress',
                         steps,
                         workspaceUrl: null
                     }
-                })
+                });
             };
 
             const { sandbox } = await getOrCreateArtifactSandbox({
@@ -141,16 +159,16 @@ export const handleChatRequest = async (c: Context) => {
                 sandbox
             });
 
-            console.log("Reporting compelted");
+            console.log('Reporting compelted');
 
             writer.write({
-                type: "data-workspace-state",
+                type: 'data-workspace-state',
                 data: {
-                    state: "completed",
+                    state: 'completed',
                     steps,
                     workspaceUrl
                 }
-            })
+            });
 
             c.set('runtimeContext', runtimeContext);
             const run = await mastra.getWorkflow('codeWorkflow').createRunAsync();
@@ -193,15 +211,13 @@ export const handleChatRequest = async (c: Context) => {
             // The last message should NEVER be by the user. If that's the case, remove it.
             const lastMessage = messages[messages.length - 1];
 
-            console.log('lastMessage', JSON.stringify(lastMessage, null, 2));
+            console.log('Saving messages');
+
+            await updateConversationHistory({
+                conversationId,
+                messages
+            });
 
-            console.log('Saving messages to imagine');
-            // await imagineClient.updateConversation(
-            //     artifactId,
-            //     imagineConvo.$id,
-            //     'Test Conversation',
-            //     messages
-            // );
             console.log('Messages saved to imagine');
         },
         onError: (error) => {

ai-service/src/handlers/conversation.ts

@@ -1,52 +1,78 @@
-import { createImagineClient } from "@/lib/imagine/create-artifact-client";
-import { Conversation } from "@/lib/imagine/imagine-api-client";
+import { daytona } from "@/lib/daytona-client";
+import { getArtifactSandbox } from "@/lib/daytona-utils";
+import { convertDBMessagesToImagineUIMessages, getOrCreateConversation } from "@/lib/message-history";
+import { prisma } from "@/lib/prisma";
+import { ImagineUIMessage } from "@/shared-types";
 import { Context } from "hono";
 
-const mapConversation = (conversation: Conversation) => {
-  return {
-    id: conversation.$id,
-    name: conversation.name,
-    messages: conversation.messages,
-  };
-};
+export type GetConversationResult = {
+  id: string;
+  title: string;
+  messages: ImagineUIMessage[];
+  createdAt: Date;
+  updatedAt: Date;
+  artifactId: string;
+  previewUrl: string | null;
+}
+
+export type GetConversationsResult = Omit<GetConversationResult, "messages" | "previewUrl"> & {};
 
 export const getConversation = async (c: Context) => {
   const { conversationId } = c.req.param();
   const token = c.req.header("X-Imagine-Token");
-
+  let previewUrl: string | null = null;
+  
   if (!token) {
     return c.json({ error: "Unauthorized" }, 401);
   }
 
-  const projectId = process.env.IMAGINE_PROJECT_ID!;
-  const artifactId = process.env.IMAGINE_ARTIFACT_ID!;
-
-  const { imagineClient } = await createImagineClient({
-    projectId,
-    token
+  const conversation = await getOrCreateConversation({
+    conversationId,
+    artifactId: conversationId,
   });
 
-  const conversation = await imagineClient.getConversation(artifactId, conversationId);
-  console.log("conversation", conversation);
-  return c.json(mapConversation(conversation));
+  // Check if sandbox is running
+  const sandbox = await getArtifactSandbox({ artifactId: conversationId });
+
+  if (sandbox) {
+    const previewLinkResult = await sandbox.getPreviewLink(3000);
+    previewUrl = previewLinkResult.url;
+  }
+
+  const messages = convertDBMessagesToImagineUIMessages(conversation.messages)
+  
+  const result: GetConversationResult = {
+    id: conversation.id,
+    title: conversation.title,
+    messages,
+    createdAt: conversation.createdAt,
+    updatedAt: conversation.updatedAt,
+    artifactId: conversationId,
+    previewUrl,
+  }
+  return c.json(result);
 }
 
 export const getConversations = async (c: Context) => {
-  console.log("getConversations");
-  const projectId = process.env.IMAGINE_PROJECT_ID!;
-  const artifactId = process.env.IMAGINE_ARTIFACT_ID!;
-  const token = c.req.header("X-Imagine-Token");
+  const artifactId = c.req.param("artifactId");
 
-  if (!token) {
-    return c.json({ error: "Unauthorized" }, 401);
+  if (!artifactId) {
+    return c.json({ error: "artifactId is required" }, 400);
   }
 
-  const { imagineClient } = await createImagineClient({
-    projectId,
-    token
+  const conversations = await prisma.conversation.findMany({
+    where: {
+      artifactId,
+    },
   });
 
-  const conversations = await imagineClient.listConversations(artifactId);
-  console.log("conversations", conversations);
-  return c.json(conversations.conversations.map(mapConversation));
+  const result: GetConversationsResult[] = conversations.map((conversation) => ({
+    id: conversation.id,
+    title: conversation.title,
+    createdAt: conversation.createdAt,
+    updatedAt: conversation.updatedAt,
+    artifactId,
+  }));
+
+  return c.json(result);
 }

ai-service/src/lib/Dockerfile-daytona

@@ -0,0 +1,16 @@
+# Use the Bun image as the base image
+FROM daytonaio/sandbox:0.4.3
+
+# Set the working directory in the container
+WORKDIR /home/daytona/workspace
+
+RUN npm -v 
+RUN npm install -g bun
+RUN bunx giget --version
+
+RUN bunx giget@latest gh:appwrite/templates-for-frameworks/base-vite-template .
+RUN bun install
+
+EXPOSE 3000
+
+ENTRYPOINT ["sleep", "infinity"]

ai-service/src/lib/daytona-image.ts

@@ -0,0 +1,29 @@
+import { Image } from '@daytonaio/sdk';
+import { daytona } from './daytona-client';
+
+export async function buildImage() {
+    const image = Image.base("daytonaio/sandbox")
+        .workdir('/home/daytona/artifact')
+        .runCommands('npm -v');
+
+    console.log(' image', image);
+
+    const snapshot = await daytona.snapshot.create(
+        {
+            name: 'imgn-vite:v3',
+            image,
+            resources: {
+                cpu: 1,
+                memory: 1,
+                disk: 3
+            }
+        },
+        {
+            onLogs: console.log
+        }
+    );
+
+    console.log('snapshot', snapshot);
+}
+
+buildImage();

ai-service/src/lib/daytona-utils.ts

@@ -3,6 +3,7 @@ import { daytona } from './daytona-client';
 import { DaytonaNotFoundError } from '@daytonaio/sdk/src/errors/DaytonaError';
 import { WriterType } from './ai/mastra/utils/runtime-context';
 import { OnStepUpdateFn, WorkspaceStepId } from './ai/custom-parts/workspace-state';
+import { prisma } from './prisma';
 
 const sandboxId = 'b62d2a47-2f7c-4367-aa53-4980f3fe6627';
 const baseDir = '/home/daytona/workspace';
@@ -196,10 +197,11 @@ const createSandbox = async ({ artifactId }: { artifactId: string }) => {
     let sandbox: Sandbox;
     try {
         sandbox = await daytona.create({
+            snapshot: "imgn-base-vite:v2",
             labels: { artifactId },
-            autoStopInterval: 600, // 10 minutes
-            autoDeleteInterval: 1200, // 20 minutes
-            language: 'typescript',
+            autoStopInterval: 10, // 10 minutes
+            autoDeleteInterval: 60, // 20 minutes
+            // language: 'typescript',
             public: true
         });
 
@@ -225,6 +227,20 @@ const createSandbox = async ({ artifactId }: { artifactId: string }) => {
     }
 };
 
+export const getArtifactSandbox = async ({ artifactId }: { artifactId: string }): Promise<Sandbox | null> => {
+    const sandboxes = await daytona.list({
+        artifactId
+    });
+
+    const sandbox = sandboxes[0];
+
+    if (!sandbox) {
+        return null;
+    }
+
+    return sandbox;
+}
+
 export const getOrCreateArtifactSandbox = async ({
     artifactId,
     onStepUpdate
@@ -232,13 +248,8 @@ export const getOrCreateArtifactSandbox = async ({
     artifactId: string;
     onStepUpdate: OnStepUpdateFn;
 }): Promise<{ sandbox: Sandbox }> => {
-    const sandboxes = await daytona.list({
-        artifactId
-    });
-
-    const existingSandbox = sandboxes[0];
-
     let sandbox: Sandbox;
+    const existingSandbox = await getArtifactSandbox({ artifactId });
 
     if (existingSandbox) {
         onStepUpdate({
@@ -246,12 +257,16 @@ export const getOrCreateArtifactSandbox = async ({
             status: 'in-progress',
             text: 'Getting existing workspace...'
         });
-        sandbox = await daytona.get(existingSandbox.id);
-        onStepUpdate({
-            id: WorkspaceStepId.CREATE_SANDBOX,
-            status: 'completed',
-            text: 'Workspace found'
-        });
+
+        const foundSandbox = await daytona.get(existingSandbox.id);
+
+        if (foundSandbox.state === "stopped") {
+            console.log('Sandbox is stopped, starting...');
+            await foundSandbox.start();
+            console.log('Sandbox started');
+        }
+
+        sandbox = foundSandbox;
     } else {
         console.log('Workspace not found, creating...');
 
@@ -265,63 +280,26 @@ export const getOrCreateArtifactSandbox = async ({
             artifactId
         });
 
+        // Save sandbox id to db
+        await prisma.conversation.update({
+            where: {
+                id: artifactId
+            },
+            data: {
+                sandboxId: sandbox.id
+            }
+        });
+
         const cwd = `/home/daytona/workspace`;
-
         console.log('Created sandbox', sandbox);
-
-        console.log('Creating artifact directory');
-        await sandbox.fs.createFolder(cwd, '755');
-
-        onStepUpdate({
-            id: WorkspaceStepId.CREATE_SANDBOX,
-            status: 'completed',
-            text: 'Workspace created'
-        });
-
-        onStepUpdate({
-            id: WorkspaceStepId.REPOSITORY_SETUP,
-            status: 'in-progress',
-            text: 'Setting up repository...'
-        });
-
-        const ls = await sandbox.process.executeCommand('npm install -g bun', cwd, {}, 30 * 1000);
-
-        console.log('ls', ls);
-
-        console.log('Cloning template');
-        const clone = await sandbox.process.executeCommand(
-            'bunx giget@latest gh:appwrite/templates-for-frameworks/base-vite-template .',
-            cwd,
-            {},
-            30
-        );
-
-        console.log('Template cloned', clone);
-
-        onStepUpdate({
-            id: WorkspaceStepId.REPOSITORY_SETUP,
-            status: 'completed',
-            text: 'Repository set up'
-        });
-
-        onStepUpdate({
-            id: WorkspaceStepId.INSTALL_DEPENDENCIES,
-            status: 'in-progress',
-            text: 'Installing dependencies...'
-        });
-
-        console.log('Installing dependencies');
-        const install = await sandbox.process.executeCommand('bun install', cwd, {}, 60);
-
-        console.log('Dependencies installed', install);
-
-        onStepUpdate({
-            id: WorkspaceStepId.INSTALL_DEPENDENCIES,
-            status: 'completed',
-            text: 'Dependencies installed'
-        });
     }
 
+    onStepUpdate({
+      id: WorkspaceStepId.CREATE_SANDBOX,
+      status: 'completed',
+      text: 'Workspace found'
+  });
+
     return { sandbox };
 };
 
@@ -377,10 +355,6 @@ export async function startDevServer({
         text: 'Dev server started'
     });
 
-    // Wait 3 seconds
-    await new Promise((resolve) => setTimeout(resolve, 3000));
-    // Read logs
-
     return {
         success: true,
         previewUrl: previewLink.url

ai-service/src/lib/message-history.ts

@@ -0,0 +1,123 @@
+import { ImagineUIMessage } from "@/shared-types";
+import { Message, Prisma } from "./generated/prisma";
+import { prisma } from "./prisma";
+
+type ConversationWithMessages = Prisma.ConversationGetPayload<{
+  include: {
+    messages: true;
+  };
+}>;
+
+export function convertDBMessagesToImagineUIMessages(messages: Message[]): ImagineUIMessage[] {
+  return messages.map((message) => ({
+    id: message.id,
+    role: message.role === "user" ? "user" : "assistant",
+    parts: message.parts as any,
+  }));
+}
+
+export async function getOrCreateConversation({
+  conversationId,
+  artifactId,
+}: {
+  conversationId: string;
+  artifactId: string;
+}): Promise<ConversationWithMessages> {
+  const conversation = await prisma.conversation.findFirst({
+    where: {
+      id: conversationId,
+      artifactId,
+    },
+    include: {
+      messages: {
+        orderBy: {
+          createdAt: "asc",
+        },
+      }
+    }
+  });
+
+  if (!conversation) {
+    return await prisma.conversation.create({
+      data: {
+        id: conversationId,
+        artifactId,
+        title: "New Conversation",
+      },
+      include: {
+        messages: true,
+      }
+    });
+  }
+
+  return conversation;
+}
+
+export async function updateConversationHistory({
+  conversationId,
+  messages,
+}: {
+  conversationId: string;
+  messages: ImagineUIMessage[];
+}) {
+  const convo = await prisma.conversation.findFirst({
+    where: {
+      id: conversationId,
+    },
+    include: {
+      messages: true,
+    },
+  });
+  
+  if (!convo) {
+    throw new Error("Conversation not found");
+  }
+
+  const incomingMessageIds = messages.map(m => m.id);
+  const existingMessages = convo.messages.filter(m => incomingMessageIds.includes(m.id));
+  const existingMessageIds = new Set(existingMessages.map(m => m.id));
+
+  const messagesToCreate = messages
+    .filter(m => !existingMessageIds.has(m.id))
+    .map((message) => ({
+      id: message.id,
+      role: message.role === "user" ? "user" as const : "assistant" as const,
+      parts: message.parts,
+      conversationId,
+      createdAt: new Date(),
+    }));
+
+  const messagesToUpdate = messages
+    .filter(m => existingMessageIds.has(m.id))
+    .map((message) => ({
+      id: message.id,
+      role: message.role === "user" ? "user" as const : "assistant" as const,
+      parts: message.parts,
+    }));
+
+  // Use bulk operations for much better performance
+  await prisma.$transaction(async (tx) => {
+    // Bulk create new messages
+    if (messagesToCreate.length > 0) {
+      await tx.message.createMany({
+        data: messagesToCreate,
+        skipDuplicates: true
+      });
+    }
+
+    // Bulk update existing messages
+    if (messagesToUpdate.length > 0) {
+      await Promise.all(
+        messagesToUpdate.map(message =>
+          tx.message.update({
+            where: { id: message.id },
+            data: {
+              role: message.role,
+              parts: message.parts,
+            }
+          })
+        )
+      );
+    }
+  });
+}

ai-service/src/shared-types/index.ts

@@ -3,6 +3,7 @@ import { thinkingUIDataPartSchema } from '@/lib/ai/custom-parts/thinking';
 import { workspaceStateUIDataPartSchema } from '@/lib/ai/custom-parts/workspace-state';
 import { fileTools } from '@/lib/ai/mastra/tools/file-tools';
 import { InferUIDataParts, InferUITool, ToolUIPart, UIMessage } from 'ai';
+export type { GetConversationResult, GetConversationsResult } from '@/handlers/conversation';
 
 export type ImagineTools = {
     readFile: InferUITool<typeof fileTools.readFileTool>;
@@ -20,4 +21,4 @@ export type ImagineUIDataParts = InferUIDataParts<{
     "workspace-state": typeof workspaceStateUIDataPartSchema;
 }>;
 
-export type ImagineUIMessage = UIMessage<never, ImagineUIDataParts, ImagineTools>;
+export type ImagineUIMessage = UIMessage<never, ImagineUIDataParts, ImagineTools>;

src/lib/components/studio/chat/chat.svelte

@@ -36,7 +36,7 @@
 
     const chat = new Chat<ImagineUIMessage>({
         maxSteps: 20,
-        id: $conversation.data.$id,
+        id: $conversation.data.id,
         transport: new DefaultChatTransport({
             api: `${VARS.AI_SERVICE_BASE_URL}/api/chat`,
         }),
@@ -93,7 +93,7 @@
             text: message,
         }, {
             body: {
-                id: $conversation.data?.$id,
+                id: $conversation.data.id,
                 projectId: page.params.project,
                 artifactId: page.params.artifact,
             },

src/lib/components/studio/chatWrapper.svelte

@@ -1,16 +1,19 @@
 <script lang="ts">
     import Chat from '$lib/components/studio/chat/chat.svelte';
     import { page } from '$app/state';
-    import { conversation, showChat } from '$lib/stores/chat.js';
+    import { conversation, showChat, workspaceState } from '$lib/stores/chat.js';
     import {
         disableBodySelect,
         enabledBodySelect,
         getChatWidthFromPrefs,
         saveImagineProjectPrefs
     } from '$lib/helpers/studioLayout.js';
-    import { sdk } from '$lib/stores/sdk.js';
+
     import { isSmallViewport } from '$lib/stores/viewport';
     import { previewFrameRef } from '$routes/(console)/project-[region]-[project]/store';
+    import { VARS } from '$lib/system';
+    import type { Conversation } from '$lib/sdk/imagine';
+    import { SvelteURL } from 'svelte/reactivity';
 
     $effect(() => {
         if ($isSmallViewport || page.params.artifact) {
@@ -93,24 +96,38 @@
     }
 
     async function getOrCreateConversation(artifactId: string) {
-        const { conversations } = await sdk
-            .forProject(page.params.region, page.params.project)
-            .imagine.listConversations(artifactId);
-        if (conversations.length === 0) {
-            const convo = await sdk
-                .forProject(page.params.region, page.params.project)
-                .imagine.createConversation(artifactId, `Conversation ${new Date().getTime()}`);
-            conversation.set(convo);
-        } else {
-            conversation.set(conversations[0]);
+        try {
+            const response = await fetch(`${VARS.AI_SERVICE_BASE_URL}/api/conversations/${artifactId}`, {
+                headers: {
+                    "X-Imagine-Token": "test",
+                    "Content-Type": "application/json",
+                }
+            });
+            const data = (await response.json()) as Conversation;
+            conversation.set(data);
+
+            if (data.previewUrl) {
+                workspaceState.set({
+                    state: "completed",
+                    workspaceUrl: new SvelteURL(data.previewUrl),
+                    steps: [],
+                });
+            }
+
+            return data;
+        } catch (error) {
+            console.error(`Could not get conversation for artifact ${artifactId}`, error);
         }
     }
 
+    let lastArtifactId: string | undefined = $state();
+
     $effect(() => {
-        if (page.params.artifact && $conversation.data?.artifactId !== page.params.artifact) {
-            getOrCreateConversation(page.params.artifact);
-        } else {
-            console.log('conversation', $conversation.data);
+        const artifactId = page.params.artifact;
+        
+        if (artifactId && artifactId !== lastArtifactId) {
+            lastArtifactId = artifactId;
+            getOrCreateConversation(artifactId);
         }
     });
 </script>

src/lib/sdk/imagine.ts

@@ -1,4 +1,4 @@
-import type { ImagineUIMessage } from '$shared-types';
+import type { GetConversationResult, ImagineUIMessage } from '$shared-types';
 import { AppwriteException } from '@appwrite.io/console';
 import type { Client, Payload } from '@appwrite.io/console';
 
@@ -534,39 +534,42 @@ export type Artifact = {
      */
     status: string;
 };
-/**
- * Conversation
- */
-export type Conversation = {
-    /**
-     * Conversation unique ID.
-     */
-    $id: string;
-    /**
-     * Conversation creation date in ISO 8601 format.
-     */
-    $createdAt: string;
-    /**
-     * Conversation update date in ISO 8601 format.
-     */
-    $updatedAt: string;
-    /**
-     * ID of the artifact this conversation belongs to.
-     */
-    artifactId: string;
-    /**
-     * Conversation name.
-     */
-    name: string;
-    /**
-     * Search index string.
-     */
-    search: string;
-    /**
-     * Messages (UI)
-     */
-    messages: ImagineUIMessage[];
-};
+
+export type Conversation = GetConversationResult;
+
+// /**
+//  * Conversation
+//  */
+// export type Conversation = {
+//     /**
+//      * Conversation unique ID.
+//      */
+//     $id: string;
+//     /**
+//      * Conversation creation date in ISO 8601 format.
+//      */
+//     $createdAt: string;
+//     /**
+//      * Conversation update date in ISO 8601 format.
+//      */
+//     $updatedAt: string;
+//     /**
+//      * ID of the artifact this conversation belongs to.
+//      */
+//     artifactId: string;
+//     /**
+//      * Conversation name.
+//      */
+//     name: string;
+//     /**
+//      * Search index string.
+//      */
+//     search: string;
+//     /**
+//      * Messages (UI)
+//      */
+//     messages: ImagineUIMessage[];
+// };
 /**
  * ConversationsMessage
  */

src/lib/stores/chat.ts

@@ -8,7 +8,6 @@ export const showChat = writable(false);
 export const showPrompt = writable(false);
 export const conversation = asyncWritable<Conversation>();
 
-
 export type WorkspaceState = Omit<WorkspaceStateUIDataPart["data"], "workspaceUrl"> & {
     workspaceUrl: SvelteURL | null;
 }