From 5b1c903f78d215f2541033e9e68537c8979acfc8 Mon Sep 17 00:00:00 2001
From: dal <dallin@buster.so>
Date: Tue, 26 Aug 2025 10:34:15 -0600
Subject: [PATCH] snowflake credentials adjustment

---
 .../communication-tools/done-tool/done-tool-start.ts | 12 +++++-------
 packages/data-source/src/adapters/snowflake.ts       | 11 ++++++++++-
 packages/data-source/src/types/credentials.ts        |  3 +++
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts
index 73bd3afc5..7d2a86d44 100644
--- a/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts
+++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts
@@ -102,14 +102,12 @@ export function createDoneToolStart(context: DoneToolContext, doneToolState: Don
 
     const doneToolResponseEntry = createDoneToolResponseMessage(doneToolState, options.toolCallId);
     const doneToolMessage = createDoneToolRawLlmMessageEntry(doneToolState, options.toolCallId);
-    
+
     // Create the tool result immediately with success: true
     // This ensures it's always present even if the stream terminates early
-    const rawToolResultEntry = createRawToolResultEntry(
-      options.toolCallId,
-      DONE_TOOL_NAME,
-      { success: true }
-    );
+    const rawToolResultEntry = createRawToolResultEntry(options.toolCallId, DONE_TOOL_NAME, {
+      success: true,
+    });
 
     const entries: UpdateMessageEntriesParams = {
       messageId: context.messageId,
@@ -126,7 +124,7 @@ export function createDoneToolStart(context: DoneToolContext, doneToolState: Don
       rawLlmMessages.push(doneToolMessage);
     }
     rawLlmMessages.push(rawToolResultEntry);
-    
+
     if (rawLlmMessages.length > 0) {
       entries.rawLlmMessages = rawLlmMessages;
     }
diff --git a/packages/data-source/src/adapters/snowflake.ts b/packages/data-source/src/adapters/snowflake.ts
index b4e022021..6777b2f9e 100644
--- a/packages/data-source/src/adapters/snowflake.ts
+++ b/packages/data-source/src/adapters/snowflake.ts
@@ -93,13 +93,22 @@ export class SnowflakeAdapter extends BaseAdapter {
 
   private async createConnection(credentials: SnowflakeCredentials): Promise<snowflake.Connection> {
     const connectionOptions: snowflake.ConnectionOptions = {
-      account: credentials.account_id,
+      account: credentials.account_id, // Always required by SDK
       username: credentials.username,
       password: credentials.password,
       warehouse: credentials.warehouse_id,
       database: credentials.default_database,
     };
 
+    // Use custom_host if provided via accessUrl
+    if (credentials.custom_host) {
+      // Ensure the URL has proper protocol
+      const host = credentials.custom_host.startsWith('http')
+        ? credentials.custom_host
+        : `https://${credentials.custom_host}`;
+      connectionOptions.accessUrl = host;
+    }
+
     if (credentials.role) {
       connectionOptions.role = credentials.role;
     }
diff --git a/packages/data-source/src/types/credentials.ts b/packages/data-source/src/types/credentials.ts
index 10c487157..ee1958ebd 100644
--- a/packages/data-source/src/types/credentials.ts
+++ b/packages/data-source/src/types/credentials.ts
@@ -38,6 +38,9 @@ export interface SnowflakeCredentials {
 
   /** Default schema to use */
   default_schema?: string;
+
+  /** Optional custom host for load balancers (e.g., "ridedvp-angel.yukicomputing.com:443") */
+  custom_host?: string;
 }
 
 /**