feat: purely in-isolate template unzip + inmemory storage

AshishKumar4 · AshishKumar4 · commit 3cd14ed58a4f · 2025-10-20T22:02:35.000-04:00
diff --git a/src/routes/chat/utils/handle-websocket-message.ts b/src/routes/chat/utils/handle-websocket-message.ts
@@ -152,10 +152,10 @@ export function createWebSocketMessageHandler(deps: HandleMessageDeps) {
                 }));
                 break;
             }
-            case 'cf_agent_state': {
-                const { state } = message;
-                logger.debug('🔄 Agent state update received:', state);
-
+            case 'agent_connected': {
+                const { state, templateDetails } = message;
+                console.log('Agent connected', state, templateDetails);
+                
                 if (!isInitialStateRestored) {
                     logger.debug('📥 Performing initial state restoration');
                     
@@ -168,11 +168,11 @@ export function createWebSocketMessageHandler(deps: HandleMessageDeps) {
                         setQuery(state.query);
                     }
 
-                    if (state.templateDetails?.allFiles && bootstrapFiles.length === 0) {
-                        const files = Object.entries(state.templateDetails.allFiles).map(([filePath, fileContents]) => ({
+                    if (templateDetails?.allFiles && bootstrapFiles.length === 0) {
+                        const files = Object.entries(templateDetails.allFiles).map(([filePath, fileContents]) => ({
                             filePath,
                             fileContents,
-                        }));
+                        })).filter((file) => templateDetails.importantFiles.includes(file.filePath));
                         logger.debug('📥 Restoring bootstrap files:', files);
                         loadBootstrapFiles(files);
                     }
@@ -254,6 +254,11 @@ export function createWebSocketMessageHandler(deps: HandleMessageDeps) {
                         sendWebSocketMessage(websocket, 'preview');
                     }
                 }
+                break;
+            }
+            case 'cf_agent_state': {
+                const { state } = message;
+                logger.debug('🔄 Agent state update received:', state);
 
                 if (state.shouldBeGenerating) {
                     logger.debug('🔄 shouldBeGenerating=true detected, auto-resuming generation');
diff --git a/worker/agents/core/simpleGeneratorAgent.ts b/worker/agents/core/simpleGeneratorAgent.ts
@@ -1,4 +1,4 @@
-import { Agent, AgentContext, Connection } from 'agents';
+import { Agent, AgentContext, Connection, ConnectionContext } from 'agents';
 import { 
     Blueprint, 
     PhaseConceptGenerationSchemaType, 
@@ -12,7 +12,7 @@ import {  GitHubExportResult } from '../../services/github/types';
 import { CodeGenState, CurrentDevState, MAX_PHASES } from './state';
 import { AllIssues, AgentSummary, AgentInitArgs, PhaseExecutionResult, UserContext } from './types';
 import { PREVIEW_EXPIRED_ERROR, WebSocketMessageResponses } from '../constants';
-import { broadcastToConnections, handleWebSocketClose, handleWebSocketMessage } from './websocket';
+import { broadcastToConnections, handleWebSocketClose, handleWebSocketMessage, sendToConnection } from './websocket';
 import { createObjectLogger, StructuredLogger } from '../../logger';
 import { ProjectSetupAssistant } from '../assistants/projectsetup';
 import { UserConversationProcessor, RenderToolCall } from '../operations/UserConversationProcessor';
@@ -82,6 +82,7 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
     protected deploymentManager!: DeploymentManager;
 
     private previewUrlCache: string = '';
+    private templateDetailsCache: TemplateDetails | null = null;
     
     // In-memory storage for user-uploaded images (not persisted in DO state)
     private pendingUserImages: ProcessedImageAttachment[] = []
@@ -131,7 +132,7 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
         generatedFilesMap: {},
         agentMode: 'deterministic',
         sandboxInstanceId: undefined,
-        templateDetails: {} as TemplateDetails,
+        templateName: '',
         commandsHistory: [],
         lastPackageJson: '',
         pendingUserInputs: [],
@@ -160,7 +161,7 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
         );
         
         // Initialize FileManager
-        this.fileManager = new FileManager(this.stateManager);
+        this.fileManager = new FileManager(this.stateManager, () => this.getTemplateDetails());
         
         // Initialize DeploymentManager first (manages sandbox client caching)
         // DeploymentManager will use its own getClient() override for caching
@@ -212,12 +213,14 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
         })
 
         const packageJson = templateInfo.templateDetails?.allFiles['package.json'];
+
+        this.templateDetailsCache = templateInfo.templateDetails;
         
         this.setState({
             ...this.initialState,
             query,
             blueprint,
-            templateDetails: templateInfo.templateDetails,
+            templateName: templateInfo.templateDetails.name,
             sandboxInstanceId: undefined,
             generatedPhases: [],
             commandsHistory: [],
@@ -253,7 +256,59 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
 
     async isInitialized() {
         return this.getAgentId() ? true : false
-    }  
+    }
+
+    async onStart(_props?: Record<string, unknown> | undefined): Promise<void> {
+        this.logger().info(`Agent ${this.getAgentId()} session: ${this.state.sessionId} onStart`);
+        // Ignore if agent not initialized
+        if (!this.isInitialized()) {
+            this.logger().info(`Agent ${this.getAgentId()} session: ${this.state.sessionId} not initialized, ignoring onStart`);
+            return;
+        }
+        this.logger().info(`Agent ${this.getAgentId()} session: ${this.state.sessionId} onStart being processed`);
+        // Fill the template cache
+        await this.ensureTemplateDetails();
+        this.logger().info(`Agent ${this.getAgentId()} session: ${this.state.sessionId} onStart processed successfully`);
+    }
+    
+    onStateUpdate(_state: CodeGenState, _source: "server" | Connection) {}
+
+    setState(state: CodeGenState): void {
+        try {
+            super.setState(state);
+        } catch (error) {
+            this.broadcastError("Error setting state", error);
+            this.logger().error("State details:", {
+                originalState: JSON.stringify(this.state, null, 2),
+                newState: JSON.stringify(state, null, 2)
+            });
+        }
+    }
+
+    onConnect(connection: Connection, ctx: ConnectionContext) {
+        this.logger().info(`Agent connected for agent ${this.getAgentId()}`, { connection, ctx });
+        sendToConnection(connection, 'agent_connected', {
+            state: this.state,
+            templateDetails: this.getTemplateDetails()
+        });
+    }
+
+    async ensureTemplateDetails() {
+        if (!this.templateDetailsCache) {
+            this.logger().info(`Template details being cached for template: ${this.state.templateName}`);
+            const results = await BaseSandboxService.getTemplateDetails(this.state.templateName);
+            if (!results.success || !results.templateDetails) {
+                throw new Error(`Failed to get template details for template: ${this.state.templateName}`);
+            }
+            this.templateDetailsCache = results.templateDetails;
+            this.logger().info(`Template details for template: ${this.state.templateName} cached successfully`);
+        }
+        return this.templateDetailsCache;
+    }
+
+    private getTemplateDetails() {
+        return this.templateDetailsCache!;
+    }
 
     /*
     * Each DO has 10 gb of sqlite storage. However, the way agents sdk works, it stores the 'state' object of the agent as a single row
@@ -333,20 +388,6 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
         });
         this.logger().info(`Agent initialized successfully for agent ${this.state.inferenceContext.agentId}`);
     }
-    
-    onStateUpdate(_state: CodeGenState, _source: "server" | Connection) {}
-
-    setState(state: CodeGenState): void {
-        try {
-            super.setState(state);
-        } catch (error) {
-            this.broadcastError("Error setting state", error);
-            this.logger().error("State details:", {
-                originalState: JSON.stringify(this.state, null, 2),
-                newState: JSON.stringify(state, null, 2)
-            });
-        }
-    }
 
     getPreviewUrlCache() {
         return this.previewUrlCache;
@@ -359,7 +400,7 @@ export class SimpleCodeGeneratorAgent extends Agent<Env, CodeGenState> {
                 agentId: this.getAgentId(),
                 query: this.state.query,
                 blueprint: this.state.blueprint,
-                template: this.state.templateDetails,
+                template: this.getTemplateDetails(),
                 inferenceContext: this.state.inferenceContext
             });
         }
diff --git a/worker/agents/core/state.ts b/worker/agents/core/state.ts
@@ -1,7 +1,6 @@
 import type { Blueprint, PhaseConceptType ,
     FileOutputType,
 } from '../schemas';
-import type { TemplateDetails } from '../../services/sandbox/sandboxTypes';
 // import type { ScreenshotData } from './types';
 import type { ConversationMessage } from '../inferutils/common';
 import type { InferenceContext } from '../inferutils/config.types';
@@ -36,7 +35,7 @@ export interface CodeGenState {
     generatedPhases: PhaseState[];
     commandsHistory?: string[]; // History of commands run
     lastPackageJson?: string; // Last package.json file contents
-    templateDetails: TemplateDetails;   // TODO: Remove this from state and rely on directly fetching from sandbox
+    templateName: string;
     sandboxInstanceId?: string;
     
     shouldBeGenerating: boolean; // Persistent flag indicating generation should be active
diff --git a/worker/agents/core/stateMigration.ts b/worker/agents/core/stateMigration.ts
@@ -1,10 +1,14 @@
 import { CodeGenState, FileState } from './state';
 import { StructuredLogger } from '../../logger';
+import { TemplateDetails } from 'worker/services/sandbox/sandboxTypes';
 
 export class StateMigration {
     static migrateIfNeeded(state: CodeGenState, logger: StructuredLogger): CodeGenState | null {
         let needsMigration = false;
         
+        //------------------------------------------------------------------------------------
+        // Migrate files from old schema
+        //------------------------------------------------------------------------------------
         const migrateFile = (file: any): any => {
             const hasOldFormat = 'file_path' in file || 'file_contents' in file || 'file_purpose' in file;
             
@@ -34,27 +38,9 @@ export class StateMigration {
             }
         }
 
-        let migratedTemplateDetails = state.templateDetails;
-        if ('files' in migratedTemplateDetails && migratedTemplateDetails?.files) {
-            const migratedTemplateFiles = (migratedTemplateDetails.files as Array<any>).map(file => {
-                const migratedFile = migrateFile(file);
-                return migratedFile;
-            });
-
-            const allFiles = migratedTemplateFiles.reduce((acc, file) => {
-                acc[file.filePath] = file;
-                return acc;
-            }, {} as Record<string, any>);
-
-            migratedTemplateDetails = {
-                ...migratedTemplateDetails,
-                allFiles
-            };
-
-            // Remove 'files' property
-            delete (migratedTemplateDetails as any).files;
-            needsMigration = true;
-        }
+        //------------------------------------------------------------------------------------
+        // Migrate conversations cleanups and internal memos
+        //------------------------------------------------------------------------------------
 
         let migratedConversationMessages = state.conversationMessages;
         const MIN_MESSAGES_FOR_CLEANUP = 25;
@@ -132,6 +118,9 @@ export class StateMigration {
             }
         }
 
+        //------------------------------------------------------------------------------------
+        // Migrate inference context from old schema
+        //------------------------------------------------------------------------------------
         let migratedInferenceContext = state.inferenceContext;
         if (migratedInferenceContext && 'userApiKeys' in migratedInferenceContext) {
             migratedInferenceContext = {
@@ -142,6 +131,9 @@ export class StateMigration {
             needsMigration = true;
         }
 
+        //------------------------------------------------------------------------------------
+        // Migrate deprecated props
+        //------------------------------------------------------------------------------------  
         const stateHasDeprecatedProps = 'latestScreenshot' in (state as any);
         if (stateHasDeprecatedProps) {
             needsMigration = true;
@@ -151,19 +143,29 @@ export class StateMigration {
         if (!stateHasProjectUpdatesAccumulator) {
             needsMigration = true;
         }
+
+        //------------------------------------------------------------------------------------
+        // Migrate Template Details -> remove template details and instead use template name
+        //------------------------------------------------------------------------------------
+        const hasTemplateDetails = 'templateDetails' in (state as any);
+        if (hasTemplateDetails) {
+            needsMigration = true;
+            const templateDetails = (state as any).templateDetails;
+            const templateName = (templateDetails as TemplateDetails).name;
+            delete (state as any).templateDetails;
+            (state as any).templateName = templateName;
+        }
         
         if (needsMigration) {
             logger.info('Migrating state: schema format, conversation cleanup, and security fixes', {
                 generatedFilesCount: Object.keys(migratedFilesMap).length,
-                templateFilesCount: migratedTemplateDetails?.allFiles?.length || 0,
                 finalConversationCount: migratedConversationMessages?.length || 0,
                 removedUserApiKeys: state.inferenceContext && 'userApiKeys' in state.inferenceContext
             });
             
             const newState = {
                 ...state,
                 generatedFilesMap: migratedFilesMap,
-                templateDetails: migratedTemplateDetails,
                 conversationMessages: migratedConversationMessages,
                 inferenceContext: migratedInferenceContext,
                 projectUpdatesAccumulator: []
diff --git a/worker/agents/services/implementations/DeploymentManager.ts b/worker/agents/services/implementations/DeploymentManager.ts
@@ -517,7 +517,7 @@ export class DeploymentManager extends BaseAgentService implements IDeploymentMa
      */
     private async createNewInstance(): Promise<BootstrapResponse | null> {
         const state = this.getState();
-        const templateName = state.templateDetails?.name || 'scratch';
+        const templateName = state.templateName || 'scratch';
         
         // Generate unique project name
         let prefix = (state.blueprint?.projectName || templateName)
@@ -532,7 +532,7 @@ export class DeploymentManager extends BaseAgentService implements IDeploymentMa
 
         // Add AI proxy vars if AI template
         let localEnvVars: Record<string, string> = {};
-        if (state.templateDetails?.name?.includes('agents')) {
+        if (state.templateName?.includes('agents')) {
             localEnvVars = {
                 "CF_AI_BASE_URL": generateAppProxyUrl(this.env),
                 "CF_AI_API_KEY": await generateAppProxyToken(
diff --git a/worker/agents/services/implementations/FileManager.ts b/worker/agents/services/implementations/FileManager.ts
@@ -2,17 +2,18 @@ import * as Diff from 'diff';
 import { IFileManager } from '../interfaces/IFileManager';
 import { IStateManager } from '../interfaces/IStateManager';
 import { FileOutputType } from '../../schemas';
-// import { TemplateDetails } from '../../../services/sandbox/sandboxTypes';
 import { FileProcessing } from '../../domain/pure/FileProcessing';
 import { FileState } from 'worker/agents/core/state';
+import { TemplateDetails } from '../../../services/sandbox/sandboxTypes';
 
 /**
  * Manages file operations for code generation
  * Handles both template and generated files
  */
 export class FileManager implements IFileManager {
     constructor(
-        private stateManager: IStateManager
+        private stateManager: IStateManager,
+        private getTemplateDetailsFunc: () => TemplateDetails
     ) {}
 
     getGeneratedFile(path: string): FileOutputType | null {
@@ -22,7 +23,7 @@ export class FileManager implements IFileManager {
 
     getAllFiles(): FileOutputType[] {
         const state = this.stateManager.getState();
-        return FileProcessing.getAllFiles(state.templateDetails, state.generatedFilesMap);
+        return FileProcessing.getAllFiles(this.getTemplateDetailsFunc(), state.generatedFilesMap);
     }
 
     saveGeneratedFile(file: FileOutputType): FileState {
@@ -86,6 +87,7 @@ export class FileManager implements IFileManager {
             generatedFilesMap: newFilesMap
         });
     }
+
     fileExists(path: string): boolean {
         return !!this.getGeneratedFile(path)
     }
diff --git a/worker/api/websocketTypes.ts b/worker/api/websocketTypes.ts
@@ -1,7 +1,7 @@
 import type { CodeReviewOutputType, FileConceptType, FileOutputType } from "../agents/schemas";
 import type { CodeGenState } from "../agents/core/state";
 import type { ConversationState } from "../agents/inferutils/common";
-import type { CodeIssue, RuntimeError, StaticAnalysisResponse } from "../services/sandbox/sandboxTypes";
+import type { CodeIssue, RuntimeError, StaticAnalysisResponse, TemplateDetails } from "../services/sandbox/sandboxTypes";
 import type { CodeFixResult } from "../services/code-fixer";
 import { IssueReport } from "../agents/domain/values/IssueReport";
 import type { RateLimitExceededError } from 'shared/types/errors';
@@ -16,6 +16,12 @@ type StateMessage = {
 	state: CodeGenState;
 };
 
+type AgentConnectedMessage = {
+    type: 'agent_connected';
+    state: CodeGenState;
+    templateDetails: TemplateDetails;
+};
+
 type ConversationStateMessage = {
     type: 'conversation_state';
     state: ConversationState;
@@ -408,6 +414,7 @@ type ServerLogMessage = {
 
 export type WebSocketMessage =
 	| StateMessage
+	| AgentConnectedMessage
 	| ConversationStateMessage
 	| GenerationStartedMessage
 	| FileGeneratingMessage
diff --git a/worker/services/sandbox/BaseSandboxService.ts b/worker/services/sandbox/BaseSandboxService.ts
