@servicetitan/titan-chatbot-api 8.0.0 → 9.0.0

package/src/stores/__tests__/chatbot-ui-backend.store.observability.test.ts

@@ -0,0 +1,105 @@
+import { expect } from '@jest/globals';
+import { ILog, Log, LogError, LogInfo, LogWarning } from '@servicetitan/log-service';
+import { Container } from '@servicetitan/react-ioc';
+import { ChatMessageModelText, ChatUiEventListener } from '@servicetitan/titan-chat-ui-common';
+import { CHATBOT_API_CLIENT, IChatbotApiClient, Models, ModelsMocks } from '../../api-client';
+import { ChatbotApiClientMock } from '../../api-client/__mocks__/chatbot-api-client.mock';
+import { AgentStreamHandlers } from '../../streaming';
+import { initTestContainer } from '../../utils/test-utils';
+import { ChatbotUiBackendStore } from '../chatbot-ui-backend.store';
+import { CHATBOT_UI_STORE_TOKEN, ChatbotUiStore } from '../chatbot-ui.store';
+
+const initContainer = initTestContainer(ChatbotUiBackendStore, container => {
+    container
+        .bind<ILog>(Log)
+        .to(
+            class implements ILog {
+                error: (entry: LogError) => void = jest.fn();
+                info: (entry: LogInfo) => void = jest.fn();
+                warning: (entry: LogWarning) => void = jest.fn();
+            }
+        )
+        .inSingletonScope();
+    container
+        .bind<IChatbotApiClient>(CHATBOT_API_CLIENT)
+        .toConstantValue(new ChatbotApiClientMock());
+    container.bind(CHATBOT_UI_STORE_TOKEN).to(ChatbotUiStore).inSingletonScope();
+});
+
+describe('[ChatbotUiBackendStore] streaming observability', () => {
+    let container: Container;
+    let store: ChatbotUiBackendStore;
+    let chatbotApi: ChatbotApiClientMock;
+    let chatUiStore: ChatbotUiStore;
+    let log: ILog;
+
+    const runListener = async (listener: ChatUiEventListener, args: unknown[]) =>
+        new Promise<void>((resolve, reject) => listener.apply(store, [resolve, reject, ...args]));
+
+    const sendStreamed = async () => {
+        chatUiStore.setCustomizationContext({ streaming: { enabled: true } });
+        await chatUiStore.sendMessageText('q');
+        const message = chatUiStore.messages.at(-1)! as ChatMessageModelText;
+        await runListener(store.handleMessageSend, [message]);
+    };
+
+    const codesLogged = () =>
+        (log.info as jest.Mock).mock.calls.map(c => (c[0] as { code: string }).code);
+
+    beforeEach(() => {
+        container = initContainer();
+        store = container.get(ChatbotUiBackendStore);
+        chatbotApi = container.get<IChatbotApiClient>(CHATBOT_API_CLIENT) as ChatbotApiClientMock;
+        chatUiStore = container.get<ChatbotUiStore>(CHATBOT_UI_STORE_TOKEN);
+        log = container.get<ILog>(Log);
+        chatbotApi.getOptions.mockResolvedValue(ModelsMocks.mockFrontendModel());
+        chatbotApi.postSession.mockResolvedValue(ModelsMocks.mockSession());
+    });
+
+    test('logs the connection lifecycle (connected → completed)', async () => {
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onConnected?.();
+                h.onCompleted?.();
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        await sendStreamed();
+
+        const codes = codesLogged();
+        expect(codes).toContain('TitanChatbot_Streaming_connected');
+        expect(codes).toContain('TitanChatbot_Streaming_completed');
+    });
+
+    test('logs disconnected and timed_out lifecycle phases', async () => {
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onConnected?.();
+                h.onDisconnected?.();
+                h.onTimeout?.();
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        await sendStreamed();
+
+        const codes = codesLogged();
+        expect(codes).toContain('TitanChatbot_Streaming_disconnected');
+        expect(codes).toContain('TitanChatbot_Streaming_timed_out');
+    });
+
+    test('tracks fallback separately and does not log a completed streaming run', async () => {
+        chatbotApi.streamMessage.mockRejectedValue(new Error('unreachable'));
+        chatbotApi.postMessage.mockResolvedValue(
+            ModelsMocks.mockBotMessage({ answer: 'fallback' })
+        );
+
+        await sendStreamed();
+
+        expect(store.streamingFallbackCount).toBe(1);
+        const codes = codesLogged();
+        expect(codes).toContain('TitanChatbot_Streaming_fallback');
+        expect(codes).not.toContain('TitanChatbot_Streaming_completed');
+    });
+});

