From 72720ef2921a8f9df5d820e41374249a01b9ba3e Mon Sep 17 00:00:00 2001
From: Lucio Lelii <lucio.lelii@isti.cnr.it>
Date: Fri, 24 Apr 2026 12:18:40 +0200
Subject: [PATCH] Persist assistant sessions per flow and auto-scroll latest
 messages

---
 .../shared/flow-assistant/flow-assistant.html |   2 +-
 .../shared/flow-assistant/flow-assistant.ts   | 165 +++++++++++++++++-
 src/app/stores/assistant-session-store.ts     | 131 ++++++++++++++
 3 files changed, 291 insertions(+), 7 deletions(-)
 create mode 100644 src/app/stores/assistant-session-store.ts

diff --git a/src/app/shared/flow-assistant/flow-assistant.html b/src/app/shared/flow-assistant/flow-assistant.html
index 21ab6a4..d29d7fc 100644
--- a/src/app/shared/flow-assistant/flow-assistant.html
+++ b/src/app/shared/flow-assistant/flow-assistant.html
@@ -1,5 +1,5 @@
 <aside class="assistant-shell">
-  <div class="assistant-scroll">
+  <div #assistantScroll class="assistant-scroll">
     <header class="assistant-header">
       <div>
         <p class="assistant-eyebrow">Workflow Copilot</p>
diff --git a/src/app/shared/flow-assistant/flow-assistant.ts b/src/app/shared/flow-assistant/flow-assistant.ts
index c6898ef..43b1fbb 100644
--- a/src/app/shared/flow-assistant/flow-assistant.ts
+++ b/src/app/shared/flow-assistant/flow-assistant.ts
@@ -1,5 +1,5 @@
 import { CommonModule } from '@angular/common';
-import { ChangeDetectionStrategy, Component, computed, inject, OnDestroy, OnInit, signal } from '@angular/core';
+import { ChangeDetectionStrategy, Component, computed, effect, ElementRef, inject, OnDestroy, OnInit, signal, ViewChild } from '@angular/core';
 import { FormsModule } from '@angular/forms';
 import { MatButtonModule } from '@angular/material/button';
 import { MatFormFieldModule } from '@angular/material/form-field';
@@ -17,6 +17,7 @@ import {
 import { Flow } from '@models/flow';
 import { AssistantService } from '@services/assistant/assistant';
 import { Authorization } from '@services/authorization/authorization';
+import { AssistantSessionStore } from '@stores/assistant-session-store';
 import { EditorStateHolder } from '@stores/flow-editor';
 import { finalize, interval, Subscription, switchMap, take } from 'rxjs';
 
@@ -28,10 +29,16 @@ import { finalize, interval, Subscription, switchMap, take } from 'rxjs';
   changeDetection: ChangeDetectionStrategy.OnPush
 })
 export class FlowAssistant implements OnInit, OnDestroy {
+  @ViewChild('assistantScroll') assistantScrollElement?: ElementRef<HTMLDivElement>;
+
   private readonly assistant = inject(AssistantService);
   private readonly editorState = inject(EditorStateHolder);
   private readonly authorization = inject(Authorization);
+  private readonly sessionStore = inject(AssistantSessionStore);
   private pollSubscription: Subscription | null = null;
+  private initialized = false;
+  private activeFlowKey: string | null = null;
+  private lastAutoScrollKey = '';
 
   readonly assistantConfig = signal<AssistantConfig | null>(null);
   readonly sessionState = signal<AssistantSessionState | null>(null);
@@ -137,11 +144,39 @@ export class FlowAssistant implements OnInit, OnDestroy {
   readonly currentFlow = this.editorState.currentFlow;
   readonly currentDraft = computed(() => this.sessionState()?.currentDraftFlow ?? null);
 
+  constructor() {
+    effect(() => {
+      const flowKey = this.resolveFlowKey(this.currentFlow()?.id ?? null);
+      if (!this.initialized || this.activeFlowKey === flowKey) return;
+
+      const previousFlowKey = this.activeFlowKey;
+      if (previousFlowKey) {
+        this.persistSnapshot(previousFlowKey);
+      }
+
+      this.activeFlowKey = flowKey;
+      void this.restoreSessionForFlow(flowKey);
+    });
+
+    effect(() => {
+      const messages = this.displayedMessages();
+      const lastMessageId = messages[messages.length - 1]?.id ?? '';
+      const busy = this.assistantBusy();
+      const nextKey = `${messages.length}:${lastMessageId}:${busy ? '1' : '0'}`;
+      if (nextKey === this.lastAutoScrollKey) return;
+      this.lastAutoScrollKey = nextKey;
+      this.scheduleScrollToLatestMessage();
+    });
+  }
+
   ngOnInit(): void {
     this.bootstrapAssistant();
   }
 
   ngOnDestroy(): void {
+    if (this.activeFlowKey) {
+      this.persistSnapshot(this.activeFlowKey);
+    }
     this.stopPolling();
   }
 
@@ -153,14 +188,17 @@ export class FlowAssistant implements OnInit, OnDestroy {
 
   useStarter(prompt: string) {
     this.prompt.set(prompt);
+    this.persistSnapshot();
   }
 
   toggleQuickPrompts() {
     this.quickPromptsOpen.update((value) => !value);
+    this.persistSnapshot();
   }
 
   toggleModelPicker() {
     this.modelPickerOpen.update((value) => !value);
+    this.persistSnapshot();
   }
 
   submitPrompt() {
@@ -176,6 +214,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
         content
       }
     ]);
+    this.persistSnapshot();
 
     this.assistant.sendMessage(sessionId, { message: content }).pipe(
       take(1)
@@ -187,6 +226,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
           status: 'QUEUED',
           phase: 'queued'
         });
