perf(common,sdk,bigquery): optimize JSON parsing, fix race condition, add query time bounds

- Optimize partial JSON delta parsing from O(n²) to O(n)
- Fix race condition in code-search Promise.allSettled logic
- Add time bounds to BigQuery queries for partition pruning

Author: brandonkachen

common/src/util/__tests__/partial-json-delta.test.ts

@@ -108,6 +108,38 @@ describe('parsePartialJsonObjectSingle', () => {
     })
   })
 
+  describe('comma search optimization', () => {
+    it('should efficiently find last valid comma in deeply nested incomplete JSON', () => {
+      // This tests the O(n) backward comma search optimization
+      const input = '{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5, "incomplete":'
+      const result = parsePartialJsonObjectSingle(input)
+      expect(result).toEqual({
+        lastParamComplete: true,
+        params: { a: 1, b: 2, c: 3, d: 4, e: 5 },
+      })
+    })
+
+    it('should handle comma inside string value when searching backwards', () => {
+      // Comma inside a string should not be treated as a separator
+      const input = '{"message": "Hello, world", "incomplete":'
+      const result = parsePartialJsonObjectSingle(input)
+      expect(result).toEqual({
+        lastParamComplete: true,
+        params: { message: 'Hello, world' },
+      })
+    })
+
+    it('should find valid comma after skipping invalid parse attempts', () => {
+      // Multiple commas, need to find the right one
+      const input = '{"x": [1, 2, 3], "y": {"a": 1, "b": 2}, "z":'
+      const result = parsePartialJsonObjectSingle(input)
+      expect(result).toEqual({
+        lastParamComplete: true,
+        params: { x: [1, 2, 3], y: { a: 1, b: 2 } },
+      })
+    })
+  })
+
   describe('edge cases', () => {
     it('should return empty object for empty string', () => {
       const input = ''

common/src/util/partial-json-delta.ts

@@ -1,4 +1,3 @@
-// TODO: optimize this to not be O(n^2)
 export function parsePartialJsonObjectSingle(content: string): {
   lastParamComplete: boolean
   params: any
@@ -26,16 +25,14 @@ export function parsePartialJsonObjectSingle(content: string): {
     } catch {}
   }
 
-  let lastIndex = content.lastIndexOf(',')
-  while (lastIndex > 0) {
+  let commaPos = content.length
+  while ((commaPos = content.lastIndexOf(',', commaPos - 1)) !== -1) {
     try {
       return {
         lastParamComplete: true,
-        params: JSON.parse(content.slice(0, lastIndex) + '}'),
+        params: JSON.parse(content.slice(0, commaPos) + '}'),
       }
     } catch {}
-
-    lastIndex = content.lastIndexOf(',', lastIndex - 1)
   }
 
   return { lastParamComplete: true, params: {} }

packages/bigquery/src/client.ts

@@ -246,27 +246,39 @@ export async function getTracesWithoutRelabels(
   userId: string | undefined = undefined,
   dataset: string = DATASET,
 ) {
-  // TODO: Optimize query, maybe only get traces in last 30 days etc
+  const thirtyDaysAgo = new Date(Date.now() - 30 * 24 * 60 * 60 * 1000)
+    .toISOString()
+    .split('T')[0]
+
   const query = `
     SELECT t.*
     FROM \`${dataset}.${TRACES_TABLE}\` t
     LEFT JOIN (
       SELECT r.agent_step_id, r.user_id, JSON_EXTRACT_SCALAR(r.payload, '$.user_input_id') as user_input_id
       FROM \`${dataset}.${RELABELS_TABLE}\` r
-      WHERE r.model = '${model}'
-      ${userId ? `AND r.user_id = '${userId}'` : ''}
+      WHERE r.model = @model
+      ${userId ? `AND r.user_id = @userId` : ''}
     ) r
     ON t.agent_step_id = r.agent_step_id
        AND t.user_id = r.user_id
        AND JSON_EXTRACT_SCALAR(t.payload, '$.user_input_id') = r.user_input_id
     WHERE t.type = 'get-relevant-files'
+      AND t.created_at >= @thirtyDaysAgo
       AND r.agent_step_id IS NULL
-      ${userId ? `AND t.user_id = '${userId}'` : ''}
+      ${userId ? `AND t.user_id = @userId` : ''}
     ORDER BY t.created_at DESC
-    LIMIT ${limit}
+    LIMIT @limit
   `
 
-  const [rows] = await getClient().query(query)
+  const [rows] = await getClient().query({
+    query,
+    params: {
+      model,
+      thirtyDaysAgo,
+      limit,
+      ...(userId ? { userId } : {}),
+    },
+  })
   // Parse the payload as JSON if it's a string
   return rows.map((row) => ({
     ...row,

sdk/src/tools/code-search.ts

@@ -104,28 +104,42 @@ export function codeSearch({
     let estimatedOutputLen = 0
     let killedForLimit = false
 
+    // Guard to prevent double-settlement from concurrent timeout and process close events
+    let killTimeoutId: ReturnType<typeof setTimeout> | null = null
+
     const settle = (payload: any) => {
       if (isResolved) return
       isResolved = true
 
-      // Clean up listeners immediately
+      // Clean up listeners immediately to prevent further events
       childProcess.stdout.removeAllListeners()
       childProcess.stderr.removeAllListeners()
       childProcess.removeAllListeners()
 
+      // Clear both the main timeout and the kill timeout to prevent late callbacks
       clearTimeout(timeoutId)
+      if (killTimeoutId) {
+        clearTimeout(killTimeoutId)
+        killTimeoutId = null
+      }
+
       resolve([{ type: 'json', value: payload }])
     }
 
     const hardKill = () => {
       try {
         childProcess.kill('SIGTERM')
       } catch {}
-      setTimeout(() => {
+      // Store timeout reference so it can be cleared if process closes normally
+      killTimeoutId = setTimeout(() => {
         try {
-          // SIGKILL doesn't exist on Windows, fall back to no-signal kill
-          childProcess.kill('SIGKILL') || childProcess.kill()
-        } catch {}
+          childProcess.kill('SIGKILL')
+        } catch {
+          try {
+            childProcess.kill()
+          } catch {}
+        }
+        killTimeoutId = null
       }, 1000)
     }