[Payment due @shubham1206agra] Render thinking + draft-streaming UI for custom agent chats

src/hooks/useAgentZeroStatusIndicator.ts

@@ -75,12 +75,17 @@ function selectNewestReportAction(reportActions: OnyxEntry<ReportActions>): Newe
  * Hook to manage AgentZero status indicator for chats where AgentZero responds.
  *
  * Callers must gate this hook at the mount level (only mount for AgentZero-enabled chats:
- * Concierge DMs or policy #admins rooms). The outer `AgentZeroStatusProvider` already
- * enforces this, so the hook assumes it's always running for an AgentZero chat.
+ * Concierge DMs, policy #admins rooms, or custom-agent chats). The outer
+ * `AgentZeroStatusProvider` already enforces this, so the hook assumes it's always running
+ * for an AgentZero chat.
  *
  * @param reportID - The report ID to monitor
+ * @param personaAccountID - The persona handling this chat (Concierge for Concierge/admin chats;
+ *   the agent's accountID for custom-agent chats). Used to decide when a final reply has
+ *   actually landed: the indicator only clears once the newest reportAction's actorAccountID
+ *   matches this persona AND the server NVP signals done.
  */
-function useAgentZeroStatusIndicator(reportID: string): AgentZeroStatusState {
+function useAgentZeroStatusIndicator(reportID: string, personaAccountID: number = CONST.ACCOUNT_ID.CONCIERGE): AgentZeroStatusState {
     // Server-driven processing label from report name-value pairs (e.g. "Looking up categories...")
     // Uses selector to only re-render when the specific field changes, not on any NVP change.
     const [serverLabel] = useOnyx(`${ONYXKEYS.COLLECTION.REPORT_NAME_VALUE_PAIRS}${reportID}`, {selector: agentZeroProcessingIndicatorSelector});
@@ -377,7 +382,7 @@ function useAgentZeroStatusIndicator(reportID: string): AgentZeroStatusState {
     const newestActorAccountID = newestReportAction?.actorAccountID;
     const newestActionID = newestReportAction?.reportActionID;
     useEffect(() => {
-        if (newestActorAccountID !== CONST.ACCOUNT_ID.CONCIERGE) {
+        if (newestActorAccountID !== personaAccountID) {
             return;
         }
         if (pendingOptimisticRequests === 0 && !serverLabel) {
@@ -386,15 +391,15 @@ function useAgentZeroStatusIndicator(reportID: string): AgentZeroStatusState {
         if (!newestActionID || newestActionID === indicatorBaselineActionIDRef.current) {
             return;
         }
-        // Server hasn't signaled done yet — this is an intermediate Concierge action, not
+        // Server hasn't signaled done yet — this is an intermediate persona action, not
         // the final reply. Wait for the NVP to clear before tearing everything down.
         if (serverLabel) {
             return;
         }
         clearAgentZeroProcessingIndicator(reportID);
         clearSafetyTimer();
         AgentZeroOptimisticStore.clear(reportID);
-    }, [newestActorAccountID, newestActionID, serverLabel, pendingOptimisticRequests, reportID, clearSafetyTimer]);
+    }, [newestActorAccountID, newestActionID, serverLabel, pendingOptimisticRequests, reportID, clearSafetyTimer, personaAccountID]);
 
     const isProcessing = !isOffline && isIndicatorActive;
 

src/libs/Navigation/AppNavigator/AuthScreensInitHandler.tsx

@@ -18,6 +18,7 @@ import {getReportIDFromLink} from '@libs/ReportUtils';
 import * as SessionUtils from '@libs/SessionUtils';
 import {endSpan, getSpan, startSpan} from '@libs/telemetry/activeSpans';
 import {getSearchParamFromUrl} from '@libs/Url';
+import {openAgentsPage} from '@userActions/Agent';
 import * as App from '@userActions/App';
 import * as Download from '@userActions/Download';
 import {clearStaleExportDownloads} from '@userActions/Export';
@@ -128,6 +129,11 @@ function AuthScreensInitHandler() {
             App.reconnectApp(initialLastUpdateIDAppliedToClient);
         }
 
+        // Hydrate the user's custom-agent prompts so AgentZeroStatusProvider can recognize
+        // custom-agent chats opened directly from a deep link or right after sign-in, without
+        // requiring a prior visit to Settings > Agents.
+        openAgentsPage();
+
         App.setUpPoliciesAndNavigate(
             session,
             introSelected,

src/libs/Pusher/types.ts

@@ -51,6 +51,12 @@ type ConciergeDraftEvent = {
     startedAt?: string;
     terminalReason?: string;
     updatedAt?: string;
+    /**
+     * Persona accountID the streamed draft should be attributed to — Concierge for Concierge
+     * runs, the custom agent's accountID for agent runs. Optional for backward compatibility;
+     * absent payloads default to Concierge.
+     */
+    actorAccountID?: number;
 };
 
 type PusherEventMap = {

src/pages/home/report/ConciergeThinkingMessage.tsx

@@ -35,7 +35,7 @@ type ConciergeThinkingMessageProps = {
 };
 
 function ConciergeThinkingMessage({report, action}: ConciergeThinkingMessageProps) {
-    const {isProcessing, reasoningHistory, statusLabel} = useAgentZeroStatus();
+    const {isProcessing, reasoningHistory, statusLabel, personaAccountID} = useAgentZeroStatus();
     const shouldSuppress = useShouldSuppressConciergeIndicators(report?.reportID);
 
     if (!isProcessing || shouldSuppress) {
@@ -48,6 +48,7 @@ function ConciergeThinkingMessage({report, action}: ConciergeThinkingMessageProp
             action={action}
             reasoningHistory={reasoningHistory}
             statusLabel={statusLabel}
+            personaAccountID={personaAccountID}
         />
     );
 }
@@ -57,11 +58,13 @@ function ConciergeThinkingMessageContent({
     action,
     reasoningHistory,
     statusLabel,
+    personaAccountID,
 }: {
     report: OnyxEntry<Report>;
     action?: OnyxEntry<ReportAction>;
     reasoningHistory: ReasoningEntry[];
     statusLabel: string;
+    personaAccountID: number;
 }) {
     const styles = useThemeStyles();
     const theme = useTheme();
@@ -116,7 +119,7 @@ function ConciergeThinkingMessageContent({
     }));
 
     const [personalDetails] = useOnyx(ONYXKEYS.PERSONAL_DETAILS_LIST);
-    const accountID = action?.actorAccountID ?? CONST.ACCOUNT_ID.CONCIERGE;
+    const accountID = action?.actorAccountID ?? personaAccountID;
     const displayName = action?.person?.[0]?.text ?? getDisplayNameOrDefault(personalDetails?.[accountID]) ?? CONST.CONCIERGE_DISPLAY_NAME;
     const actorIcon = personalDetails?.[accountID]?.avatar ? {source: personalDetails[accountID].avatar, name: displayName, type: CONST.ICON_TYPE_AVATAR} : undefined;
 
@@ -156,7 +159,7 @@ function ConciergeThinkingMessageContent({
                         reportID={report?.reportID}
                         chatReportID={report?.chatReportID ?? report?.reportID}
                         action={action}
-                        accountIDs={[CONST.ACCOUNT_ID.CONCIERGE]}
+                        accountIDs={[accountID]}
                     />
                 </OfflineWithFeedback>
             </View>

src/pages/inbox/AgentZeroStatusContext.tsx

@@ -1,3 +1,4 @@
+import {getAgentAccountIDFlags, getReportParticipantAccountIDs} from '@selectors/AgentZeroChat';
 import {getReportChatType} from '@selectors/Report';
 import React, {createContext, useContext, useEffect} from 'react';
 import useAgentZeroStatusIndicator from '@hooks/useAgentZeroStatusIndicator';
@@ -16,6 +17,13 @@ type AgentZeroStatusState = {
 
     /** Debounced label shown in the thinking bubble */
     statusLabel: string;
+
+    /**
+     * The accountID of the AgentZero persona handling this chat. Concierge for Concierge DMs and
+     * #admins rooms; the agent's own accountID for custom-agent chats. Consumers use it to render
+     * the thinking-bubble avatar and to decide when a reply has actually landed.
+     */
+    personaAccountID: number;
 };
 
 type AgentZeroStatusActions = {
@@ -27,6 +35,7 @@ const defaultState: AgentZeroStatusState = {
     isProcessing: false,
     reasoningHistory: [],
     statusLabel: '',
+    personaAccountID: CONST.ACCOUNT_ID.CONCIERGE,
 };
 
 const defaultActions: AgentZeroStatusActions = {
@@ -37,17 +46,26 @@ const AgentZeroStatusStateContext = createContext<AgentZeroStatusState>(defaultS
 const AgentZeroStatusActionsContext = createContext<AgentZeroStatusActions>(defaultActions);
 
 /**
- * Cheap outer guard — only subscribes to the scalar CONCIERGE_REPORT_ID.
- * For non-AgentZero reports (the common case), returns children directly.
+ * Cheap outer guard — only subscribes to the scalar CONCIERGE_REPORT_ID and the report's chat
+ * metadata. For non-AgentZero reports (the common case), returns children directly.
  *
- * AgentZero chats include Concierge DMs and policy #admins rooms.
+ * AgentZero chats include Concierge DMs, policy #admins rooms, and custom-agent chats (any
+ * report with a participant whose accountID has a `SHARED_NVP_AGENT_PROMPT_<accountID>` entry,
+ * populated by `OpenAgentsPage` for agents the current user owns).
  */
 function AgentZeroStatusProvider({reportID, children}: React.PropsWithChildren<{reportID: string | undefined}>) {
     const [chatType] = useOnyx(`${ONYXKEYS.COLLECTION.REPORT}${reportID}`, {selector: getReportChatType});
+    const [participantAccountIDs] = useOnyx(`${ONYXKEYS.COLLECTION.REPORT}${reportID}`, {selector: getReportParticipantAccountIDs});
+    const [agentAccountIDFlags] = useOnyx(ONYXKEYS.COLLECTION.SHARED_NVP_AGENT_PROMPT, {selector: getAgentAccountIDFlags});
     const [conciergeReportID] = useOnyx(ONYXKEYS.CONCIERGE_REPORT_ID);
+
     const isConciergeChat = reportID === conciergeReportID;
     const isAdmin = chatType === CONST.REPORT.CHAT_TYPE.POLICY_ADMINS;
-    const isAgentZeroChat = isConciergeChat || isAdmin;
+    // First participant whose accountID has a SHARED_NVP_AGENT_PROMPT entry. Both gates and
+    // identifies the persona in one pass.
+    const agentParticipantAccountID = participantAccountIDs?.find((accountID) => !!agentAccountIDFlags?.[accountID]);
+    const isCustomAgentChat = agentParticipantAccountID !== undefined;
+    const isAgentZeroChat = isConciergeChat || isAdmin || isCustomAgentChat;
 
     if (!reportID || !isAgentZeroChat) {
         return children;
@@ -57,14 +75,16 @@ function AgentZeroStatusProvider({reportID, children}: React.PropsWithChildren<{
         <AgentZeroStatusGate
             key={reportID}
             reportID={reportID}
+            personaAccountID={agentParticipantAccountID ?? CONST.ACCOUNT_ID.CONCIERGE}
         >
             {children}
         </AgentZeroStatusGate>
     );
 }
 
-function AgentZeroStatusGate({reportID, children}: React.PropsWithChildren<{reportID: string}>) {
-    const {kickoffWaitingIndicator, ...stateValue} = useAgentZeroStatusIndicator(reportID);
+function AgentZeroStatusGate({reportID, personaAccountID, children}: React.PropsWithChildren<{reportID: string; personaAccountID: number}>) {
+    const {kickoffWaitingIndicator, ...indicatorState} = useAgentZeroStatusIndicator(reportID, personaAccountID);
+    const stateValue = {...indicatorState, personaAccountID};
     const actionsValue = {kickoffWaitingIndicator};
 
     // Auto-kickoff "thinking" indicator when opened from search (where kickoffWaitingIndicator isn't accessible)

src/pages/inbox/ConciergeDraftContext.tsx

@@ -1,3 +1,4 @@
+import {getAgentAccountIDFlags, getReportParticipantAccountIDs} from '@selectors/AgentZeroChat';
 import {getReportChatType} from '@selectors/Report';
 import React, {createContext, useContext, useEffect, useState} from 'react';
 import useOnyx from '@hooks/useOnyx';
@@ -41,10 +42,16 @@ const ConciergeDraftActionsContext = createContext<ConciergeDraftActions>(defaul
 
 function ConciergeDraftProvider({reportID, children}: React.PropsWithChildren<{reportID: string | undefined}>) {
     const [chatType] = useOnyx(`${ONYXKEYS.COLLECTION.REPORT}${reportID}`, {selector: getReportChatType});
+    const [participantAccountIDs] = useOnyx(`${ONYXKEYS.COLLECTION.REPORT}${reportID}`, {selector: getReportParticipantAccountIDs});
+    const [agentAccountIDFlags] = useOnyx(ONYXKEYS.COLLECTION.SHARED_NVP_AGENT_PROMPT, {selector: getAgentAccountIDFlags});
     const [conciergeReportID] = useOnyx(ONYXKEYS.CONCIERGE_REPORT_ID);
+
     const isConciergeChat = reportID === conciergeReportID;
     const isAdmin = chatType === CONST.REPORT.CHAT_TYPE.POLICY_ADMINS;
-    const isAgentZeroChat = isConciergeChat || isAdmin;
+    // See AgentZeroStatusContext for the rationale: agentAccountIDFlags reflects agents the
+    // user owns (populated by `OpenAgentsPage`).
+    const isCustomAgentChat = participantAccountIDs?.some((accountID) => !!agentAccountIDFlags?.[accountID]);
+    const isAgentZeroChat = isConciergeChat || isAdmin || isCustomAgentChat;
 
     if (!reportID || !isAgentZeroChat) {
         return children;

src/pages/inbox/conciergeDraftState.ts

@@ -13,6 +13,7 @@ type ConciergeDraft = {
 };
 
 type BuildConciergeDraftReportActionParams = {
+    actorAccountID?: number;
     bodyMarkdown?: string;
     created: string;
     finalRenderedHTML?: string;
@@ -182,18 +183,21 @@ function stripIncompleteMarkdown(markdown: string): string {
     return result;
 }
 
-function buildConciergeDraftReportAction({bodyMarkdown, created, finalRenderedHTML, reportActionID, reportID}: BuildConciergeDraftReportActionParams): ReportAction | null {
+function buildConciergeDraftReportAction({actorAccountID, bodyMarkdown, created, finalRenderedHTML, reportActionID, reportID}: BuildConciergeDraftReportActionParams): ReportAction | null {
     const html = finalRenderedHTML ?? (bodyMarkdown ? getParsedComment(stripIncompleteMarkdown(bodyMarkdown), {reportID}) : '');
 
     if (!html) {
         return null;
     }
 
+    // Default to Concierge so existing call sites that don't pass an actor stay byte-identical.
+    const resolvedActorAccountID = actorAccountID ?? CONST.ACCOUNT_ID.CONCIERGE;
+
     return {
         reportActionID,
         reportID,
         actionName: CONST.REPORT.ACTIONS.TYPE.ADD_COMMENT,
-        actorAccountID: CONST.ACCOUNT_ID.CONCIERGE,
+        actorAccountID: resolvedActorAccountID,
         person: [{style: 'strong', text: CONST.CONCIERGE_DISPLAY_NAME, type: 'TEXT'}],
         created,
         message: [{type: CONST.REPORT.MESSAGE.TYPE.COMMENT, html, text: Parser.htmlToText(html)}],
@@ -243,6 +247,7 @@ function applyConciergeDraftEvent(currentDraft: ConciergeDraft | null, event: Co
 
     const nextReportAction =
         buildConciergeDraftReportAction({
+            actorAccountID: event.actorAccountID,
             bodyMarkdown: event.bodyMarkdown,
             created: event.created,
             finalRenderedHTML: event.finalRenderedHTML,

src/selectors/AgentZeroChat.ts

@@ -0,0 +1,26 @@
+import type {OnyxCollection, OnyxEntry} from 'react-native-onyx';
+import ONYXKEYS from '@src/ONYXKEYS';
+import type {AgentPrompt} from '@src/types/onyx';
+import type Report from '@src/types/onyx/Report';
+
+const getReportParticipantAccountIDs = (report: OnyxEntry<Report>): number[] => (report?.participants ? Object.keys(report.participants).map(Number) : []);
+
+/**
+ * Reduces the SHARED_NVP_AGENT_PROMPT collection to a `Record<accountID, true>` so callers
+ * can do O(1) lookups without re-rendering on every prompt-content edit.
+ */
+const getAgentAccountIDFlags = (agentPrompts: OnyxCollection<AgentPrompt>): Record<number, true> => {
+    if (!agentPrompts) {
+        return {};
+    }
+    const flags: Record<number, true> = {};
+    for (const key of Object.keys(agentPrompts)) {
+        const accountID = Number(key.slice(ONYXKEYS.COLLECTION.SHARED_NVP_AGENT_PROMPT.length));
+        if (!Number.isNaN(accountID)) {
+            flags[accountID] = true;
+        }
+    }
+    return flags;
+};
+
+export {getReportParticipantAccountIDs, getAgentAccountIDFlags};

tests/ui/AuthScreensInitHandlerTest.tsx

@@ -72,6 +72,10 @@ jest.mock('@userActions/App', () => ({
     setLocale: jest.fn(),
 }));
 
+jest.mock('@userActions/Agent', () => ({
+    openAgentsPage: jest.fn(),
+}));
+
 jest.mock('@userActions/Download', () => ({
     clearDownloads: jest.fn(),
 }));

tests/unit/ConciergeThinkingMessageAvatarTest.tsx

@@ -19,12 +19,18 @@ jest.mock('@components/ReportActionAvatars', () => {
     };
 });
 
-// Mock the AgentZero context to make isProcessing=true so the component renders
+// Admin and announce rooms surface Concierge as the persona, so the mock returns
+// Concierge's accountID here. `mock` prefix lets jest's hoist plugin reference
+// this from inside the factory below.
+const mockPersonaAccountID = CONST.ACCOUNT_ID.CONCIERGE;
+
+// Mock the AgentZero context to make isProcessing=true so the component renders.
 jest.mock('@pages/inbox/AgentZeroStatusContext', () => ({
     useAgentZeroStatus: () => ({
         isProcessing: true,
         reasoningHistory: [],
         statusLabel: 'Thinking...',
+        personaAccountID: mockPersonaAccountID,
     }),
 }));