package/src/stores/__tests__/chatbot-ui-backend.store.streaming.test.ts

@@ -0,0 +1,261 @@
+import { expect } from '@jest/globals';
+import { ILog, Log, LogError, LogInfo, LogWarning } from '@servicetitan/log-service';
+import { Container } from '@servicetitan/react-ioc';
+import { ChatMessageModelText, ChatUiEventListener } from '@servicetitan/titan-chat-ui-common';
+import { CHATBOT_API_CLIENT, IChatbotApiClient, Models, ModelsMocks } from '../../api-client';
+import { ChatbotApiClientMock } from '../../api-client/__mocks__/chatbot-api-client.mock';
+import { AgentStreamError, AgentStreamHandlers } from '../../streaming';
+import { initTestContainer } from '../../utils/test-utils';
+import { ChatbotUiBackendStore } from '../chatbot-ui-backend.store';
+import { CHATBOT_UI_STORE_TOKEN, ChatbotUiStore } from '../chatbot-ui.store';
+
+const initContainer = initTestContainer(ChatbotUiBackendStore, container => {
+    container
+        .bind<ILog>(Log)
+        .to(
+            class implements ILog {
+                error: (entry: LogError) => void = jest.fn();
+                info: (entry: LogInfo) => void = jest.fn();
+                warning: (entry: LogWarning) => void = jest.fn();
+            }
+        )
+        .inSingletonScope();
+    container
+        .bind<IChatbotApiClient>(CHATBOT_API_CLIENT)
+        .toConstantValue(new ChatbotApiClientMock());
+    container.bind(CHATBOT_UI_STORE_TOKEN).to(ChatbotUiStore).inSingletonScope();
+});
+
+describe('[ChatbotUiBackendStore] streaming', () => {
+    let container: Container;
+    let store: ChatbotUiBackendStore;
+    let chatbotApi: ChatbotApiClientMock;
+    let chatUiStore: ChatbotUiStore;
+
+    const runListener = async <T = void>(listener: ChatUiEventListener<T>, args: unknown[]) =>
+        new Promise<T>((resolve, reject) => listener.apply(store, [resolve, reject, ...args]));
+
+    const mockSession = () => {
+        chatbotApi.getOptions.mockResolvedValue(ModelsMocks.mockFrontendModel());
+        chatbotApi.postSession.mockResolvedValue(ModelsMocks.mockSession());
+    };
+
+    const enableStreaming = () =>
+        chatUiStore.setCustomizationContext({ streaming: { enabled: true } });
+
+    const sendQuestion = async (text = 'user question') => {
+        await chatUiStore.sendMessageText(text);
+        return chatUiStore.messages.at(-1)! as ChatMessageModelText;
+    };
+
+    beforeEach(() => {
+        container = initContainer();
+        store = container.get(ChatbotUiBackendStore);
+        chatbotApi = container.get<IChatbotApiClient>(CHATBOT_API_CLIENT) as ChatbotApiClientMock;
+        chatUiStore = container.get<ChatbotUiStore>(CHATBOT_UI_STORE_TOKEN);
+        mockSession();
+    });
+
+    test('uses the streaming path when enabled, not postMessage', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockResolvedValue(
+            ModelsMocks.mockBotMessage({ answer: 'streamed' })
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(chatbotApi.streamMessage).toHaveBeenCalledTimes(1);
+        expect(chatbotApi.postMessage).not.toHaveBeenCalled();
+        expect((chatUiStore.messages.at(-1) as ChatMessageModelText).message).toBe('streamed');
+        expect(chatUiStore.isAgentTyping).toBe(false);
+    });
+
+    test('uses the non-streaming path when disabled', async () => {
+        chatbotApi.postMessage.mockResolvedValue(ModelsMocks.mockBotMessage({ answer: 'regular' }));
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(chatbotApi.postMessage).toHaveBeenCalledTimes(1);
+        expect(chatbotApi.streamMessage).not.toHaveBeenCalled();
+        expect((chatUiStore.messages.at(-1) as ChatMessageModelText).message).toBe('regular');
+    });
+
+    test('maps progress events onto the observable progress model', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onStatus?.('Thinking…');
+                h.onText?.('✓ searched');
+                h.onPlan?.([{ id: '1', title: 'Plan' }]);
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(store.streamingProgress.logLines).toContain('✓ searched');
+        expect(store.streamingProgress.steps.map(s => s.id)).toEqual(['1']);
+    });
+
+    test('scrolls the chat after each streaming progress update', async () => {
+        enableStreaming();
+        const triggerScroll = jest.spyOn(chatUiStore, 'triggerScroll');
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                /*
+                 * Each progress event must scroll exactly once (isolated from the scrolls that
+                 * delivering the final answer also performs).
+                 */
+                const scrollsFor = (fn: () => void) => {
+                    triggerScroll.mockClear();
+                    fn();
+                    return triggerScroll.mock.calls.length;
+                };
+                expect(scrollsFor(() => h.onStatus?.('Thinking…'))).toBe(1);
+                expect(scrollsFor(() => h.onText?.('✓ searched'))).toBe(1);
+                expect(scrollsFor(() => h.onPlan?.([{ id: '1', title: 'Plan' }]))).toBe(1);
+                expect(scrollsFor(() => h.onStepActive?.('1'))).toBe(1);
+                expect(scrollsFor(() => h.onInactivity?.())).toBe(1);
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+        expect(triggerScroll).toHaveBeenCalled();
+    });
+
+    test('onStepActive marks the active plan step (earlier done, later pending)', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onPlan?.([
+                    { id: '1', title: 'Look up' },
+                    { id: '2', title: 'Search' },
+                    { id: '3', title: 'Answer' },
+                ]);
+                h.onStepActive?.('2');
+                // Capture mid-run state before run.finished completes all steps.
+                expect(store.streamingProgress.steps.map(s => s.status)).toEqual([
+                    'done',
+                    'active',
+                    'pending',
+                ]);
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+    });
+
+    test('marks all plan steps done once the run finishes', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onPlan?.([
+                    { id: '1', title: 'Look up' },
+                    { id: '2', title: 'Answer' },
+                ]);
+                h.onStepActive?.('1');
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(store.streamingProgress.steps.map(s => s.status)).toEqual(['done', 'done']);
+    });
+
+    test('shows the "Still working on it…" keepalive on inactivity', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onInactivity?.();
+                expect(store.streamingProgress.keepaliveText).toBe('Still working on it…');
+                return ModelsMocks.mockBotMessage({ answer: 'done' });
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+    });
+
+    test('passes the configured inactivity timeout (defaulting to 16000)', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockResolvedValue(ModelsMocks.mockBotMessage({ answer: 'done' }));
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        const handlers = chatbotApi.streamMessage.mock.calls[0][1] as AgentStreamHandlers;
+        expect(handlers.inactivityTimeoutMs).toBe(16_000);
+    });
+
+    test('falls back to non-streaming silently when the stream is unreachable at connect time', async () => {
+        enableStreaming();
+        // Rejects WITHOUT ever calling onConnected → connect-time failure.
+        chatbotApi.streamMessage.mockRejectedValue(new Error('connect failed'));
+        chatbotApi.postMessage.mockResolvedValue(
+            ModelsMocks.mockBotMessage({ answer: 'fallback answer' })
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(chatbotApi.postMessage).toHaveBeenCalledTimes(1);
+        expect((chatUiStore.messages.at(-1) as ChatMessageModelText).message).toBe(
+            'fallback answer'
+        );
+        expect(chatUiStore.isError).toBe(false);
+        expect(store.streamingFallbackCount).toBe(1);
+    });
+
+    test('shows the drop error when the connection fails after connecting', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onConnected?.();
+                throw new Error('socket dropped');
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(chatbotApi.postMessage).not.toHaveBeenCalled();
+        expect(chatUiStore.isError).toBe(true);
+        expect(chatUiStore.error?.message).toContain('Something went wrong during this step');
+    });
+
+    test('surfaces the agent step error message from run.error', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockImplementation(
+            (_m: Models.IUserMessage, h: AgentStreamHandlers) => {
+                h.onConnected?.();
+                throw new AgentStreamError('Search step failed');
+            }
+        );
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        expect(chatUiStore.isError).toBe(true);
+        expect(chatUiStore.error?.message).toContain('Search step failed');
+    });
+
+    test('sends the active session id with the streamed message (session isolation)', async () => {
+        enableStreaming();
+        chatbotApi.streamMessage.mockResolvedValue(ModelsMocks.mockBotMessage({ answer: 'done' }));
+
+        const message = await sendQuestion();
+        await runListener(store.handleMessageSend, [message]);
+
+        const sentBody = chatbotApi.streamMessage.mock.calls[0][0] as Models.IUserMessage;
+        expect(sentBody.sessionId).toBe(store.session?.id);
+    });
+});

