From 33dd65f677c08df00440ba4585bdf6507cc0b8f0 Mon Sep 17 00:00:00 2001
From: Jinwoo-H <jinwoo0825@gmail.com>
Date: Fri, 12 Jun 2026 19:26:57 -0700
Subject: [PATCH] Fix agent draft paste readiness race

Co-authored-by: Orca <help@stably.ai>
---
 .../src/lib/agent-paste-draft.test.ts         | 46 ++++++++++++++++++-
 src/renderer/src/lib/agent-paste-draft.ts     | 44 +++++++++++++-----
 2 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/src/renderer/src/lib/agent-paste-draft.test.ts b/src/renderer/src/lib/agent-paste-draft.test.ts
index 5b950aae0c5..114d02f4146 100644
--- a/src/renderer/src/lib/agent-paste-draft.test.ts
+++ b/src/renderer/src/lib/agent-paste-draft.test.ts
@@ -1,14 +1,23 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { pasteDraftWhenAgentReady, sendBracketedPasteToRunningAgent } from './agent-paste-draft'
 
+type TestAppState = {
+  settings: Record<string, unknown>
+  ptyIdsByTabId: Record<string, string[]>
+  runtimePaneTitlesByTabId: Record<string, Record<string, string>>
+  tabsByWorktree: Record<string, { id: string; title?: string }[]>
+}
+
 const testState = vi.hoisted(() => ({
   appState: {
     settings: {},
     ptyIdsByTabId: { 'tab-1': ['pty-1'] },
     runtimePaneTitlesByTabId: {},
     tabsByWorktree: {}
-  },
+  } as TestAppState,
   ptyObserver: null as ((data: string) => void) | null,
+  storeSubscriber: null as ((state: TestAppState) => void) | null,
+  unsubscribeStore: vi.fn(),
   unsubscribe: vi.fn(),
   subscribeToPtyData: vi.fn(),
   isRemoteRuntimePtyId: vi.fn(),
@@ -19,7 +28,11 @@ const testState = vi.hoisted(() => ({
 
 vi.mock('@/store', () => ({
   useAppStore: {
-    getState: () => testState.appState
+    getState: () => testState.appState,
+    subscribe: (subscriber: (state: TestAppState) => void) => {
+      testState.storeSubscriber = subscriber
+      return testState.unsubscribeStore
+    }
   }
 }))
 
@@ -54,6 +67,8 @@ describe('pasteDraftWhenAgentReady', () => {
     testState.appState.runtimePaneTitlesByTabId = {}
     testState.appState.tabsByWorktree = {}
     testState.ptyObserver = null
+    testState.storeSubscriber = null
+    testState.unsubscribeStore.mockReset()
     testState.unsubscribe.mockReset()
     testState.subscribeToPtyData.mockReset()
     testState.subscribeToPtyData.mockImplementation(
@@ -126,6 +141,33 @@ describe('pasteDraftWhenAgentReady', () => {
     )
   })
 
+  it('subscribes to PTY output as soon as the tab PTY binding appears', async () => {
+    testState.appState.ptyIdsByTabId = {}
+
+    const promise = pasteDraftWhenAgentReady({
+      tabId: 'tab-1',
+      content: ISSUE_URL,
+      agent: 'codex'
+    })
+    await flushMicrotasks()
+    expect(testState.subscribeToPtyData).not.toHaveBeenCalled()
+
+    testState.appState.ptyIdsByTabId = { 'tab-1': ['pty-1'] }
+    testState.storeSubscriber?.(testState.appState)
+    await flushMicrotasks()
+
+    expect(testState.unsubscribeStore).toHaveBeenCalledTimes(1)
+    expect(testState.subscribeToPtyData).toHaveBeenCalledWith('pty-1', expect.any(Function))
+    testState.ptyObserver?.(`${DECSET_BRACKETED_PASTE}${CODEX_COMPOSER_PROMPT_RENDER}`)
+
+    await expect(promise).resolves.toBe(true)
+    expect(testState.sendRuntimePtyInputVerified).toHaveBeenCalledWith(
+      {},
+      'pty-1',
+      PASTED_ISSUE_URL
+    )
+  })
+
   it('keeps the render-quiet wait for agents without the Codex ready signal', async () => {
     const promise = pasteDraftWhenAgentReady({
       tabId: 'tab-1',
diff --git a/src/renderer/src/lib/agent-paste-draft.ts b/src/renderer/src/lib/agent-paste-draft.ts
index 6515144a38a..0281230b8e8 100644
--- a/src/renderer/src/lib/agent-paste-draft.ts
+++ b/src/renderer/src/lib/agent-paste-draft.ts
@@ -281,19 +281,39 @@ function waitForInputBoxReady(
 }
 
 /**
- * Why: activation creates the tab synchronously but the PTY spawn is
- * async. Poll the store until the primary PTY id appears or the budget
- * expires. Tight interval because the wait is normally <200ms — only the
- * first launch on a cold app reaches the tail of this.
+ * Why: activation creates the tab synchronously but the PTY spawn is async.
+ * Subscribe to the tab→PTY binding instead of polling so paste-readiness
+ * sidecars attach in the same renderer turn that connectPanePty publishes the
+ * PTY id, before the transport registers its primary data handler.
  */
 async function waitForPtyId(tabId: string, timeoutMs: number): Promise<string | null> {
-  const deadline = Date.now() + timeoutMs
-  while (Date.now() < deadline) {
-    const ptyId = useAppStore.getState().ptyIdsByTabId[tabId]?.[0]
-    if (ptyId) {
-      return ptyId
-    }
-    await new Promise<void>((resolve) => window.setTimeout(resolve, 50))
+  const initialPtyId = useAppStore.getState().ptyIdsByTabId[tabId]?.[0]
+  if (initialPtyId) {
+    return initialPtyId
   }
-  return null
+
+  return new Promise<string | null>((resolve) => {
+    let settled = false
+    let unsubscribe: (() => void) | null = null
+    let timer: number | null = null
+    const finish = (ptyId: string | null): void => {
+      if (settled) {
+        return
+      }
+      settled = true
+      unsubscribe?.()
+      if (timer !== null) {
+        window.clearTimeout(timer)
+      }
+      resolve(ptyId)
+    }
+    timer = window.setTimeout(() => finish(null), timeoutMs)
+
+    unsubscribe = useAppStore.subscribe((state) => {
+      const ptyId = state.ptyIdsByTabId[tabId]?.[0]
+      if (ptyId) {
+        finish(ptyId)
+      }
+    })
+  })
 }