[Agent Builder] Lift streaming provider + remove 'new' cache key (elastic#267324)

Author: chrisbmar

x-pack/platform/plugins/shared/agent_builder/CONTRIBUTOR_GUIDE.md

@@ -1280,3 +1280,98 @@ setupDeps.agentBuilder.sml.registerType(visualizationSmlType);
 
 The full implementation is ~130 lines and serves as the reference for new types.
 
+## Streams lifecycle (frontend)
+
+The chat streaming layer lives in `public/application/context/send_message/`. This section
+documents how mutations, lifted state, and the React Query cache fit together. Read this
+before touching any of:
+`send_message_context.tsx`, `use_send_message.ts`, `use_send_message_mutation.ts`,
+`use_resume_round_mutation.ts`, `use_subscribe_to_chat_events.ts`,
+`use_is_any_conversation_streaming.ts`.
+
+### The lift
+
+`<SendMessageProvider>` is mounted **once** above the routes/sidebar:
+
+- Routed app: in `mount.tsx`, above `<AgentBuilderRoutes>`. The sidebar is part of the
+  routes, so it can read streaming state directly.
+- Embeddable: in `embeddable_conversations_provider.tsx`, one provider per embeddable
+  instance because each instance has its own `QueryClient`.
+
+The sidebar uses `useIsAnyConversationStreaming()` and `useSendMessageContext()` directly.
+Anything inside the conversation tree should use the per-conversation scoped hook,
+`useSendMessage()` (in `use_send_message.ts`).
+
+### Lifted state
+
+`SendMessageProvider` owns:
+
+- `activeStream: { conversationId, type, agentReasoning } | undefined` — points at the
+  conversation that is currently streaming. Set synchronously when each mutation kicks
+  off; cleared in the mutation's `finally`.
+- `byConversationId: Record<string, StreamRecord>` — per-conversation `pendingMessage`,
+  `error`, `errorSteps`. Persists across stream end so a user can hit Retry after a
+  failure.
+
+### Mutations: single-scope `mutationFn`
+
+`useSendMessageMutation` and `useResumeRoundMutation` use a **single-scope `mutationFn`
+with `try / catch / finally`**, not React Query's lifecycle methods (`onMutate`,
+`onSuccess`, `onError`, `onSettled`). The shape is:
+
+```ts
+mutationFn: async (vars) => {
+  // setup phase (sync, before any await): seed the optimistic round, set pending message.
+  // Note: `activeStream` is set by the provider's `mutateSendMessage` wrapper *before*
+  // `mutate()` returns — not here.
+  const streamActions = createConversationActions({ conversationId: vars.conversationId, ... });
+
+  try {
+    await subscribeToChatEvents({ events$, conversationActions: streamActions, ... });
+    // success cleanup
+  } catch (err) {
+    // error cleanup
+    throw err;
+  } finally {
+    // cleanup: invalidate cache (skipped if round paused on HITL),
+    // clear `activeStream`, clear abort ref.
+  }
+}
+```
+
+**Why not lifecycle methods?** Streams aren't typical mutations — the bulk of the work
+happens *during* `mutationFn`, with state mutations flowing for many seconds. Splitting
+the work across lifecycle callbacks forces you to bridge state between scopes via refs
+or React Query's `context` return — neither is clean for streaming. With single-scope,
+`streamActions` and `vars` are visible throughout. No refs to bridge phases. Reads
+top-to-bottom.
+
+### Each conversation owns its streaming lifecycle
+
+Every `mutationFn` invocation builds its **own** `streamActions` instance via
+`createConversationActions({ conversationId: vars.conversationId, ... })`. That instance
+is closure-bound to the mutation's conversation id. **Stream events target the
+conversation the mutation was started for, regardless of where the user has navigated.**
+
+If the user submits on conversation A and immediately switches to B, the stream events
+keep writing to A's cache. B loads cleanly from the server.
+
+### Per-conversation `useConversation` gate
+
+`useConversation` is disabled for a conversation when (a) a stream is currently writing
+to its cache, or (b) the cache shows it's paused on a HITL prompt. The cache is
+authoritative in both cases, so a refetch would race with optimistic chunks (streaming)
+or with the resume mutation about to fire on Approve (HITL). Other conversations stay
+free to refetch — switch to conversation B while A streams and B loads cleanly. See the
+inline comment on the `enabled` predicate for details.
+
+### Single-stream vs concurrent streams
+
+Today the app enforces single-stream-at-a-time. The global gates (HITL Approve,
+submit button, page-leave guard) all read `useIsAnyConversationStreaming()`.
+
+The architecture supports concurrent streams in principle — per-conversation cache,
+mutation-scoped `streamActions`, lifted `byConversationId`. The follow-up PR removes
+the global gates, moves the abort controller into a per-conversation slot so each
+stream can be cancelled independently, and enables concurrent streams.
+

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation.tsx

