From 8d2c70689aedb9ab1901e7bbff5246e3cf840478 Mon Sep 17 00:00:00 2001
From: Felarof <nithin.sonti@gmail.com>
Date: Wed, 14 Jan 2026 17:30:39 -0800
Subject: [PATCH] feat: fixes to the jtbd agent (#231)

* feat: add support for jtbd agent to accept max turns and experiment id as query params

* fix: add jtbd agent integration with workflow

* fix: change message threshold to 5
---
 apps/agent/entrypoints/options/App.tsx         | 15 ++++++++++++++-
 .../options/create-graph/GraphChat.tsx         | 12 ++++++++++--
 .../options/jtbd-agent/SurveyPage.tsx          |  9 +++++++--
 .../options/jtbd-agent/useSurveyChat.ts        | 17 +++++++++++++----
 .../entrypoints/sidepanel/index/ChatError.tsx  |  2 +-
 apps/agent/lib/jtbd-popup/constants.ts         |  2 +-
 apps/agent/lib/jtbd-popup/useJtbdPopup.ts      | 18 ++++++++++++------
 7 files changed, 58 insertions(+), 17 deletions(-)

diff --git a/apps/agent/entrypoints/options/App.tsx b/apps/agent/entrypoints/options/App.tsx
index 7ac65db4..7e646d20 100644
--- a/apps/agent/entrypoints/options/App.tsx
+++ b/apps/agent/entrypoints/options/App.tsx
@@ -19,8 +19,18 @@ function getInitialRoute(): string {
   return '/ai'
 }
 
+// Get survey params from query string
+function getSurveyParams(): { maxTurns?: number; experimentId?: string } {
+  const params = new URLSearchParams(window.location.search)
+  const maxTurnsStr = params.get('maxTurns')
+  const experimentId = params.get('experimentId') ?? 'default'
+  const maxTurns = maxTurnsStr ? Number.parseInt(maxTurnsStr, 10) : 7
+  return { maxTurns, experimentId }
+}
+
 export const App: FC = () => {
   const initialRoute = getInitialRoute()
+  const surveyParams = getSurveyParams()
 
   return (
     <HashRouter>
@@ -40,7 +50,10 @@ export const App: FC = () => {
             />
             <Route path="scheduled" element={<ScheduledTasksPage />} />
             <Route path="workflows" element={<WorkflowsPage />} />
-            <Route path="jtbd-agent" element={<SurveyPage />} />
+            <Route
+              path="jtbd-agent"
+              element={<SurveyPage {...surveyParams} />}
+            />
           </Route>
           <Route path="create-graph" element={<CreateGraph />} />
         </Routes>
diff --git a/apps/agent/entrypoints/options/create-graph/GraphChat.tsx b/apps/agent/entrypoints/options/create-graph/GraphChat.tsx
index 83bc1aed..0bbe3ea5 100644
--- a/apps/agent/entrypoints/options/create-graph/GraphChat.tsx
+++ b/apps/agent/entrypoints/options/create-graph/GraphChat.tsx
@@ -40,11 +40,14 @@ export const GraphChat: FC<GraphChatProps> = ({
 
   const {
     popupVisible,
+    recordMessageSent,
     triggerIfEligible,
-    onTakeSurvey,
+    onTakeSurvey: onTakeSurveyBase,
     onDismiss: onDismissJtbdPopup,
   } = useJtbdPopup()
 
+  const onTakeSurvey = () => onTakeSurveyBase(5, 'workflow_survey')
+
   // Trigger JTBD popup when AI finishes responding
   const previousChatStatus = useRef(status)
   // biome-ignore lint/correctness/useExhaustiveDependencies: intentionally only trigger on status change
@@ -93,6 +96,11 @@ export const GraphChat: FC<GraphChatProps> = ({
     }))
   }
 
+  const handleSubmit: FormEventHandler<HTMLFormElement> = (e) => {
+    recordMessageSent()
+    onSubmit(e)
+  }
+
   const handleKeyDown = (e: KeyboardEvent<HTMLTextAreaElement>) => {
     if (
       e.key === 'Enter' &&
@@ -128,7 +136,7 @@ export const GraphChat: FC<GraphChatProps> = ({
       {chatError && <ChatError error={chatError} />}
       <div className="shrink-0 border-border/40 border-t bg-background/80 p-2 backdrop-blur-md">
         <form
-          onSubmit={onSubmit}
+          onSubmit={handleSubmit}
           className="relative flex w-full items-end gap-2"
         >
           <textarea
diff --git a/apps/agent/entrypoints/options/jtbd-agent/SurveyPage.tsx b/apps/agent/entrypoints/options/jtbd-agent/SurveyPage.tsx
index 0bd88638..7a211691 100644
--- a/apps/agent/entrypoints/options/jtbd-agent/SurveyPage.tsx
+++ b/apps/agent/entrypoints/options/jtbd-agent/SurveyPage.tsx
@@ -7,6 +7,11 @@ import { Header } from './SurveyHeader'
 import { Welcome } from './SurveyWelcome'
 import { useChat } from './useSurveyChat'
 
+interface SurveyPageProps {
+  maxTurns?: number
+  experimentId?: string
+}
+
 const ThankYouCard: FC<{ onReset: () => void }> = ({ onReset }) => (
   <div className="rounded-xl border border-border bg-card p-8 text-center shadow-sm">
     <div className="mx-auto mb-4 flex h-16 w-16 items-center justify-center rounded-full bg-green-500/10">
@@ -43,8 +48,8 @@ const ErrorCard: FC<{ error: Error; onRetry: () => void }> = ({
   </div>
 )
 
-export const SurveyPage: FC = () => {
-  const chat = useChat()
+export const SurveyPage: FC<SurveyPageProps> = ({ maxTurns, experimentId }) => {
+  const chat = useChat({ maxTurns, experimentId })
 
   const handleStart = async () => {
     const current = await jtbdPopupStorage.getValue()
diff --git a/apps/agent/entrypoints/options/jtbd-agent/useSurveyChat.ts b/apps/agent/entrypoints/options/jtbd-agent/useSurveyChat.ts
index 49d83bb7..4b4b5a76 100644
--- a/apps/agent/entrypoints/options/jtbd-agent/useSurveyChat.ts
+++ b/apps/agent/entrypoints/options/jtbd-agent/useSurveyChat.ts
@@ -4,7 +4,13 @@ import { BROWSEROS_PREFS } from '@/lib/browseros/prefs'
 
 const JTBD_API_URL = 'https://jtbd-agent.fly.dev'
 // const LOCAL_JTBD_API_URL = 'http://localhost:3001'
-const EXPERIMENT_ID = 'jtbd_jan26'
+const DEFAULT_MAX_TURNS = 7
+const DEFAULT_EXPERIMENT_ID = 'default'
+
+export interface SurveyChatOptions {
+  maxTurns?: number
+  experimentId?: string
+}
 
 async function getInstallId(): Promise<string> {
   try {
@@ -77,7 +83,10 @@ async function* streamSSE(
   }
 }
 
-export function useChat() {
+export function useChat(options: SurveyChatOptions = {}) {
+  const maxTurns = options.maxTurns ?? DEFAULT_MAX_TURNS
+  const experimentId = options.experimentId ?? DEFAULT_EXPERIMENT_ID
+
   const [phase, setPhase] = useState<Phase>('idle')
   const [messages, setMessages] = useState<Message[]>([])
   const [isStreaming, setIsStreaming] = useState(false)
@@ -116,7 +125,7 @@ export function useChat() {
       const response = await fetch(`${JTBD_API_URL}/api/interview/start`, {
         method: 'POST',
         headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({ installId, experimentId: EXPERIMENT_ID }),
+        body: JSON.stringify({ installId, experimentId, maxTurns }),
         signal: abortControllerRef.current.signal,
       })
 
@@ -170,7 +179,7 @@ export function useChat() {
         {
           method: 'POST',
           headers: { 'Content-Type': 'application/json' },
-          body: JSON.stringify({ response: text }),
+          body: JSON.stringify({ response: text, maxTurns }),
           signal: abortControllerRef.current.signal,
         },
       )
diff --git a/apps/agent/entrypoints/sidepanel/index/ChatError.tsx b/apps/agent/entrypoints/sidepanel/index/ChatError.tsx
index 1cc54597..7e5d7ffb 100644
--- a/apps/agent/entrypoints/sidepanel/index/ChatError.tsx
+++ b/apps/agent/entrypoints/sidepanel/index/ChatError.tsx
@@ -61,7 +61,7 @@ export const ChatError: FC<ChatErrorProps> = ({ error, onRetry }) => {
           </a>
           {' or '}
           <a
-            href="/options.html?page=survey"
+            href="/options.html?page=survey&maxTurns=5&experimentId=daily_limit"
             target="_blank"
             rel="noopener noreferrer"
             className="underline hover:text-foreground"
diff --git a/apps/agent/lib/jtbd-popup/constants.ts b/apps/agent/lib/jtbd-popup/constants.ts
index 6dad110d..c547aef3 100644
--- a/apps/agent/lib/jtbd-popup/constants.ts
+++ b/apps/agent/lib/jtbd-popup/constants.ts
@@ -1,6 +1,6 @@
 export const JTBD_POPUP_CONSTANTS = {
   // Show popup after this many messages
-  MESSAGE_THRESHOLD: 15,
+  MESSAGE_THRESHOLD: 5,
   // Show to 1 in N users (samplingId % N === 0)
   // Set to 1 to show to everyone
   SAMPLING_DIVISOR: 1,
diff --git a/apps/agent/lib/jtbd-popup/useJtbdPopup.ts b/apps/agent/lib/jtbd-popup/useJtbdPopup.ts
index 7a3e4273..3c4e635b 100644
--- a/apps/agent/lib/jtbd-popup/useJtbdPopup.ts
+++ b/apps/agent/lib/jtbd-popup/useJtbdPopup.ts
@@ -44,12 +44,18 @@ export function useJtbdPopup() {
     }
   }, [])
 
-  const onTakeSurvey = useCallback(async () => {
-    const current = await jtbdPopupStorage.getValue()
-    track(JTBD_POPUP_CLICKED_EVENT, { messageCount: current.messageCount })
-    setPopupVisible(false)
-    window.open('/options.html?page=survey', '_blank')
-  }, [])
+  const onTakeSurvey = useCallback(
+    async (maxTurns = 15, experimentId = 'popup_survey') => {
+      const current = await jtbdPopupStorage.getValue()
+      track(JTBD_POPUP_CLICKED_EVENT, { messageCount: current.messageCount })
+      setPopupVisible(false)
+      window.open(
+        `/options.html?page=survey&maxTurns=${maxTurns}&experimentId=${experimentId}`,
+        '_blank',
+      )
+    },
+    [],
+  )
 
   const onDismiss = useCallback(async () => {
     const current = await jtbdPopupStorage.getValue()