package/src/stores/chatbot-ui-backend.store.ts

@@ -9,16 +9,25 @@ import {
     ChatUiEvent,
     ChatUiEventListener,
     IChatUiBackendStore,
+    StreamingProgressModel,
 } from '@servicetitan/titan-chat-ui-common';
 import axios from 'axios';
 import { makeObservable, runInAction } from 'mobx';
 import { CHATBOT_API_CLIENT, IChatbotApiClient, Models, ModelsUtils } from '../api-client';
+import {
+    DEFAULT_STREAMING_INACTIVITY_TIMEOUT_MS,
+    isChatbotStreamingEnabled,
+} from '../models/chatbot-customizations';
+import { AgentStreamError, AgentStreamHandlers } from '../streaming';
 import { withTimeout } from '../utils/axios-utils';
 import { CHATBOT_UI_STORE_TOKEN, ChatbotUiEvent, IChatbotUiStore } from './chatbot-ui.store';
 import { InitializeStore } from './initialize.store';
 
 const DEFAULT_CHATBOT_ASK_BOT_TIMEOUT_MS = 31000; // 31 seconds
 
+const STREAMING_KEEPALIVE_TEXT = 'Still working on it…';
+const STREAMING_STEP_ERROR = 'Something went wrong during this step. Please try again.';
+
 const isAbortedError = (error: unknown): boolean =>
     axios.isCancel(error) || (error as { name?: string } | null)?.name === 'AbortError';
 