@@ -26,6 +26,7 @@ import { NewConversationPrompt } from './new_conversation_prompt';
 import { useConversationId } from '../../context/conversation/use_conversation_id';
 import { useShouldStickToBottom } from '../../context/conversation/use_should_stick_to_bottom';
 import { useSendMessage } from '../../context/send_message/send_message_context';
+import { useIsAnyConversationStreaming } from '../../hooks/use_is_any_conversation_streaming';
 import { useConversationScrollActions } from '../../hooks/use_conversation_scroll_actions';
 import { useConversationStatus } from '../../hooks/use_conversation';
 import { useSendPredefinedInitialMessage } from '../../hooks/use_initial_message';
@@ -53,6 +54,7 @@ export const Conversation: React.FC<{}> = () => {
   const conversationId = useConversationId();
   const hasActiveConversation = useHasActiveConversation();
   const { isResponseLoading } = useSendMessage();
+  const isAnyStreaming = useIsAnyConversationStreaming();
   const conversationRounds = useConversationRounds();
   const lastRound = conversationRounds.at(-1);
   const { isFetched } = useConversationStatus();
@@ -65,9 +67,10 @@ export const Conversation: React.FC<{}> = () => {
   const [dismissStaleAttachments, setDismissStaleAttachments] = useState(false);
   useSendPredefinedInitialMessage();
 
+  // Page-leave guard fires for any in-flight stream, not just this conversation's.
   useNavigationAbort({
     onAppLeave,
-    isResponseLoading,
+    isResponseLoading: isAnyStreaming,
   });
 
   const scrollContainerRef = useRef<HTMLDivElement | null>(null);

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_input/conversation_input.tsx

@@ -18,9 +18,12 @@ import type { PropsWithChildren } from 'react';
 import React, { useEffect, useMemo } from 'react';
 import { useConversationId } from '../../../context/conversation/use_conversation_id';
 import { useSendMessage } from '../../../context/send_message/send_message_context';
+import { useSubmitMessage } from '../../../hooks/use_submit_message';
 import { useAgentBuilderAgents } from '../../../hooks/agents/use_agents';
 import { useValidateAgentId } from '../../../hooks/agents/use_validate_agent_id';
-import { useIsSendingMessage } from '../../../hooks/use_is_sending_message';
+// Submit is gated globally on any-conversation streaming until concurrent streams are
+// unblocked in a future PR — at which point it becomes a per-conversation check.
+import { useIsAnyConversationStreaming } from '../../../hooks/use_is_any_conversation_streaming';
 import {
   useAgentId,
   useConversationTitle,
@@ -144,8 +147,8 @@ export const ConversationInput: React.FC<ConversationInputProps> = ({
   onSubmit,
   onEditorFocus,
 }) => {
-  const isSendingMessage = useIsSendingMessage();
-  const { sendMessage, pendingMessage, error, isResuming } = useSendMessage();
+  const isSendingMessage = useIsAnyConversationStreaming();
+  const { pendingMessage, error, isResuming } = useSendMessage();
   const { isFetched } = useAgentBuilderAgents();
   const agentId = useAgentId();
   const conversationId = useConversationId();
@@ -158,6 +161,7 @@ export const ConversationInput: React.FC<ConversationInputProps> = ({
   const isAwaitingPrompt = useIsAwaitingPrompt();
   const { attachments, initialMessage, autoSendInitialMessage, resetInitialMessage } =
     useConversationContext();
+  const submitMessage = useSubmitMessage();
 
   const validateAgentId = useValidateAgentId();
   const isAgentIdValid = validateAgentId(agentId);
@@ -241,7 +245,7 @@ export const ConversationInput: React.FC<ConversationInputProps> = ({
       }
       return;
     }
-    sendMessage({ message: content });
+    submitMessage(content);
     messageEditorController.clear();
     onSubmit?.();
   };

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_input/input_actions/connector_selector/connector_selector.test.tsx

@@ -18,8 +18,8 @@ jest.mock('../../../../../hooks/use_kibana', () => ({
   useKibana: jest.fn(),
 }));
 
-jest.mock('../../../../../context/send_message/send_message_context', () => ({
-  useSendMessage: jest.fn(),
+jest.mock('../../../../../hooks/chat/use_connector_selection', () => ({
+  useConnectorSelection: jest.fn(),
 }));
 
 jest.mock('../../../../../hooks/chat/use_default_connector', () => ({
@@ -74,13 +74,15 @@ jest.mock('./connector_icon', () => ({
 
 import { useLoadConnectors } from '@kbn/inference-connectors';
 import { useKibana } from '../../../../../hooks/use_kibana';
-import { useSendMessage } from '../../../../../context/send_message/send_message_context';
+import { useConnectorSelection } from '../../../../../hooks/chat/use_connector_selection';
 import { useDefaultConnector } from '../../../../../hooks/chat/use_default_connector';
 import { ConnectorSelector } from './connector_selector';
 
 const mockUseLoadConnectors = useLoadConnectors as jest.MockedFunction<typeof useLoadConnectors>;
 const mockUseKibana = useKibana as jest.MockedFunction<typeof useKibana>;
-const mockUseSendMessage = useSendMessage as jest.MockedFunction<typeof useSendMessage>;
+const mockUseConnectorSelection = useConnectorSelection as jest.MockedFunction<
+  typeof useConnectorSelection
+>;
 const mockUseDefaultConnector = useDefaultConnector as jest.MockedFunction<
   typeof useDefaultConnector
 >;
@@ -137,14 +139,12 @@ const setup = ({
 
   const selectConnector = jest.fn();
 
-  mockUseSendMessage.mockReturnValue({
-    connectorSelection: {
-      selectedConnector,
-      selectConnector,
-      defaultConnectorId,
-      defaultConnectorOnly,
-    },
-  } as any);
+  mockUseConnectorSelection.mockReturnValue({
+    selectedConnector,
+    selectConnector,
+    defaultConnectorId,
+    defaultConnectorOnly,
+  });
 
   const utils = render(
     <IntlProvider locale="en">
@@ -154,17 +154,15 @@ const setup = ({
   return {
     ...utils,
     selectConnector,
-    // Helper to re-render with a new send-message context (simulates admin changing a setting).
+    // Helper to re-render with a new connector selection (simulates admin changing a setting).
     updateContext: (next: Partial<RenderOptions>) => {
-      mockUseSendMessage.mockReturnValue({
-        connectorSelection: {
-          selectedConnector: next.selectedConnector ?? selectedConnector,
-          selectConnector,
-          defaultConnectorId:
-            'defaultConnectorId' in next ? next.defaultConnectorId : defaultConnectorId,
-          defaultConnectorOnly: next.defaultConnectorOnly ?? defaultConnectorOnly,
-        },
-      } as any);
+      mockUseConnectorSelection.mockReturnValue({
+        selectedConnector: next.selectedConnector ?? selectedConnector,
+        selectConnector,
+        defaultConnectorId:
+          'defaultConnectorId' in next ? next.defaultConnectorId : defaultConnectorId,
+        defaultConnectorOnly: next.defaultConnectorOnly ?? defaultConnectorOnly,
+      });
       act(() => {
         utils.rerender(
           <IntlProvider locale="en">

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_input/input_actions/connector_selector/connector_selector.tsx

@@ -21,7 +21,7 @@ import { FormattedMessage } from '@kbn/i18n-react';
 import React, { useEffect, useMemo, useRef, useState } from 'react';
 import { useUiPrivileges } from '../../../../../hooks/use_ui_privileges';
 import { useNavigation } from '../../../../../hooks/use_navigation';
-import { useSendMessage } from '../../../../../context/send_message/send_message_context';
+import { useConnectorSelection } from '../../../../../hooks/chat/use_connector_selection';
 import { useDefaultConnector } from '../../../../../hooks/chat/use_default_connector';
 import { useKibana } from '../../../../../hooks/use_kibana';
 import {
@@ -165,13 +165,11 @@ export const ConnectorSelector: React.FC<{}> = () => {
     services: { http, settings },
   } = useKibana();
   const {
-    connectorSelection: {
-      selectConnector: onSelectConnector,
-      selectedConnector: selectedConnectorId,
-      defaultConnectorId,
-      defaultConnectorOnly,
-    },
-  } = useSendMessage();
+    selectConnector: onSelectConnector,
+    selectedConnector: selectedConnectorId,
+    defaultConnectorId,
+    defaultConnectorOnly,
+  } = useConnectorSelection();
   const [isPopoverOpen, setIsPopoverOpen] = useState(false);
 
   const { data: aiConnectors, isLoading } = useLoadConnectors({

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_rounds/round_layout.tsx

@@ -21,6 +21,7 @@ import { RoundInput } from './round_input';
 import { RoundThinking } from './round_thinking/round_thinking';
 import { RoundResponse } from './round_response/round_response';
 import { useSendMessage } from '../../../context/send_message/send_message_context';
+import { useIsAnyConversationStreaming } from '../../../hooks/use_is_any_conversation_streaming';
 import { RoundError } from './round_error/round_error';
 import { ConfirmationPrompt } from './round_prompt';
 import { RoundAttachmentReferences } from './round_attachment_references';
@@ -88,6 +89,12 @@ export const RoundLayout: React.FC<RoundLayoutProps> = ({
     resumeRound,
     isResuming,
   } = useSendMessage();
+  // Approve / Cancel for HITL must be gated on global streaming state: while ANY other
+  // conversation is streaming, racing two mutations against the same single-stream
+  // backend would corrupt cache state. This becomes a per-conversation check in the
+  // concurrent-streams follow-up PR.
+  const isAnyStreaming = useIsAnyConversationStreaming();
+  const isHitlDisabled = isAnyStreaming && !isResuming;
 
   const isLoadingCurrentRound = isResponseLoading && isCurrentRound;
   const isErrorCurrentRound = Boolean(error) && isCurrentRound;
@@ -193,6 +200,7 @@ export const RoundLayout: React.FC<RoundLayoutProps> = ({
               onConfirm={() => handlePromptResponse(prompt.id, true)}
               onCancel={() => handlePromptResponse(prompt.id, false)}
               isLoading={isResuming}
+              isDisabled={isHitlDisabled}
               isAnswered={promptResponses[prompt.id] !== undefined}
               answeredValue={promptResponses[prompt.id]?.allow}
             />

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_rounds/round_prompt/confirmation_prompt.tsx

@@ -50,6 +50,7 @@ export interface ConfirmationPromptProps {
   onConfirm: () => void;
   onCancel: () => void;
   isLoading?: boolean;
+  isDisabled?: boolean;
   isAnswered?: boolean;
   answeredValue?: boolean;
 }
@@ -59,6 +60,7 @@ export const ConfirmationPrompt: React.FC<ConfirmationPromptProps> = ({
   onConfirm,
   onCancel,
   isLoading = false,
+  isDisabled = false,
   isAnswered = false,
   answeredValue,
 }) => {
@@ -120,7 +122,7 @@ export const ConfirmationPrompt: React.FC<ConfirmationPromptProps> = ({
           <EuiFlexItem grow={false}>
             <EuiButtonEmpty
               onClick={onCancel}
-              disabled={isLoading || isAnswered}
+              disabled={isDisabled || isLoading || isAnswered}
               size="s"
               color={isAnswered && answeredValue === false ? 'danger' : 'text'}
               data-test-subj="agentBuilderConfirmationPromptCancelButton"
@@ -132,7 +134,7 @@ export const ConfirmationPrompt: React.FC<ConfirmationPromptProps> = ({
             <EuiButton
               onClick={onConfirm}
               isLoading={isLoading}
-              disabled={isAnswered}
+              disabled={isDisabled || isAnswered}
               fill={!isAnswered || answeredValue === true}
               size="s"
               color={isAnswered && answeredValue === true ? 'success' : color}

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_rounds/round_response/chat_message_text.test.tsx

@@ -112,11 +112,9 @@ describe('chat_message_text', () => {
       isEmbeddedContext: false,
       browserApiTools: undefined,
       conversationActions: {
-        removeNewConversationQuery: jest.fn(),
         invalidateConversation: jest.fn(),
         addOptimisticRound: jest.fn(),
         removeOptimisticRound: jest.fn(),
-        setAgentId: jest.fn(),
         addReasoningStep: jest.fn(),
         addToolCall: jest.fn(),
         setToolCallProgress: jest.fn(),

x-pack/platform/plugins/shared/agent_builder/public/application/components/conversations/conversation_rounds/round_response/streaming_text.tsx

@@ -31,10 +31,12 @@ export const StreamingText = ({
   attachmentRefs,
   conversationId,
 }: StreamingTextProps) => {
-  const [displayedText, setDisplayedText] = useState('');
+  // Initial state derives from the content already in the cache so navigating away and back
+  // mid-stream doesn't replay the full text. Only chunks arriving AFTER mount get animated.
+  const [displayedText, setDisplayedText] = useState(content);
   const tokenQueueRef = useRef<string[]>([]);
   const intervalRef = useRef<NodeJS.Timeout | null>(null);
-  const previousContentLengthRef = useRef(0);
+  const previousContentLengthRef = useRef(content.length);
 
   useEffect(() => {
     const previousContentLength = previousContentLengthRef.current;