feat: Enhance session management with prompt tracking

- Added prompt_number to observations and session summaries for better tracking.
- Implemented prompt counter in SDK sessions to manage user prompts effectively.
- Updated database schema to include prompt tracking columns and removed unique constraints on session summaries.
- Modified hooks to utilize prompt_number in observations and summaries.
- Changed worker service to handle summarize requests instead of finalize, keeping the SDK agent active.
- Improved logging for better debugging and tracking of prompt numbers across sessions.

src/hooks/context.ts

@@ -42,10 +42,10 @@ export function contextHook(input?: SessionStartInput): void {
       output.push('');
 
       // Group observations by type
-      const byType: Record<string, Array<{text: string; created_at: string}>> = {};
+      const byType: Record<string, Array<{text: string; prompt_number: number | null; created_at: string}>> = {};
       for (const obs of observations) {
         if (!byType[obs.type]) byType[obs.type] = [];
-        byType[obs.type].push({ text: obs.text, created_at: obs.created_at });
+        byType[obs.type].push({ text: obs.text, prompt_number: obs.prompt_number, created_at: obs.created_at });
       }
 
       // Display each type
@@ -54,7 +54,8 @@ export function contextHook(input?: SessionStartInput): void {
         if (byType[type] && byType[type].length > 0) {
           output.push(`### ${type.charAt(0).toUpperCase() + type.slice(1)}s`);
           for (const obs of byType[type]) {
-            output.push(`- ${obs.text}`);
+            const promptLabel = obs.prompt_number ? ` (prompt #${obs.prompt_number})` : '';
+            output.push(`- ${obs.text}${promptLabel}`);
           }
           output.push('');
         }
@@ -76,6 +77,10 @@ export function contextHook(input?: SessionStartInput): void {
       output.push('---');
       output.push('');
 
+      const promptLabel = summary.prompt_number ? ` (Prompt #${summary.prompt_number})` : '';
+      output.push(`**Summary${promptLabel}**`);
+      output.push('');
+
       if (summary.request) {
         output.push(`**Request:** ${summary.request}`);
       }

src/hooks/new.ts

@@ -48,26 +48,31 @@ export async function newHook(input?: UserPromptSubmitInput): Promise<void> {
     // Check for any existing session (active, failed, or completed)
     let existing = db.findActiveSDKSession(session_id);
     let sessionDbId: number;
+    let isNewSession = false;
 
     if (existing) {
-      // Session already active, just continue
+      // Session already active, increment prompt counter
       sessionDbId = existing.id;
-      console.log(createHookResponse('UserPromptSubmit', true));
-      return;
-    }
-
-    // Check for inactive sessions we can reuse
-    const inactive = db.findAnySDKSession(session_id);
-
-    if (inactive) {
-      // Reactivate the existing session
-      sessionDbId = inactive.id;
-      db.reactivateSession(sessionDbId, prompt);
-      console.error(`[new-hook] Reactivated session ${sessionDbId} for Claude session ${session_id}`);
+      const promptNumber = db.incrementPromptCounter(sessionDbId);
+      console.error(`[new-hook] Continuing session ${sessionDbId}, prompt #${promptNumber}`);
     } else {
-      // Create new session
-      sessionDbId = db.createSDKSession(session_id, project, prompt);
-      console.error(`[new-hook] Created new session ${sessionDbId} for Claude session ${session_id}`);
+      // Check for inactive sessions we can reuse
+      const inactive = db.findAnySDKSession(session_id);
+
+      if (inactive) {
+        // Reactivate the existing session
+        sessionDbId = inactive.id;
+        db.reactivateSession(sessionDbId, prompt);
+        const promptNumber = db.incrementPromptCounter(sessionDbId);
+        isNewSession = true;
+        console.error(`[new-hook] Reactivated session ${sessionDbId}, prompt #${promptNumber}`);
+      } else {
+        // Create new session
+        sessionDbId = db.createSDKSession(session_id, project, prompt);
+        const promptNumber = db.incrementPromptCounter(sessionDbId);
+        isNewSession = true;
+        console.error(`[new-hook] Created new session ${sessionDbId}, prompt #${promptNumber}`);
+      }
     }
 
     // Find worker service port
@@ -78,16 +83,19 @@ export async function newHook(input?: UserPromptSubmitInput): Promise<void> {
       return;
     }
 
-    // Initialize session via HTTP
-    const response = await fetch(`http://127.0.0.1:${port}/sessions/${sessionDbId}/init`, {
-      method: 'POST',
-      headers: { 'Content-Type': 'application/json' },
-      body: JSON.stringify({ project, userPrompt: prompt }),
-      signal: AbortSignal.timeout(5000)
-    });
+    // Only initialize worker on new sessions
+    if (isNewSession) {
+      // Initialize session via HTTP
+      const response = await fetch(`http://127.0.0.1:${port}/sessions/${sessionDbId}/init`, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ project, userPrompt: prompt }),
+        signal: AbortSignal.timeout(5000)
+      });
 
-    if (!response.ok) {
-      console.error('[new-hook] Failed to init session:', await response.text());
+      if (!response.ok) {
+        console.error('[new-hook] Failed to init session:', await response.text());
+      }
     }
 
     console.log(createHookResponse('UserPromptSubmit', true));

src/hooks/save.ts

@@ -34,19 +34,24 @@ export async function saveHook(input?: PostToolUseInput): Promise<void> {
 
   const db = new HooksDatabase();
   const session = db.findActiveSDKSession(session_id);
-  db.close();
 
   if (!session) {
+    db.close();
     console.log(createHookResponse('PostToolUse', true));
     return;
   }
 
   if (!session.worker_port) {
+    db.close();
     console.error('[save-hook] No worker port for session', session.id);
     console.log(createHookResponse('PostToolUse', true));
     return;
   }
 
+  // Get current prompt number for this session
+  const promptNumber = db.getPromptCounter(session.id);
+  db.close();
+
   try {
     const response = await fetch(`http://127.0.0.1:${session.worker_port}/sessions/${session.id}/observations`, {
       method: 'POST',
@@ -54,7 +59,8 @@ export async function saveHook(input?: PostToolUseInput): Promise<void> {
       body: JSON.stringify({
         tool_name,
         tool_input: JSON.stringify(tool_input),
-        tool_output: JSON.stringify(tool_output)
+        tool_output: JSON.stringify(tool_output),
+        prompt_number: promptNumber
       }),
       signal: AbortSignal.timeout(2000)
     });

src/hooks/summary.ts

@@ -9,7 +9,7 @@ export interface StopInput {
 
 /**
  * Summary Hook - Stop
- * Sends FINALIZE message to worker via HTTP POST
+ * Sends SUMMARIZE message to worker via HTTP POST (not finalize - keeps SDK agent running)
  */
 export async function summaryHook(input?: StopInput): Promise<void> {
   if (!input) {
@@ -19,28 +19,34 @@ export async function summaryHook(input?: StopInput): Promise<void> {
   const { session_id } = input;
   const db = new HooksDatabase();
   const session = db.findActiveSDKSession(session_id);
-  db.close();
 
   if (!session) {
+    db.close();
     console.log(createHookResponse('Stop', true));
     return;
   }
 
   if (!session.worker_port) {
+    db.close();
     console.error('[summary-hook] No worker port for session', session.id);
     console.log(createHookResponse('Stop', true));
     return;
   }
 
+  // Get current prompt number
+  const promptNumber = db.getPromptCounter(session.id);
+  db.close();
+
   try {
-    const response = await fetch(`http://127.0.0.1:${session.worker_port}/sessions/${session.id}/finalize`, {
+    const response = await fetch(`http://127.0.0.1:${session.worker_port}/sessions/${session.id}/summarize`, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({ prompt_number: promptNumber }),
       signal: AbortSignal.timeout(2000)
     });
 
     if (!response.ok) {
-      console.error('[summary-hook] Failed to finalize:', await response.text());
+      console.error('[summary-hook] Failed to generate summary:', await response.text());
     }
   } catch (error: any) {
     console.error('[summary-hook] Error:', error.message);

src/services/sqlite/HooksDatabase.ts

@@ -18,8 +18,10 @@ export class HooksDatabase {
     this.db.pragma('synchronous = NORMAL');
     this.db.pragma('foreign_keys = ON');
 
-    // Run migration to add worker_port column if it doesn't exist
+    // Run migrations
     this.ensureWorkerPortColumn();
+    this.ensurePromptTrackingColumns();
+    this.removeSessionSummariesUniqueConstraint();
   }
 
   /**
@@ -40,6 +42,125 @@ export class HooksDatabase {
     }
   }
 
+  /**
+   * Ensure prompt tracking columns exist (migration 006)
+   */
+  private ensurePromptTrackingColumns(): void {
+    try {
+      // Check sdk_sessions for prompt_counter
+      const sessionsInfo = this.db.pragma('table_info(sdk_sessions)');
+      const hasPromptCounter = (sessionsInfo as any[]).some((col: any) => col.name === 'prompt_counter');
+
+      if (!hasPromptCounter) {
+        this.db.exec('ALTER TABLE sdk_sessions ADD COLUMN prompt_counter INTEGER DEFAULT 0');
+        console.error('[HooksDatabase] Added prompt_counter column to sdk_sessions table');
+      }
+
+      // Check observations for prompt_number
+      const observationsInfo = this.db.pragma('table_info(observations)');
+      const obsHasPromptNumber = (observationsInfo as any[]).some((col: any) => col.name === 'prompt_number');
+
+      if (!obsHasPromptNumber) {
+        this.db.exec('ALTER TABLE observations ADD COLUMN prompt_number INTEGER');
+        console.error('[HooksDatabase] Added prompt_number column to observations table');
+      }
+
+      // Check session_summaries for prompt_number
+      const summariesInfo = this.db.pragma('table_info(session_summaries)');
+      const sumHasPromptNumber = (summariesInfo as any[]).some((col: any) => col.name === 'prompt_number');
+
+      if (!sumHasPromptNumber) {
+        this.db.exec('ALTER TABLE session_summaries ADD COLUMN prompt_number INTEGER');
+        console.error('[HooksDatabase] Added prompt_number column to session_summaries table');
+      }
+
+      // Remove UNIQUE constraint on session_summaries.sdk_session_id
+      // SQLite doesn't support dropping constraints, so we need to check if it exists first
+      const summariesIndexes = this.db.pragma('index_list(session_summaries)');
+      const hasUniqueConstraint = (summariesIndexes as any[]).some((idx: any) => idx.unique === 1);
+
+    } catch (error: any) {
+      console.error('[HooksDatabase] Prompt tracking migration error:', error.message);
+    }
+  }
+
+  /**
+   * Remove UNIQUE constraint from session_summaries.sdk_session_id (migration 007)
+   */
+  private removeSessionSummariesUniqueConstraint(): void {
+    try {
+      // Check if UNIQUE constraint exists
+      const summariesIndexes = this.db.pragma('index_list(session_summaries)');
+      const hasUniqueConstraint = (summariesIndexes as any[]).some((idx: any) => idx.unique === 1);
+
+      if (!hasUniqueConstraint) {
+        // Already migrated
+        return;
+      }
+
+      console.error('[HooksDatabase] Removing UNIQUE constraint from session_summaries.sdk_session_id...');
+
+      // Begin transaction
+      this.db.exec('BEGIN TRANSACTION');
+
+      try {
+        // Create new table without UNIQUE constraint
+        this.db.exec(`
+          CREATE TABLE session_summaries_new (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            sdk_session_id TEXT NOT NULL,
+            project TEXT NOT NULL,
+            request TEXT,
+            investigated TEXT,
+            learned TEXT,
+            completed TEXT,
+            next_steps TEXT,
+            files_read TEXT,
+            files_edited TEXT,
+            notes TEXT,
+            prompt_number INTEGER,
+            created_at TEXT NOT NULL,
+            created_at_epoch INTEGER NOT NULL,
+            FOREIGN KEY(sdk_session_id) REFERENCES sdk_sessions(sdk_session_id) ON DELETE CASCADE
+          )
+        `);
+
+        // Copy data from old table
+        this.db.exec(`
+          INSERT INTO session_summaries_new
+          SELECT id, sdk_session_id, project, request, investigated, learned,
+                 completed, next_steps, files_read, files_edited, notes,
+                 prompt_number, created_at, created_at_epoch
+          FROM session_summaries
+        `);
+
+        // Drop old table
+        this.db.exec('DROP TABLE session_summaries');
+
+        // Rename new table
+        this.db.exec('ALTER TABLE session_summaries_new RENAME TO session_summaries');
+
+        // Recreate indexes
+        this.db.exec(`
+          CREATE INDEX idx_session_summaries_sdk_session ON session_summaries(sdk_session_id);
+          CREATE INDEX idx_session_summaries_project ON session_summaries(project);
+          CREATE INDEX idx_session_summaries_created ON session_summaries(created_at_epoch DESC);
+        `);
+
+        // Commit transaction
+        this.db.exec('COMMIT');
+
+        console.error('[HooksDatabase] Successfully removed UNIQUE constraint from session_summaries.sdk_session_id');
+      } catch (error: any) {
+        // Rollback on error
+        this.db.exec('ROLLBACK');
+        throw error;
+      }
+    } catch (error: any) {
+      console.error('[HooksDatabase] Migration error (remove UNIQUE constraint):', error.message);
+    }
+  }
+
   /**
    * Get recent session summaries for a project
    */
@@ -52,12 +173,13 @@ export class HooksDatabase {
     files_read: string | null;
     files_edited: string | null;
     notes: string | null;
+    prompt_number: number | null;
     created_at: string;
   }> {
     const stmt = this.db.prepare(`
       SELECT
         request, investigated, learned, completed, next_steps,
-        files_read, files_edited, notes, created_at
+        files_read, files_edited, notes, prompt_number, created_at
       FROM session_summaries
       WHERE project = ?
       ORDER BY created_at_epoch DESC
@@ -73,10 +195,11 @@ export class HooksDatabase {
   getRecentObservations(project: string, limit: number = 20): Array<{
     type: string;
     text: string;
+    prompt_number: number | null;
     created_at: string;
   }> {
     const stmt = this.db.prepare(`
-      SELECT type, text, created_at
+      SELECT type, text, prompt_number, created_at
       FROM observations
       WHERE project = ?
       ORDER BY created_at_epoch DESC
@@ -132,6 +255,36 @@ export class HooksDatabase {
     stmt.run(userPrompt, id);
   }
 
+  /**
+   * Increment prompt counter and return new value
+   */
+  incrementPromptCounter(id: number): number {
+    const stmt = this.db.prepare(`
+      UPDATE sdk_sessions
+      SET prompt_counter = COALESCE(prompt_counter, 0) + 1
+      WHERE id = ?
+    `);
+
+    stmt.run(id);
+
+    const result = this.db.prepare(`
+      SELECT prompt_counter FROM sdk_sessions WHERE id = ?
+    `).get(id) as { prompt_counter: number } | undefined;
+
+    return result?.prompt_counter || 1;
+  }
+
+  /**
+   * Get current prompt counter for a session
+   */
+  getPromptCounter(id: number): number {
+    const result = this.db.prepare(`
+      SELECT prompt_counter FROM sdk_sessions WHERE id = ?
+    `).get(id) as { prompt_counter: number | null } | undefined;
+
+    return result?.prompt_counter || 0;
+  }
+
   /**
    * Create a new SDK session
    */
@@ -151,15 +304,20 @@ export class HooksDatabase {
 
   /**
    * Update SDK session ID (captured from init message)
+   * Only updates if current sdk_session_id is NULL to avoid breaking foreign keys
    */
   updateSDKSessionId(id: number, sdkSessionId: string): void {
     const stmt = this.db.prepare(`
       UPDATE sdk_sessions
       SET sdk_session_id = ?
-      WHERE id = ?
+      WHERE id = ? AND sdk_session_id IS NULL
     `);
 
-    stmt.run(sdkSessionId, id);
+    const result = stmt.run(sdkSessionId, id);
+
+    if (result.changes === 0) {
+      console.error(`[HooksDatabase] Skipped updating sdk_session_id for session ${id} - already set (prevents FOREIGN KEY constraint violation)`);
+    }
   }
 
   /**
@@ -197,18 +355,19 @@ export class HooksDatabase {
     sdkSessionId: string,
     project: string,
     type: string,
-    text: string
+    text: string,
+    promptNumber?: number
   ): void {
     const now = new Date();
     const nowEpoch = now.getTime();
 
     const stmt = this.db.prepare(`
       INSERT INTO observations
-      (sdk_session_id, project, text, type, created_at, created_at_epoch)
-      VALUES (?, ?, ?, ?, ?, ?)
+      (sdk_session_id, project, text, type, prompt_number, created_at, created_at_epoch)
+      VALUES (?, ?, ?, ?, ?, ?, ?)
     `);
 
-    stmt.run(sdkSessionId, project, text, type, now.toISOString(), nowEpoch);
+    stmt.run(sdkSessionId, project, text, type, promptNumber || null, now.toISOString(), nowEpoch);
   }
 
   /**
@@ -226,7 +385,8 @@ export class HooksDatabase {
       files_read?: string;
       files_edited?: string;
       notes?: string;
-    }
+    },
+    promptNumber?: number
   ): void {
     const now = new Date();
     const nowEpoch = now.getTime();
@@ -234,8 +394,8 @@ export class HooksDatabase {
     const stmt = this.db.prepare(`
       INSERT INTO session_summaries
       (sdk_session_id, project, request, investigated, learned, completed,
-       next_steps, files_read, files_edited, notes, created_at, created_at_epoch)
-      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+       next_steps, files_read, files_edited, notes, prompt_number, created_at, created_at_epoch)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
     `);
 
     stmt.run(
@@ -249,6 +409,7 @@ export class HooksDatabase {
       summary.files_read || null,
       summary.files_edited || null,
       summary.notes || null,
+      promptNumber || null,
       now.toISOString(),
       nowEpoch
     );

src/services/worker-service.ts

@@ -20,13 +20,15 @@ interface ObservationMessage {
   tool_name: string;
   tool_input: string;
   tool_output: string;
+  prompt_number: number;
 }
 
-interface FinalizeMessage {
-  type: 'finalize';
+interface SummarizeMessage {
+  type: 'summarize';
+  prompt_number: number;
 }
 
-type WorkerMessage = ObservationMessage | FinalizeMessage;
+type WorkerMessage = ObservationMessage | SummarizeMessage;
 
 /**
  * Active session state
@@ -36,10 +38,10 @@ interface ActiveSession {
   sdkSessionId: string | null;
   project: string;
   userPrompt: string;
-  isFinalized: boolean;
   pendingMessages: WorkerMessage[];
   abortController: AbortController;
   generatorPromise: Promise<void> | null;
+  lastPromptNumber: number; // Track which prompt_number we last sent to SDK
 }
 
 class WorkerService {
@@ -57,7 +59,7 @@ class WorkerService {
     // Session endpoints
     this.app.post('/sessions/:sessionDbId/init', this.handleInit.bind(this));
     this.app.post('/sessions/:sessionDbId/observations', this.handleObservation.bind(this));
-    this.app.post('/sessions/:sessionDbId/finalize', this.handleFinalize.bind(this));
+    this.app.post('/sessions/:sessionDbId/summarize', this.handleSummarize.bind(this));
     this.app.get('/sessions/:sessionDbId/status', this.handleStatus.bind(this));
     this.app.delete('/sessions/:sessionDbId', this.handleDelete.bind(this));
   }
@@ -133,10 +135,10 @@ class WorkerService {
       sdkSessionId: null,
       project,
       userPrompt,
-      isFinalized: false,
       pendingMessages: [],
       abortController: new AbortController(),
-      generatorPromise: null
+      generatorPromise: null,
+      lastPromptNumber: 0
     };
 
     this.sessions.set(sessionDbId, session);
@@ -164,11 +166,11 @@ class WorkerService {
 
   /**
    * POST /sessions/:sessionDbId/observations
-   * Body: { tool_name, tool_input, tool_output }
+   * Body: { tool_name, tool_input, tool_output, prompt_number }
    */
   private handleObservation(req: Request, res: Response): void {
     const sessionDbId = parseInt(req.params.sessionDbId, 10);
-    const { tool_name, tool_input, tool_output } = req.body;
+    const { tool_name, tool_input, tool_output, prompt_number } = req.body;
 
     const session = this.sessions.get(sessionDbId);
     if (!session) {
@@ -176,28 +178,26 @@ class WorkerService {
       return;
     }
 
-    if (session.isFinalized) {
-      res.status(400).json({ error: 'Session already finalized' });
-      return;
-    }
-
     console.error(`[WorkerService] Queueing observation for session ${sessionDbId}:`, tool_name);
 
     session.pendingMessages.push({
       type: 'observation',
       tool_name,
       tool_input,
-      tool_output
+      tool_output,
+      prompt_number
     });
 
     res.json({ status: 'queued', queueLength: session.pendingMessages.length });
   }
 
   /**
-   * POST /sessions/:sessionDbId/finalize
+   * POST /sessions/:sessionDbId/summarize
+   * Body: { prompt_number }
    */
-  private handleFinalize(req: Request, res: Response): void {
+  private handleSummarize(req: Request, res: Response): void {
     const sessionDbId = parseInt(req.params.sessionDbId, 10);
+    const { prompt_number } = req.body;
 
     const session = this.sessions.get(sessionDbId);
     if (!session) {
@@ -205,16 +205,14 @@ class WorkerService {
       return;
     }
 
-    if (session.isFinalized) {
-      res.status(400).json({ error: 'Session already finalized' });
-      return;
-    }
+    console.error(`[WorkerService] Requesting summary for session ${sessionDbId}, prompt #${prompt_number}`);
 
-    console.error(`[WorkerService] Finalizing session ${sessionDbId}`);
+    session.pendingMessages.push({
+      type: 'summarize',
+      prompt_number
+    });
 
-    session.pendingMessages.push({ type: 'finalize' });
-
-    res.json({ status: 'finalizing' });
+    res.json({ status: 'queued', queueLength: session.pendingMessages.length });
   }
 
   /**
@@ -233,7 +231,6 @@ class WorkerService {
       sessionDbId,
       sdkSessionId: session.sdkSessionId,
       project: session.project,
-      isFinalized: session.isFinalized,
       pendingMessages: session.pendingMessages.length
     });
   }
@@ -263,12 +260,10 @@ class WorkerService {
       ]);
     }
 
-    // Mark as failed if not completed
-    if (!session.isFinalized) {
-      const db = new HooksDatabase();
-      db.markSessionFailed(sessionDbId);
-      db.close();
-    }
+    // Mark as failed since we're aborting
+    const db = new HooksDatabase();
+    db.markSessionFailed(sessionDbId);
+    db.close();
 
     this.sessions.delete(sessionDbId);
 
@@ -315,10 +310,10 @@ class WorkerService {
             ? content.filter((c: any) => c.type === 'text').map((c: any) => c.text).join('\n')
             : typeof content === 'string' ? content : '';
 
-          console.error(`[WorkerService] SDK response (${textContent.length} chars)`);
+          console.error(`[WorkerService] SDK response (${textContent.length} chars) for prompt #${session.lastPromptNumber}`);
 
-          // Parse and store
-          this.handleAgentMessage(session, textContent);
+          // Parse and store with prompt number
+          this.handleAgentMessage(session, textContent, session.lastPromptNumber);
         }
       }
 
@@ -342,6 +337,7 @@ class WorkerService {
 
   /**
    * Create async message generator for SDK streaming
+   * Keeps running continuously - no finalize, agent stays alive for entire Claude Code session
    */
   private async* createMessageGenerator(session: ActiveSession): AsyncIterable<SDKUserMessage> {
     const claudeSessionId = `session-${session.sessionDbId}`;
@@ -359,8 +355,12 @@ class WorkerService {
       }
     };
 
-    // Process messages as they arrive
-    while (!session.isFinalized) {
+    // Process messages continuously until session is deleted
+    while (true) {
+      if (session.abortController.signal.aborted) {
+        break;
+      }
+
       if (session.pendingMessages.length === 0) {
         await new Promise(resolve => setTimeout(resolve, 100));
         continue;
@@ -369,9 +369,9 @@ class WorkerService {
       while (session.pendingMessages.length > 0) {
         const message = session.pendingMessages.shift()!;
 
-        if (message.type === 'finalize') {
-          console.error(`[WorkerService] Processing FINALIZE for session ${session.sessionDbId}`);
-          session.isFinalized = true;
+        if (message.type === 'summarize') {
+          console.error(`[WorkerService] Processing SUMMARIZE for session ${session.sessionDbId}, prompt #${message.prompt_number}`);
+          session.lastPromptNumber = message.prompt_number;
 
           const db = new HooksDatabase();
           const dbSession = db.db.prepare(`
@@ -382,8 +382,24 @@ class WorkerService {
           db.close();
 
           if (dbSession) {
-            const finalizePrompt = buildFinalizePrompt(dbSession);
-            console.error(`[WorkerService] Yielding finalize prompt (${finalizePrompt.length} chars)`);
+            const summarizePrompt = `You have been processing tool observations for this session. Please generate a summary of what you've learned from prompt #${message.prompt_number}.
+
+Use this XML format:
+
+<session_summary>
+  <request>What was the user trying to accomplish in this prompt?</request>
+  <investigated>What code/systems did you explore?</investigated>
+  <learned>What did you learn about the codebase?</learned>
+  <completed>What was done or determined?</completed>
+  <next_steps>What should happen next?</next_steps>
+  <files_read>["file1.ts", "file2.ts"]</files_read>
+  <files_edited>["file3.ts"]</files_edited>
+  <notes>Any additional context or insights</notes>
+</session_summary>
+
+Respond ONLY with the XML block. Be concise and specific.`;
+
+            console.error(`[WorkerService] Yielding summarize prompt`);
 
             yield {
               type: 'user',
@@ -391,14 +407,13 @@ class WorkerService {
               parent_tool_use_id: null,
               message: {
                 role: 'user',
-                content: finalizePrompt
+                content: summarizePrompt
               }
             };
           }
-          break;
-        }
+        } else if (message.type === 'observation') {
+          session.lastPromptNumber = message.prompt_number;
 
-        if (message.type === 'observation') {
           const observationPrompt = buildObservationPrompt({
             id: 0,
             tool_name: message.tool_name,
@@ -407,7 +422,7 @@ class WorkerService {
             created_at_epoch: Date.now()
           });
 
-          console.error(`[WorkerService] Yielding observation: ${message.tool_name}`);
+          console.error(`[WorkerService] Yielding observation: ${message.tool_name} (prompt #${message.prompt_number})`);
 
           yield {
             type: 'user',
@@ -425,23 +440,24 @@ class WorkerService {
 
   /**
    * Handle agent message - parse and store observations/summaries
+   * Gets prompt_number from the message that triggered this response
    */
-  private handleAgentMessage(session: ActiveSession, content: string): void {
+  private handleAgentMessage(session: ActiveSession, content: string, promptNumber: number): void {
     // Parse observations
     const observations = parseObservations(content);
-    console.error(`[WorkerService] Parsed ${observations.length} observations`);
+    console.error(`[WorkerService] Parsed ${observations.length} observations for prompt #${promptNumber}`);
 
     const db = new HooksDatabase();
     for (const obs of observations) {
       if (session.sdkSessionId) {
-        db.storeObservation(session.sdkSessionId, session.project, obs.type, obs.text);
+        db.storeObservation(session.sdkSessionId, session.project, obs.type, obs.text, promptNumber);
       }
     }
 
     // Parse summary
     const summary = parseSummary(content);
     if (summary && session.sdkSessionId) {
-      console.error(`[WorkerService] Parsed summary for session ${session.sessionDbId}`);
+      console.error(`[WorkerService] Parsed summary for session ${session.sessionDbId}, prompt #${promptNumber}`);
 
       const summaryWithArrays = {
         request: summary.request,
@@ -454,7 +470,7 @@ class WorkerService {
         notes: summary.notes
       };
 
-      db.storeSummary(session.sdkSessionId, session.project, summaryWithArrays);
+      db.storeSummary(session.sdkSessionId, session.project, summaryWithArrays, promptNumber);
     }
 
     db.close();