@@ -28,6 +37,10 @@ export const CHATBOT_UI_BACKEND_STORE_TOKEN = symbolToken<IChatbotUiBackendStore
 
 export interface IChatbotUiBackendStore extends IChatUiBackendStore {
     session?: Models.ISession;
+    /** Live progress for the in-flight streamed run. */
+    streamingProgress: StreamingProgressModel;
+    /** Count of streaming → non-streaming fallbacks. */
+    streamingFallbackCount: number;
     initialize: () => void;
     saveCurrentState: () => void;
     deleteFromSessionStorage: () => void;
@@ -40,6 +53,12 @@ export interface IChatbotUiBackendStore extends IChatUiBackendStore {
 export class ChatbotUiBackendStore extends InitializeStore implements IChatbotUiBackendStore {
     session?: Models.ISession;
     abortController!: AbortController;
+
+    /** Live progress for the in-flight streamed run (status line, activity log, steps, keepalive). */
+    readonly streamingProgress = new StreamingProgressModel();
+    /** Count of streaming → non-streaming fallbacks, tracked separately for regression detection. */
+    streamingFallbackCount = 0;
+
     @inject(CHATBOT_API_CLIENT) protected readonly chatbotApiClient!: IChatbotApiClient;
     @inject(CHATBOT_UI_STORE_TOKEN) protected readonly chatUiStore!: IChatbotUiStore;
 
@@ -112,7 +131,7 @@ export class ChatbotUiBackendStore extends InitializeStore implements IChatbotUi
         selections: Models.Selections | undefined,
         timeoutMs?: number
     ) => {
-        await this.sendMessage(messageModel, selections, timeoutMs);
+        await this.routeSend(messageModel, selections, timeoutMs);
         resolve();
     };
 
@@ -127,7 +146,7 @@ export class ChatbotUiBackendStore extends InitializeStore implements IChatbotUi
         }
         const selections = messageModel.data?.selections as Models.Selections | undefined;
         const timeoutMs = messageModel.data?.timeoutMs as number | undefined;
-        await this.sendMessage(messageModel as ChatMessageModelText, selections, timeoutMs);
+        await this.routeSend(messageModel as ChatMessageModelText, selections, timeoutMs);
         resolve();
     };
 
