fix(snowflake): cancel Cortex Analyst's generated SQL with the run

The generated SQL used to run as a nested tool call, which inherited the run's
abort signal; running it directly dropped that. postProcess now receives the
executor's effective abort signal as an optional fourth argument (the same
signal the tool's own request used), and Cortex Analyst forwards it to the DNS
check, the guarded fetch, and the response read. Other tools' postProcess
hooks are unchanged.

Author: waleedlatif1

apps/sim/tools/index.test.ts

@@ -3737,6 +3737,51 @@ describe('Internal Route Trust', () => {
       }
     )
 
+    it('aborts Cortex Analyst’s generated SQL when the run is cancelled', async () => {
+      ;(tools as Record<string, unknown>)[snowflakeCortexAnalystAskTool.id] =
+        snowflakeCortexAnalystAskTool
+      const controller = new AbortController()
+      let statementSignal: AbortSignal | undefined
+      try {
+        mockResolveExecutorCredentialToken.mockResolvedValue({
+          credentialType: 'service_account',
+          domain: CREDENTIAL_HOST,
+          accessToken: 'pat-secret',
+        })
+        mockSecureFetchWithPinnedIP
+          .mockResolvedValueOnce(
+            toSecureFetchResponse(
+              Response.json({
+                message: { role: 'analyst', content: [{ type: 'sql', statement: 'SELECT 1' }] },
+              })
+            )
+          )
+          .mockImplementationOnce(
+            async (_url: string, _ip: string, init: { signal?: AbortSignal }) => {
+              statementSignal = init.signal
+              controller.abort(new Error('run cancelled'))
+              throw new Error('aborted')
+            }
+          )
+
+        const result = await executeTool(
+          snowflakeCortexAnalystAskTool.id,
+          {
+            oauthCredential: 'snowflake-credential',
+            question: 'How many?',
+            semanticView: 'DB.SCHEMA.VIEW',
+            executeSql: true,
+          },
+          { signal: controller.signal }
+        )
+
+        expect(statementSignal?.aborted).toBe(true)
+        expect(result.success).toBe(false)
+      } finally {
+        Reflect.deleteProperty(tools, snowflakeCortexAnalystAskTool.id)
+      }
+    })
+
     it('sends a keyless credential’s token through workload identity federation', async () => {
       const { headers } = await snowflakeRequestHeaders(
         { accessToken: 'signed-jwt', authStyle: 'snowflake-workload-identity' },
@@ -4520,7 +4565,8 @@ describe('Internal Route Trust', () => {
       expect(postProcess).toHaveBeenCalledWith(
         expect.any(Object),
         expect.objectContaining({ prompt: 'direct-secret', apiKey: 'direct-secret' }),
-        expect.any(Function)
+        expect.any(Function),
+        expect.any(Object)
       )
       expect(params).toEqual(originalParams)
     } finally {

apps/sim/tools/index.ts

@@ -2369,7 +2369,9 @@ async function executeToolImplementation(
       let finalResult = result
       if (tool.postProcess && result.success && !skipPostProcess) {
         try {
-          finalResult = await tool.postProcess(result, contextParams, executeNestedTool)
+          finalResult = await tool.postProcess(result, contextParams, executeNestedTool, {
+            signal: effectiveSignal,
+          })
         } catch (error) {
           logger.error(`Post-processing error for ${toolId}`, {
             requestId,
@@ -2462,7 +2464,9 @@ async function executeToolImplementation(
     let finalResult = result
     if (tool.postProcess && result.success && !skipPostProcess) {
       try {
-        finalResult = await tool.postProcess(result, contextParams, executeNestedTool)
+        finalResult = await tool.postProcess(result, contextParams, executeNestedTool, {
+          signal: effectiveSignal,
+        })
       } catch (error) {
         const normalizedError = toError(error)
         logger.error(

apps/sim/tools/snowflake/cortex_analyst_ask.ts

@@ -143,7 +143,7 @@ export const cortexAnalystAskTool: ToolConfig<
       ),
     }
   },
-  postProcess: async (result, params) => {
+  postProcess: async (result, params, _executeTool, context) => {
     const { sql } = result.output
     if (!sql || !shouldExecuteSql(params.executeSql)) return result
 
@@ -159,7 +159,8 @@ export const cortexAnalystAskTool: ToolConfig<
       const execution = await runSnowflakeStatement(
         params,
         sql,
-        cortexAnalystSqlContext(params, result.output.semanticModelSelection?.index)
+        cortexAnalystSqlContext(params, result.output.semanticModelSelection?.index),
+        context?.signal
       )
       return { ...result, output: { ...result.output, execution } }
     } catch (error) {

apps/sim/tools/snowflake/run-statement.server.ts

@@ -22,15 +22,16 @@ import {
 export async function runSnowflakeStatement(
   params: SnowflakeResultParams,
   statement: string,
-  context: { database?: string; schema?: string }
+  context: { database?: string; schema?: string },
+  signal?: AbortSignal
 ): Promise<SnowflakeStatementOutput> {
   const url = `${getSnowflakeBaseUrl(params)}/api/v2/statements`
   const body = buildSnowflakeStatementBody(
     params,
     { statement },
     { context, warehouse: params.warehouse, maxRows: params.maxRows }
   )
-  const validation = await validateUrlWithDNS(url, 'snowflakeUrl', 'requestTarget')
+  const validation = await validateUrlWithDNS(url, 'snowflakeUrl', 'requestTarget', { signal })
   if (!validation.isValid) {
     throw new Error(`Invalid Snowflake URL: ${validation.error ?? 'DNS resolution failed'}`)
   }
@@ -41,12 +42,14 @@ export async function runSnowflakeStatement(
     body: JSON.stringify(body),
     maxRedirects: 0,
     maxResponseBytes: SNOWFLAKE_MAX_RESPONSE_BYTES,
+    signal,
   })
   return readSnowflakeResult(
     new Response(response.body, {
       status: response.status,
       statusText: response.statusText,
       headers: response.headers.toRecord(),
-    })
+    }),
+    { signal }
   )
 }

apps/sim/tools/types.ts

@@ -331,7 +331,9 @@ export interface ToolConfig<P = any, R = any> {
   postProcess?: (
     result: R extends ToolResponse ? R : ToolResponse,
     params: P,
-    executeTool: (toolId: string, params: Record<string, any>) => Promise<ToolResponse>
+    executeTool: (toolId: string, params: Record<string, any>) => Promise<ToolResponse>,
+    /** The run's abort signal, the same one the tool's own request used. */
+    context?: { signal?: AbortSignal }
   ) => Promise<R extends ToolResponse ? R : ToolResponse>
 
   // Response handling