+        this.persistSnapshot();
         this.startPolling(callId, sessionId);
       },
       error: (err) => {
@@ -230,7 +270,9 @@ export class FlowAssistant implements OnInit, OnDestroy {
           this.modelsError.set('No assistant model is available.');
           return;
         }
-        void this.openSession(selectedModel);
+        this.activeFlowKey = this.resolveFlowKey(this.currentFlow()?.id ?? null);
+        this.initialized = true;
+        void this.restoreSessionForFlow(this.activeFlowKey);
       },
       error: (err) => {
         console.error('Assistant model loading failed', err);
@@ -242,7 +284,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
     });
   }
 
-  private async openSession(model: string) {
+  private async openSession(model: string, flowKey = this.activeFlowKey ?? this.resolveFlowKey(this.currentFlow()?.id ?? null)) {
     this.stopPolling();
     this.currentCall.set(null);
     this.localMessages.set([]);
@@ -254,6 +296,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
     ).subscribe({
       next: (session) => {
         this.applySessionState(session);
+        this.persistSnapshot(flowKey);
       },
       error: (err) => {
         console.error('Assistant session creation failed', err);
@@ -269,6 +312,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
     ).subscribe({
       next: (callState) => {
         this.currentCall.set(callState);
+        this.persistSnapshot();
         if (callState.status === 'COMPLETED' || callState.status === 'FAILED') {
           this.stopPolling();
           void this.reloadSession(sessionId, callState.status === 'FAILED' ? callState.errorMessage : undefined);
@@ -289,6 +333,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
       next: (session) => {
         this.applySessionState(session);
         this.currentCall.set(null);
+        this.persistSnapshot();
         if (failureMessage) {
           this.pushLocalAssistantMessage(failureMessage);
         }
@@ -301,7 +346,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
     });
   }
 
-  private applySessionState(session: AssistantSessionState) {
+  private applySessionState(session: AssistantSessionState, options?: { clearLocalMessages?: boolean; syncDraftToEditor?: boolean }) {
     const normalizedSession = session.messages.length
       ? session
       : {
@@ -311,8 +356,13 @@ export class FlowAssistant implements OnInit, OnDestroy {
 
     this.sessionState.set(normalizedSession);
     this.selectedModel.set(normalizedSession.selectedModel || this.selectedModel());
-    this.localMessages.set([]);
-    this.syncDraftToEditor(normalizedSession.currentDraftFlow);
+    if (options?.clearLocalMessages !== false) {
+      this.localMessages.set([]);
+    }
+    if (options?.syncDraftToEditor !== false) {
+      this.syncDraftToEditor(normalizedSession.currentDraftFlow);
+    }
+    this.persistSnapshot();
   }
 
   private syncDraftToEditor(draft: AssistantDraftPayload | null) {
@@ -326,6 +376,11 @@ export class FlowAssistant implements OnInit, OnDestroy {
       return;
     }
 
+    const currentFlowKey = this.activeFlowKey ?? this.resolveFlowKey(currentFlow?.id ?? null);
+    const nextFlowKey = this.resolveFlowKey(nextFlow.id);
+    this.persistSnapshot(currentFlowKey);
+    this.sessionStore.cloneSnapshot(currentFlowKey, nextFlowKey);
+
     void this.editorState.openDocument(nextFlow, { skipDirtyCheck: false }).then((opened) => {
       if (!opened) {
         this.pushLocalAssistantMessage('The draft is ready, but I did not load it because the current flow has unsaved changes.');
@@ -363,6 +418,7 @@ export class FlowAssistant implements OnInit, OnDestroy {
         content
       }
     ]);
+    this.persistSnapshot();
   }
 
   private systemWelcomeMessage(): AssistantChatMessage {
@@ -375,6 +431,103 @@ export class FlowAssistant implements OnInit, OnDestroy {
     };
   }
 
+  private async restoreSessionForFlow(flowKey: string) {
+    this.stopPolling();
+    this.currentCall.set(null);
+
+    const snapshot = this.sessionStore.getSnapshot(flowKey);
+    if (!snapshot) {
+      this.prompt.set('');
+      this.modelPickerOpen.set(false);
+      this.quickPromptsOpen.set(true);
+      this.sessionState.set(null);
+      this.localMessages.set([]);
+      const model = this.selectedModel();
+      if (!model) {
+        this.sessionLoading.set(false);
+        return;
+      }
+      void this.openSession(model, flowKey);
+      return;
+    }
+
+    this.prompt.set(snapshot.prompt);
+    this.modelPickerOpen.set(snapshot.modelPickerOpen);
+    this.quickPromptsOpen.set(snapshot.quickPromptsOpen);
+    this.localMessages.set(snapshot.localMessages);
+    this.currentCall.set(null);
+    if (snapshot.selectedModel) {
+      this.selectedModel.set(snapshot.selectedModel);
+    }
+    if (snapshot.sessionState) {
+      this.sessionState.set(snapshot.sessionState);
+    } else {
+      this.sessionState.set(null);
+    }
+
+    if (!snapshot.sessionId) {
+      const model = this.selectedModel();
+      if (!model) {
+        this.sessionLoading.set(false);
+        return;
+      }
+      void this.openSession(model, flowKey);
+      return;
+    }
+
+    this.sessionLoading.set(true);
+    this.assistant.getSession(snapshot.sessionId).pipe(
+      take(1),
+      finalize(() => this.sessionLoading.set(false))
+    ).subscribe({
+      next: (session) => {
+        this.applySessionState(session, {
+          clearLocalMessages: false,
+          syncDraftToEditor: !this.currentFlow()?.id
+        });
+      },
+      error: (err) => {
+        console.error('Assistant session restore failed', err);
+        const model = this.selectedModel();
+        if (!model) return;
+        void this.openSession(model, flowKey);
+      }
+    });
+  }
+
+  private resolveFlowKey(flowId: string | null | undefined): string {
+    return this.sessionStore.flowKey(flowId);
+  }
+
+  private persistSnapshot(flowKey = this.activeFlowKey ?? this.resolveFlowKey(this.currentFlow()?.id ?? null)) {
+    this.sessionStore.setSnapshot(flowKey, {
+      sessionId: this.sessionState()?.id ?? null,
+      selectedModel: this.selectedModel(),
+      prompt: this.prompt(),
+      modelPickerOpen: this.modelPickerOpen(),
+      quickPromptsOpen: this.quickPromptsOpen(),
+      localMessages: this.localMessages(),
+      currentCall: this.currentCall(),
+      sessionState: this.sessionState()
+    });
+  }
+
+  private scheduleScrollToLatestMessage() {
+    queueMicrotask(() => {
+      if (typeof requestAnimationFrame === 'function') {
+        requestAnimationFrame(() => this.scrollToLatestMessage());
+        return;
+      }
+      setTimeout(() => this.scrollToLatestMessage(), 0);
+    });
+  }
+
+  private scrollToLatestMessage() {
+    const container = this.assistantScrollElement?.nativeElement;
+    if (!container) return;
+    container.scrollTop = container.scrollHeight;
+  }
+
   private stopPolling() {
     this.pollSubscription?.unsubscribe();
     this.pollSubscription = null;
diff --git a/src/app/stores/assistant-session-store.ts b/src/app/stores/assistant-session-store.ts
new file mode 100644
index 0000000..12bfa9a
--- /dev/null
+++ b/src/app/stores/assistant-session-store.ts
@@ -0,0 +1,131 @@
+import { Injectable } from '@angular/core';
+import { AssistantCallState, AssistantChatMessage, AssistantSessionState } from '@models/assistant';
+
+export type AssistantSessionSnapshot = {
+  sessionId: string | null;
+  selectedModel: string;
+  prompt: string;
+  modelPickerOpen: boolean;
+  quickPromptsOpen: boolean;
+  localMessages: AssistantChatMessage[];
+  currentCall: AssistantCallState | null;
+  sessionState: AssistantSessionState | null;
+};
+
+@Injectable({ providedIn: 'root' })
+export class AssistantSessionStore {
+  private static readonly NO_FLOW_KEY = '__assistant:no-flow__';
+  private static readonly STORAGE_KEY = 'assistant-session-store:v1';
+  private static readonly STORAGE_TARGETS: Array<'localStorage' | 'sessionStorage'> = ['localStorage', 'sessionStorage'];
+  private readonly snapshots = new Map<string, AssistantSessionSnapshot>();
+
+  constructor() {
+    this.hydrateFromStorage();
+  }
+
+  flowKey(flowId: string | null | undefined): string {
+    return flowId && flowId.length > 0 ? flowId : AssistantSessionStore.NO_FLOW_KEY;
+  }
+
+  getSnapshot(flowKey: string): AssistantSessionSnapshot | null {
+    const snapshot = this.snapshots.get(flowKey);
+    return snapshot ? structuredClone(snapshot) : null;
+  }
+
+  setSnapshot(flowKey: string, snapshot: AssistantSessionSnapshot) {
+    this.snapshots.set(flowKey, structuredClone(snapshot));
+    this.persistToStorage();
+  }
+
+  cloneSnapshot(fromFlowKey: string, toFlowKey: string) {
+    const snapshot = this.snapshots.get(fromFlowKey);
+    if (!snapshot) return;
+    this.snapshots.set(toFlowKey, structuredClone(snapshot));
+    this.persistToStorage();
+  }
+
+  private hydrateFromStorage() {
+    const raw = this.readStoredPayload();
+    if (!raw) return;
+
+    try {
+      const parsed = JSON.parse(raw) as unknown;
+      if (!parsed || typeof parsed !== 'object') return;
+      const root = parsed as Record<string, unknown>;
+      const snapshots = root['snapshots'];
+      if (!snapshots || typeof snapshots !== 'object') return;
+
+      for (const [flowKey, rawSnapshot] of Object.entries(snapshots as Record<string, unknown>)) {
+        const normalized = this.normalizeSnapshot(rawSnapshot);
+        if (!normalized) continue;
+        this.snapshots.set(flowKey, normalized);
+      }
+    } catch {
+      this.clearStoredPayload();
+    }
+  }
+
+  private normalizeSnapshot(rawSnapshot: unknown): AssistantSessionSnapshot | null {
+    if (!rawSnapshot || typeof rawSnapshot !== 'object') return null;
+    const snapshot = rawSnapshot as Record<string, unknown>;
+
+    return {
+      sessionId: typeof snapshot['sessionId'] === 'string' ? snapshot['sessionId'] : null,
+      selectedModel: typeof snapshot['selectedModel'] === 'string' ? snapshot['selectedModel'] : '',
+      prompt: typeof snapshot['prompt'] === 'string' ? snapshot['prompt'] : '',
+      modelPickerOpen: snapshot['modelPickerOpen'] === true,
+      quickPromptsOpen: snapshot['quickPromptsOpen'] !== false,
+      localMessages: Array.isArray(snapshot['localMessages'])
+        ? (snapshot['localMessages'] as AssistantChatMessage[])
+        : [],
+      currentCall: snapshot['currentCall'] && typeof snapshot['currentCall'] === 'object'
+        ? (snapshot['currentCall'] as AssistantCallState)
+        : null,
+      sessionState: snapshot['sessionState'] && typeof snapshot['sessionState'] === 'object'
+        ? (snapshot['sessionState'] as AssistantSessionState)
+        : null
+    };
+  }
+
+  private persistToStorage() {
+    if (typeof window === 'undefined') return;
+
+    const payload = JSON.stringify({
+      snapshots: Object.fromEntries(this.snapshots.entries())
+    });
+
+    for (const target of AssistantSessionStore.STORAGE_TARGETS) {
+      try {
+        window[target].setItem(AssistantSessionStore.STORAGE_KEY, payload);
+      } catch {
+        // Ignore storage write errors (quota, private mode, unavailable storage).
+      }
+    }
+  }
+
+  private readStoredPayload(): string | null {
+    if (typeof window === 'undefined') return null;
+
+    for (const target of AssistantSessionStore.STORAGE_TARGETS) {
+      try {
+        const value = window[target].getItem(AssistantSessionStore.STORAGE_KEY);
+        if (value) return value;
+      } catch {
+        // Ignore storage read errors and keep trying the next target.
+      }
+    }
+    return null;
+  }
+
+  private clearStoredPayload() {
+    if (typeof window === 'undefined') return;
+
+    for (const target of AssistantSessionStore.STORAGE_TARGETS) {
+      try {
+        window[target].removeItem(AssistantSessionStore.STORAGE_KEY);
+      } catch {
+        // Ignore storage cleanup errors.
+      }
+    }
+  }
+}