fix(audio): keep the selected command on voice/STT-dictated messages

backend/chainlit/message.py

@@ -58,10 +58,18 @@ def __post_init__(self) -> None:
         if not getattr(self, "id", None):
             self.id = str(uuid.uuid4())
 
+        # Auto-attach the command selected in the UI to user messages the app
+        # creates during an audio turn (e.g. the transcription in on_audio_end).
+        # Deserialized payloads are authoritative about their command and reset
+        # it in from_dict, so a command-less typed message (or resumed history)
+        # never inherits the audio turn's command.
+        if self.type == "user_message" and self.command is None:
+            self.command = context.session.current_command
+
     @classmethod
     def from_dict(self, _dict: StepDict):
         type = _dict.get("type", "assistant_message")
-        return Message(
+        message = Message(
             id=_dict["id"],
             parent_id=_dict.get("parentId"),
             created_at=_dict["createdAt"],
@@ -73,6 +81,12 @@ def from_dict(self, _dict: StepDict):
             language=_dict.get("language"),
             metadata=_dict.get("metadata", {}),
         )
+        # A deserialized payload is authoritative about its command: reset it from
+        # the payload so messages rebuilt here (incoming client messages, thread
+        # resume) never inherit an active audio turn's command that
+        # __post_init__ applies to command-less user messages.
+        message.command = _dict.get("command")
+        return message
 
     def to_dict(self) -> StepDict:
         _dict: StepDict = {

backend/chainlit/session.py

@@ -105,6 +105,7 @@ class BaseSession:
     thread_id_to_resume: Optional[str] = None
     client_type: ClientType
     current_task: Optional[asyncio.Task] = None
+    current_command: Optional[str] = None
     chat_started: bool = False
 
     def __init__(

backend/chainlit/socket.py

@@ -428,13 +428,17 @@ async def window_message(sid, data):
 
 
 @sio.on("audio_start")  # pyright: ignore [reportOptionalCall]
-async def audio_start(sid):
+async def audio_start(sid, payload=None):
     """Handle audio init."""
     session = WebsocketSession.require(sid)
 
     context = init_ws_context(session)
     config: ChainlitConfig = session.get_config()  # type: ignore
 
+    # Remember the command selected in the UI so it can be auto-attached to the
+    # user message produced from the transcribed audio (see Message.__post_init__).
+    session.current_command = payload.get("command") if payload else None
+
     if config.features.audio and config.features.audio.enabled:
         connected = bool(await config.code.on_audio_start())
         connection_state = "on" if connected else "off"
@@ -484,6 +488,8 @@ async def audio_end(sid):
             author="Error", content=str(e) or e.__class__.__name__
         ).send()
     finally:
+        # The command only applies to the audio turn that just ended.
+        session.current_command = None
         await context.emitter.task_end()
 
 

backend/tests/test_message.py

@@ -24,6 +24,9 @@ def mock_chainlit_context(session=None):
     mock_loop = Mock(spec=asyncio.AbstractEventLoop)
     mock_session = session or Mock()
     mock_session.thread_id = "thread_123"
+    if session is None:
+        # Mirror a real session where no command is selected by default.
+        mock_session.current_command = None
 
     with patch("asyncio.get_running_loop", return_value=mock_loop):
         mock_emitter = AsyncMock()
@@ -755,3 +758,75 @@ def test_message_to_dict_with_none_metadata(self):
             result = msg.to_dict()
 
             assert result["metadata"] == {}
+
+
+class TestUserMessageCommandAutoAttach:
+    """Auto-attaching the session command to user messages (e.g. from audio)."""
+
+    @staticmethod
+    def _session_with_command(command):
+        session = Mock()
+        session.thread_id = "thread_123"
+        session.current_command = command
+        return session
+
+    def test_user_message_auto_attaches_current_command(self):
+        """A user message without a command inherits the session command."""
+        session = self._session_with_command("search")
+        with mock_chainlit_context(session=session):
+            msg = Message(content="hello", type="user_message")
+            assert msg.command == "search"
+
+    def test_user_message_keeps_explicit_command(self):
+        """An explicitly provided command is never overridden."""
+        session = self._session_with_command("search")
+        with mock_chainlit_context(session=session):
+            msg = Message(content="hello", type="user_message", command="picture")
+            assert msg.command == "picture"
+
+    def test_user_message_without_current_command_stays_none(self):
+        """No session command means the user message command stays None."""
+        session = self._session_with_command(None)
+        with mock_chainlit_context(session=session):
+            msg = Message(content="hello", type="user_message")
+            assert msg.command is None
+
+    def test_assistant_message_does_not_auto_attach_command(self):
+        """Only user messages inherit the session command."""
+        session = self._session_with_command("search")
+        with mock_chainlit_context(session=session):
+            msg = Message(content="hello", type="assistant_message")
+            assert msg.command is None
+
+    def test_from_dict_user_message_does_not_inherit_session_command(self):
+        """Deserialized payloads (client messages, resume) never inherit it.
+
+        Guards against a concurrent command-less typed message (or resumed
+        history) picking up the active audio turn's command.
+        """
+        session = self._session_with_command("search")
+        step_dict = {
+            "id": "00000000-0000-4000-8000-000000000000",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "output": "typed while an audio turn was active",
+            "name": "User",
+            "type": "user_message",
+        }
+        with mock_chainlit_context(session=session):
+            msg = MessageBase.from_dict(step_dict)
+            assert msg.command is None
+
+    def test_from_dict_user_message_keeps_payload_command(self):
+        """A deserialized command stays authoritative over the session command."""
+        session = self._session_with_command("search")
+        step_dict = {
+            "id": "00000000-0000-4000-8000-000000000000",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "output": "typed with a different command",
+            "name": "User",
+            "type": "user_message",
+            "command": "picture",
+        }
+        with mock_chainlit_context(session=session):
+            msg = MessageBase.from_dict(step_dict)
+            assert msg.command == "picture"

frontend/src/components/chat/MessageComposer/VoiceButton.tsx

@@ -1,5 +1,6 @@
 import { X } from 'lucide-react';
 import { useHotkeys } from 'react-hotkeys-hook';
+import { useRecoilValue } from 'recoil';
 
 import { useAudio, useConfig } from '@chainlit/react-client';
 
@@ -12,6 +13,8 @@ import {
 } from '@/components/ui/tooltip';
 import { Translator } from 'components/i18n';
 
+import { persistentCommandState } from '@/state/chat';
+
 import { Loader } from '../../Loader';
 import { VoiceLines } from '../../icons/VoiceLines';
 import { Button } from '../../ui/button';
@@ -23,6 +26,7 @@ interface Props {
 const VoiceButton = ({ disabled }: Props) => {
   const { config } = useConfig();
   const { startConversation, endConversation, audioConnection } = useAudio();
+  const selectedCommand = useRecoilValue(persistentCommandState);
   const isEnabled = !!config?.features.audio.enabled;
 
   useHotkeys(
@@ -56,13 +60,19 @@ const VoiceButton = ({ disabled }: Props) => {
       }
 
       if (audioConnection === 'on') return endConversation();
-      return startConversation();
+      return startConversation(selectedCommand?.id);
     },
     {
       enableOnFormTags: false,
       preventDefault: false // Don't prevent default - let letters be typed
     },
-    [isEnabled, audioConnection, startConversation, endConversation]
+    [
+      isEnabled,
+      audioConnection,
+      startConversation,
+      endConversation,
+      selectedCommand
+    ]
   );
 
   if (!isEnabled) return null;
@@ -90,7 +100,7 @@ const VoiceButton = ({ disabled }: Props) => {
                 audioConnection === 'on'
                   ? endConversation
                   : audioConnection === 'off'
-                    ? startConversation
+                    ? () => startConversation(selectedCommand?.id)
                     : undefined
               }
             >

libs/react-client/src/useAudio.ts

@@ -18,10 +18,13 @@ const useAudio = () => {
 
   const { startAudioStream, endAudioStream } = useChatInteract();
 
-  const startConversation = useCallback(async () => {
-    setAudioConnection('connecting');
-    await startAudioStream();
-  }, [startAudioStream]);
+  const startConversation = useCallback(
+    async (command?: string) => {
+      setAudioConnection('connecting');
+      await startAudioStream(command);
+    },
+    [startAudioStream]
+  );
 
   const endConversation = useCallback(async () => {
     setAudioConnection('off');

libs/react-client/src/useChatInteract.ts

@@ -133,9 +133,12 @@ const useChatInteract = () => {
     [session?.socket]
   );
 
-  const startAudioStream = useCallback(() => {
-    session?.socket.emit('audio_start');
-  }, [session?.socket]);
+  const startAudioStream = useCallback(
+    (command?: string) => {
+      session?.socket.emit('audio_start', { command });
+    },
+    [session?.socket]
+  );
 
   const sendAudioChunk = useCallback(
     (