fix: apply prompt suggestion toggle immediately

src/crates/assembly/core/src/agentic/execution/prompt_suggestion.rs

@@ -19,6 +19,7 @@ use crate::infrastructure::ai::reasoning_catalog::{
     load_models_dev_reasoning_catalog_without_refresh, project_model_reasoning_catalog,
 };
 use crate::infrastructure::ai::AIClient;
+use crate::service::config::{subscribe_config_updates, ConfigUpdateEvent};
 use crate::util::types::{Message as AIMessage, ToolDefinition};
 use log::{debug, info, warn};
 use openbitfun_agent_runtime::prompt_suggestion::{
@@ -132,35 +133,70 @@ async fn generate_prompt_suggestion(
     request: PromptSuggestionRequest,
     session_manager: Arc<SessionManager>,
 ) -> anyhow::Result<Option<String>> {
-    if !prompt_suggestion_enabled().await {
-        debug!(
-            "Prompt suggestion disabled by configuration: session_id={}, turn_id={}",
-            request.session_id, request.turn_id
-        );
+    // Subscribe before reading enablement so a saved setting cannot fall between
+    // the initial check and the in-flight request's cancellation boundary.
+    let Some(updates) = subscribe_config_updates() else {
         return Ok(None);
-    }
-
-    let ai_client = suggestion_client(&request).await;
-    let ai_messages = build_prompt_suggestion_messages(&request).await?;
-
-    let trace_dir = session_manager
-        .persistent_model_exchange_trace_dir(&request.session_id)
+    };
+    generate_while_enabled(updates, prompt_suggestion_enabled, async {
+        let ai_client = suggestion_client(&request).await;
+        let ai_messages = build_prompt_suggestion_messages(&request).await?;
+
+        let trace_dir = session_manager
+            .persistent_model_exchange_trace_dir(&request.session_id)
+            .await;
+        let trace = prepare_model_exchange_trace_for_workspace(
+            &request.session_id,
+            &request.turn_id,
+            request.workspace.as_ref(),
+            trace_dir.as_deref(),
+            ModelExchangeTraceOperation {
+                kind: "prompt_suggestion",
+                id: &request.turn_id,
+                trigger: None,
+            },
+            request.ai_client.as_ref(),
+        )
         .await;
-    let trace = prepare_model_exchange_trace_for_workspace(
-        &request.session_id,
-        &request.turn_id,
-        request.workspace.as_ref(),
-        trace_dir.as_deref(),
-        ModelExchangeTraceOperation {
-            kind: "prompt_suggestion",
-            id: &request.turn_id,
-            trigger: None,
-        },
-        request.ai_client.as_ref(),
-    )
-    .await;
 
-    request_prompt_suggestion_response(&request, ai_client, ai_messages, trace).await
+        request_prompt_suggestion_response(&request, ai_client, ai_messages, trace).await
+    })
+    .await
+}
+
+/// Drop both the active provider request and any fallback when the host disables
+/// suggestions. The same config notifications cover local, peer and relay hosts.
+async fn generate_while_enabled<F, E, EF>(
+    mut updates: tokio::sync::broadcast::Receiver<ConfigUpdateEvent>,
+    enabled: E,
+    generation: F,
+) -> anyhow::Result<Option<String>>
+where
+    F: Future<Output = anyhow::Result<Option<String>>>,
+    E: Fn() -> EF,
+    EF: Future<Output = bool>,
+{
+    if !enabled().await {
+        return Ok(None);
+    }
+    tokio::pin!(generation);
+    loop {
+        tokio::select! {
+            biased;
+            update = updates.recv() => {
+                if matches!(update, Err(tokio::sync::broadcast::error::RecvError::Closed)) {
+                    return Ok(None);
+                }
+                if !enabled().await {
+                    return Ok(None);
+                }
+            }
+            result = &mut generation => {
+                // Also fence a completed response against a simultaneous save.
+                return if enabled().await { result } else { Ok(None) };
+            }
+        }
+    }
 }
 
 async fn request_prompt_suggestion_response(
@@ -327,6 +363,100 @@ mod tests {
         assert_eq!(queue.stats().await.total_enqueued, 0);
     }
 
+    #[tokio::test]
+    async fn disabling_suggestions_drops_inflight_request_without_publishing() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+
+        struct DropWitness(Arc<AtomicBool>);
+        impl Drop for DropWitness {
+            fn drop(&mut self) {
+                self.0.store(true, Ordering::SeqCst);
+            }
+        }
+        let enabled = Arc::new(AtomicBool::new(true));
+        let check = enabled.clone();
+        let dropped = Arc::new(AtomicBool::new(false));
+        let witness = DropWitness(dropped.clone());
+        let (updates, receiver) = tokio::sync::broadcast::channel(4);
+        let (started, ready) = tokio::sync::oneshot::channel();
+        let runtime = PromptSuggestionRuntime::new();
+        let queue = Arc::new(EventQueue::new(EventQueueConfig::default()));
+        let task = spawn_prompt_suggestion_generation(
+            &runtime,
+            queue.clone(),
+            "session-1".into(),
+            "turn-1".into(),
+            generate_while_enabled(
+                receiver,
+                move || std::future::ready(check.load(Ordering::SeqCst)),
+                async move {
+                    let _witness = witness;
+                    started.send(()).unwrap();
+                    std::future::pending().await
+                },
+            ),
+        );
+        ready.await.unwrap();
+        enabled.store(false, Ordering::SeqCst);
+        updates.send(ConfigUpdateEvent::AppUpdated).unwrap();
+        tokio::time::timeout(
+            Duration::from_secs(1),
+            assert_silent_completion(&runtime, &queue, task),
+        )
+        .await
+        .expect("disabling must not wait for the provider response");
+        assert!(dropped.load(Ordering::SeqCst));
+    }
+
+    #[tokio::test]
+    async fn disabled_generation_is_not_polled_and_reenable_allows_a_new_request() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+
+        let enabled = AtomicBool::new(false);
+        let (updates, receiver) = tokio::sync::broadcast::channel(4);
+        let result = generate_while_enabled(
+            receiver,
+            || std::future::ready(enabled.load(Ordering::SeqCst)),
+            async { panic!("disabled suggestions must not call the provider") },
+        )
+        .await
+        .unwrap();
+        assert_eq!(result, None);
+
+        enabled.store(true, Ordering::SeqCst);
+        let receiver = updates.subscribe();
+        updates.send(ConfigUpdateEvent::EditorUpdated).unwrap();
+        assert_eq!(
+            generate_while_enabled(
+                receiver,
+                || std::future::ready(enabled.load(Ordering::SeqCst)),
+                async { Ok(Some("Run the tests".into())) },
+            )
+            .await
+            .unwrap(),
+            Some("Run the tests".into()),
+        );
+    }
+
+    #[tokio::test]
+    async fn completed_prediction_rechecks_enablement_before_returning() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+
+        let enabled = AtomicBool::new(true);
+        let (_updates, receiver) = tokio::sync::broadcast::channel(4);
+        let result = generate_while_enabled(
+            receiver,
+            || std::future::ready(enabled.load(Ordering::SeqCst)),
+            async {
+                enabled.store(false, Ordering::SeqCst);
+                Ok(Some("Late suggestion".into()))
+            },
+        )
+        .await
+        .unwrap();
+        assert_eq!(result, None);
+    }
+
     #[tokio::test(start_paused = true)]
     async fn timeout_at_eighty_seconds_drops_generation_and_emits_nothing() {
         use std::sync::atomic::{AtomicBool, Ordering};

src/web-ui/src/app/scenes/settings/pages/ai/PromptSuggestionSection.test.tsx

@@ -3,6 +3,7 @@ import { act } from 'react';
 import { createRoot, type Root } from 'react-dom/client';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import PromptSuggestionSection from './PromptSuggestionSection';
+import { promptSuggestionKey, usePromptSuggestionStore } from '@/flow_chat/store/promptSuggestionStore';
 
 const mocks = vi.hoisted(() => ({ get: vi.fn(), save: vi.fn(), error: vi.fn(), success: vi.fn() }));
 vi.mock('@/infrastructure/config/services/AIExperienceConfigService', () => ({
@@ -25,11 +26,16 @@ beforeEach(() => {
   vi.clearAllMocks();
   mocks.get.mockResolvedValue({ enable_prompt_suggestion: true });
   mocks.save.mockResolvedValue(undefined);
+  usePromptSuggestionStore.setState({ suggestions: {}, enabledBySurface: {} });
   container = document.createElement('div');
   document.body.append(container);
   root = createRoot(container);
 });
-afterEach(async () => { await act(async () => root.unmount()); container.remove(); });
+afterEach(async () => {
+  await act(async () => root.unmount());
+  container.remove();
+  usePromptSuggestionStore.setState({ suggestions: {}, enabledBySurface: {} });
+});
 
 describe('Prompt suggestion settings', () => {
   it('keeps older hosts without the optional setting explicitly unsupported', async () => {
@@ -41,15 +47,21 @@ describe('Prompt suggestion settings', () => {
   });
 
   it('saves only the prompt suggestion field and keeps the previous value on failure', async () => {
+    const candidate = { suggestion: 'Run the tests', sourceTurnId: 't1' };
+    usePromptSuggestionStore.getState().setSuggestion('session', candidate);
     await act(async () => root.render(<PromptSuggestionSection />));
     expect(toggle().checked).toBe(true);
     mocks.save.mockRejectedValueOnce(new Error('Host unavailable'));
     await act(async () => toggle().click());
     expect(toggle().checked).toBe(true);
     expect(mocks.error).toHaveBeenCalledWith('promptSuggestion.messages.saveFailed');
+    expect(usePromptSuggestionStore.getState().suggestions[promptSuggestionKey('session')]).toEqual(candidate);
     await act(async () => toggle().click());
     expect(mocks.save).toHaveBeenLastCalledWith({ enable_prompt_suggestion: false });
     expect(toggle().checked).toBe(false);
+    expect(usePromptSuggestionStore.getState().suggestions).toEqual({});
+    usePromptSuggestionStore.getState().setSuggestion('session', candidate);
+    expect(usePromptSuggestionStore.getState().suggestions).toEqual({});
     expect(mocks.success).toHaveBeenCalledWith('promptSuggestion.messages.saveSuccess');
   });
 

src/web-ui/src/app/scenes/settings/pages/ai/PromptSuggestionSection.tsx

@@ -5,6 +5,8 @@ import { createLogger } from '@/shared/utils/logger';
 import { Switch } from '@openbitfun/ui';
 import { useCallback, useEffect, useRef, useState } from 'react';
 import { useI18n } from '@/infrastructure/i18n';
+import { getActiveSurfaceScope } from '@/infrastructure/peer-device/deviceSurface';
+import { usePromptSuggestionStore } from '@/flow_chat/store/promptSuggestionStore';
 
 const log = createLogger('PromptSuggestionSection');
 
@@ -41,10 +43,14 @@ export default function PromptSuggestionSection() {
 
   const handleToggle = async (next: boolean) => {
     if (saveInFlight.current || enabled === undefined) return;
+    const scope = getActiveSurfaceScope();
     saveInFlight.current = true;
     setSaving(true);
     try {
       await aiExperienceConfigService.saveSettings({ enable_prompt_suggestion: next });
+      if (!scope.isCurrent()) return;
+      // The composer may be unmounted while Settings is open.
+      usePromptSuggestionStore.getState().setEnabled(next, scope.surfaceId);
       setEnabled(next);
       notification.success(t('promptSuggestion.messages.saveSuccess'));
     } catch (error) {

src/web-ui/src/flow_chat/components/ChatInput.tsx

@@ -35,6 +35,7 @@ import { SessionExecutionEvent, SessionExecutionState } from '../state-machine/t
 import { ModelSelector, type ModelSelectorAvailability } from './ModelSelector';
 import { FlowChatStore } from '../store/FlowChatStore';
 import { useAcpPlan } from '../hooks/useAcpPlan';
+import { usePromptSuggestionsEnabled } from '../hooks/usePromptSuggestionsEnabled';
 import { filterSlashCommands, useAcpSlashCommands } from '../hooks/useAcpSlashCommands';
 import { acpSessionRef, acpSlashCommandText } from '../utils/acpSession';
 import { AcpPlanPanel } from './AcpPlanPanel';
@@ -830,7 +831,9 @@ export const ChatInput: React.FC<ChatInputProps> = ({
       : undefined
   ));
   const latestSuggestionTurn = effectiveTargetSession?.dialogTurns.at(-1);
+  const promptSuggestionsEnabled = usePromptSuggestionsEnabled();
   const promptSuggestion = visiblePromptSuggestion(storedPromptSuggestion, {
+    enabled: promptSuggestionsEnabled,
     editorEmpty: showPlaceholder,
     value: inputState.value,
     canCompose: targetCanCompose,

src/web-ui/src/flow_chat/hooks/usePromptSuggestionsEnabled.test.tsx

@@ -0,0 +1,92 @@
+// @vitest-environment jsdom
+import React, { act } from 'react';
+import { createRoot, type Root } from 'react-dom/client';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { activateSurface } from '@/infrastructure/peer-device/deviceSurface';
+import { usePromptSuggestionsEnabled } from './usePromptSuggestionsEnabled';
+import { promptSuggestionKey, usePromptSuggestionStore, visiblePromptSuggestion } from '../store/promptSuggestionStore';
+
+const service = vi.hoisted(() => ({
+  cached: vi.fn(), read: vi.fn(), listeners: new Set<(value: unknown) => void>(),
+}));
+vi.mock('@/infrastructure/config/services/AIExperienceConfigService', () => ({ aiExperienceConfigService: {
+  getCachedSettings: service.cached,
+  getSettingsAsync: service.read,
+  addChangeListener: (listener: (value: unknown) => void) => {
+    service.listeners.add(listener);
+    return () => service.listeners.delete(listener);
+  },
+} }));
+
+const candidate = { suggestion: 'Run the tests', sourceTurnId: 't1' };
+function Composer() {
+  const enabled = usePromptSuggestionsEnabled();
+  const stored = usePromptSuggestionStore(state => state.suggestions[promptSuggestionKey('session')]);
+  return <div data-enabled={enabled}>{visiblePromptSuggestion(stored, {
+    enabled, editorEmpty: true, value: '', canCompose: true, processing: false,
+    latestTurnId: 't1', latestTurnStatus: 'completed',
+  })}</div>;
+}
+
+describe('prompt suggestion preference', () => {
+  let root: Root;
+  let container: HTMLDivElement;
+  beforeEach(() => {
+    Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true });
+    activateSurface('local');
+    service.cached.mockReset().mockReturnValue({ enable_prompt_suggestion: true });
+    service.read.mockReset().mockResolvedValue({ enable_prompt_suggestion: true });
+    service.listeners.clear();
+    usePromptSuggestionStore.setState({ suggestions: {}, enabledBySurface: {} });
+    container = document.createElement('div');
+    root = createRoot(container);
+  });
+  afterEach(() => {
+    act(() => root.unmount());
+    activateSurface('local');
+    usePromptSuggestionStore.setState({ suggestions: {}, enabledBySurface: {} });
+  });
+
+  it('immediately hides and clears hints, rejects late events, and enables only new hints', async () => {
+    usePromptSuggestionStore.getState().setSuggestion('session', candidate);
+    await act(async () => root.render(<Composer />));
+    expect(container.textContent).toBe('Run the tests');
+    act(() => service.listeners.forEach(listener => listener({ enable_prompt_suggestion: false })));
+    expect(container.textContent).toBe('');
+    act(() => usePromptSuggestionStore.getState().setSuggestion('session', candidate));
+    act(() => service.listeners.forEach(listener => listener({ enable_prompt_suggestion: true })));
+    expect(container.textContent).toBe('');
+    act(() => usePromptSuggestionStore.getState().setSuggestion('session', candidate));
+    expect(container.textContent).toBe('Run the tests');
+  });
+
+  it('does not show suggestions before settings load or for an older unsupported host', async () => {
+    service.cached.mockReturnValue(null);
+    let resolve!: (value: unknown) => void;
+    service.read.mockReturnValue(new Promise(done => { resolve = done; }));
+    usePromptSuggestionStore.getState().setSuggestion('session', candidate);
+    act(() => root.render(<Composer />));
+    expect(container.textContent).toBe('');
+    await act(async () => resolve({}));
+    expect(container.firstElementChild?.getAttribute('data-enabled')).toBe('false');
+    expect(usePromptSuggestionStore.getState().suggestions).toEqual({});
+  });
+
+  it('does not apply one host preference or late read to another host', async () => {
+    let resolveLocal!: (value: unknown) => void;
+    let resolvePeer!: (value: unknown) => void;
+    service.read.mockReturnValueOnce(new Promise(done => { resolveLocal = done; }))
+      .mockReturnValueOnce(new Promise(done => { resolvePeer = done; }));
+    usePromptSuggestionStore.getState().setSuggestion('session', candidate, 'local');
+    usePromptSuggestionStore.getState().setSuggestion('session', candidate, 'peer');
+    act(() => root.render(<Composer />));
+    expect(container.textContent).toBe('Run the tests');
+    act(() => activateSurface('peer'));
+    expect(container.textContent).toBe('');
+    await act(async () => resolveLocal({ enable_prompt_suggestion: false }));
+    expect(usePromptSuggestionStore.getState().enabledBySurface.peer).toBeUndefined();
+    await act(async () => resolvePeer({ enable_prompt_suggestion: false }));
+    expect(usePromptSuggestionStore.getState().suggestions[promptSuggestionKey('session', 'peer')]).toBeUndefined();
+    expect(usePromptSuggestionStore.getState().suggestions[promptSuggestionKey('session', 'local')]).toEqual(candidate);
+  });
+});