@@ -299,6 +318,21 @@ export class ChatbotUiBackendStore extends InitializeStore implements IChatbotUi
         this.chatUiStore.addMessage(true, answer.answer, answer);
     }
 
+    /**
+     * Routes a send to the SSE streaming path when enabled (and supported by the client),
+     * otherwise to the existing non-streaming path. Default is non-streaming.
+     */
+    protected routeSend(
+        messageModel: ChatMessageModelText,
+        selections?: Models.Selections,
+        timeoutMs?: number
+    ): Promise<void> {
+        if (isChatbotStreamingEnabled(this.customizations) && this.chatbotApiClient.streamMessage) {
+            return this.sendMessageStreaming(messageModel, selections, timeoutMs);
+        }
+        return this.sendMessage(messageModel, selections, timeoutMs);
+    }
+
     protected async sendMessage(
         messageModel: ChatMessageModelText,
         selections?: Models.Selections,
@@ -360,6 +394,147 @@ export class ChatbotUiBackendStore extends InitializeStore implements IChatbotUi
         }
     }
 
+    /**
+     * Streaming send-path: consumes the SSE agent-progress stream, rendering live progress and
+     * resolving the final answer on `run.finished`. Applies an inactivity-only timeout (no hard
+     * request timeout). Falls back transparently to {@link sendMessage} when the stream is
+     * unreachable at connect time — propagating the original per-message `timeoutMs` so the
+     * fallback behaves identically to the non-streaming path; surfaces a step error if the
+     * connection fails after connecting.
+     */
+    protected async sendMessageStreaming(
+        messageModel: ChatMessageModelText,
+        selections?: Models.Selections,
+        timeoutMs?: number
+    ) {
+        const question = messageModel.message;
+        const perRequestController = new AbortController();
+        const forwardStoreAbort = () =>
+            perRequestController.abort(this.abortController.signal.reason);
+        this.abortController.signal.addEventListener('abort', forwardStoreAbort, { once: true });
+
+        let connected = false;
+        try {
+            await this.startSession();
+            this.chatUiStore.setAgentTyping(true);
+            runInAction(() => this.streamingProgress.reset());
+
+            const inactivityTimeoutMs =
+                this.customizations?.streaming?.inactivityTimeoutMs ??
+                DEFAULT_STREAMING_INACTIVITY_TIMEOUT_MS;
+
+            // Apply a progress mutation, then scroll so the newest event stays in view.
+            const updateProgress = (mutate: () => void) => {
+                runInAction(mutate);
+                this.chatUiStore.triggerScroll();
+            };
+
+            const handlers: AgentStreamHandlers = {
+                inactivityTimeoutMs,
+                onStatus: text =>
+                    updateProgress(() => {
+                        this.streamingProgress.clearKeepalive();
+                        this.streamingProgress.setStatus(text);
+                    }),
+                onText: text =>
+                    updateProgress(() => {
+                        this.streamingProgress.clearKeepalive();
+                        this.streamingProgress.appendLog(text);
+                    }),
+                onPlan: steps =>
+                    updateProgress(() => {
+                        this.streamingProgress.clearKeepalive();
+                        this.streamingProgress.setSteps(steps);
+                    }),
+                onStepActive: id =>
+                    updateProgress(() => {
+                        this.streamingProgress.clearKeepalive();
+                        this.streamingProgress.setActiveStep(id);
+                    }),
+                onInactivity: () =>
+                    updateProgress(() =>
+                        this.streamingProgress.setKeepalive(STREAMING_KEEPALIVE_TEXT)
+                    ),
+                onConnected: () => {
+                    connected = true;
+                    this.logStreaming('connected');
+                },
+                onDisconnected: () => this.logStreaming('disconnected'),
+                onTimeout: () => this.logStreaming('timed_out'),
+                onCompleted: () => this.logStreaming('completed'),
+            };
+
+            const answer = await this.api(
+                'streamMessage',
+                new Models.UserMessage({
+                    sessionId: this.session!.id!,
+                    question,
+                    experience: Models.Experience.MultiTurn,
+                    selections,
+                }),
+                handlers,
+                perRequestController.signal
+            );
+
+            this.withSaveState(() => {
+                this.chatUiStore.setMessageState(messageModel, ChatMessageState.Delivered, {
+                    answerId: answer.id,
+                    selections,
+                });
+                this.processBotAnswer(answer);
+                this.chatUiStore.resetError(ChatRunState.Started);
+            });
+            /*
+             * Mark every plan step done and clear the transient keepalive; the rendered trail is
+             * faded by the UI once the final answer is shown. The next streamed send resets the model.
+             */
+            runInAction(() => {
+                this.streamingProgress.completeAllSteps();
+                this.streamingProgress.clearKeepalive();
+            });
+        } catch (error: unknown) {
+            if (isAbortedError(error)) {
+                return;
+            }
+
+            // Never connected → the endpoint was unreachable at connect time: fall back silently.
+            if (!connected) {
+                this.streamingFallbackCount += 1;
+                this.logStreaming('fallback');
+                runInAction(() => this.streamingProgress.reset());
+                await this.sendMessage(messageModel, selections, timeoutMs);
+                return;
+            }
+
+            // Connected then failed: surface the step error (the agent's message if provided).
+            this.withSaveState(() => {
+                this.chatUiStore.setMessageState(messageModel, ChatMessageState.Failed, {
+                    selections,
+                });
+            });
+            runInAction(() => this.streamingProgress.reset());
+            const message =
+                error instanceof AgentStreamError ? error.message : STREAMING_STEP_ERROR;
+            if (error instanceof ChatError) {
+                this.setChatError('FailedToSendMessage', error);
+            } else {
+                this.setError('FailedToSendMessage', 'Send message', message, error);
+            }
+        } finally {
+            this.abortController.signal.removeEventListener('abort', forwardStoreAbort);
+            this.chatUiStore.setAgentTyping(false);
+        }
+    }
+
+    /** Logs a streaming connection lifecycle phase (connected, disconnected, timed_out, completed, fallback). */
+    protected logStreaming(phase: string): void {
+        this.logService.info({
+            message: `Streaming ${phase}`,
+            category: 'TitanChatbot',
+            code: `TitanChatbot_Streaming_${phase}`,
+        });
+    }
+
     protected setSession(session?: Models.ISession) {
         this.session = session;
     }
@@ -400,8 +575,8 @@ export class ChatbotUiBackendStore extends InitializeStore implements IChatbotUi
 
     protected async api<
         K extends keyof IChatbotApiClient,
-        Args extends Parameters<IChatbotApiClient[K]>,
-        Result extends ReturnType<IChatbotApiClient[K]>,
+        Args extends Parameters<NonNullable<IChatbotApiClient[K]>>,
+        Result extends ReturnType<NonNullable<IChatbotApiClient[K]>>,
     >(methodName: K, ...args: Args): Promise<Awaited<Result>> {
         const method = this.chatbotApiClient[methodName];
         if (method) {