deepline 0.1.180 → 0.1.182

package/dist/bundling-sources/shared_libs/play-runtime/context.ts

@@ -41,6 +41,7 @@ import { InMemoryRateStateBackend } from './governor/in-memory-rate-state-backen
 import { pacingPolicyForTool } from './pacing';
 import {
   cloneToolExecuteResultWithExecution,
+  attachToolResultListDataset,
   createToolExecuteResult,
   parseToolExecuteResponse,
   isToolExecuteResult,
@@ -69,6 +70,12 @@ import {
   buildDurableToolCallAuthScopeDigest,
   buildDurableToolCallCacheKey,
 } from './durable-call-cache';
+import {
+  QUERY_RESULT_DATASET_PAGE_SIZE,
+  isCustomerDbDatasetTool,
+  isQueryResultDatasetReadRequest,
+  isQueryResultDatasetTool,
+} from './query-result-dataset';
 import {
   RuntimeReceiptWaitTimeoutError,
   executeWithDurableRuntimeReceipt,
@@ -266,7 +273,44 @@ function publicToolResponseEnvelope(value: unknown): {
 
 const EXECUTE_TOOL_METADATA_HEADER = 'x-deepline-include-tool-metadata';
 const EXECUTE_RESPONSE_CONTRACT_HEADER = 'x-deepline-execute-response-contract';
+const EXECUTE_RESPONSE_INTENT_HEADER = 'x-deepline-execute-response-intent';
 const V2_EXECUTE_RESPONSE_CONTRACT = 'v2-tool-response';
+function recordOrNull(value: unknown): Record<string, unknown> | null {
+  return value && typeof value === 'object' && !Array.isArray(value)
+    ? (value as Record<string, unknown>)
+    : null;
+}
+
+function rowsFromUnknown(value: unknown): Record<string, unknown>[] {
+  if (!Array.isArray(value)) return [];
+  return value.map((row) =>
+    row && typeof row === 'object' && !Array.isArray(row)
+      ? (row as Record<string, unknown>)
+      : { value: row },
+  );
+}
+
+function finiteNonNegativeInteger(value: unknown): number | null {
+  if (typeof value !== 'number' || !Number.isFinite(value) || value < 0) {
+    return null;
+  }
+  return Math.floor(value);
+}
+
+function finitePositiveInteger(value: unknown): number | null {
+  const integer = finiteNonNegativeInteger(value);
+  return integer !== null && integer > 0 ? integer : null;
+}
+
+type ToolExecutionApiOptions = {
+  timeoutMs?: number;
+  customerDbDataset?: {
+    limit: number;
+    offset: number;
+    pageSize: number;
+    totalRows: number;
+  };
+};
 const IN_MEMORY_STEP_RESULT_PREVIEW_LIMIT = 25;
 const BATCH_SIZE_LOG_SAMPLE_LIMIT = 10;
 const STEP_PROGRAM_MAP_DEFINITION = Symbol('deepline.stepProgramMapDefinition');
@@ -1695,24 +1739,146 @@ export class PlayContextImpl {
     metadata?: ToolResultMetadataInput | null;
     meta?: Record<string, unknown>;
     execution: ToolResultExecutionMetadata;
+    requestInput?: Record<string, unknown>;
   }): Promise<ToolExecuteResult> {
     if (isToolExecuteResult(input.result)) {
-      return cloneToolExecuteResultWithExecution(input.result, input.execution);
+      return this.attachCustomerDbDatasetResult(
+        input.toolId,
+        input.requestInput,
+        cloneToolExecuteResultWithExecution(input.result, input.execution),
+      );
     }
     const publicToolResult = publicToolResponseEnvelope(input.result);
-    return createToolExecuteResult({
-      status: publicToolResult?.status ?? input.status,
-      jobId: input.jobId,
-      result: publicToolResult
-        ? {
-            data: publicToolResult.raw,
-            ...(publicToolResult.meta ? { meta: publicToolResult.meta } : {}),
-          }
-        : input.result,
-      metadata:
-        input.metadata ?? (await this.resolveToolResultMetadata(input.toolId)),
-      execution: input.execution,
-      meta: input.meta,
+    return this.attachCustomerDbDatasetResult(
+      input.toolId,
+      input.requestInput,
+      createToolExecuteResult({
+        status: publicToolResult?.status ?? input.status,
+        jobId: input.jobId,
+        result: publicToolResult
+          ? {
+              data: publicToolResult.raw,
+              ...(publicToolResult.meta ? { meta: publicToolResult.meta } : {}),
+            }
+          : input.result,
+        metadata:
+          input.metadata ??
+          (await this.resolveToolResultMetadata(input.toolId)),
+        execution: input.execution,
+        meta: input.meta,
+      }),
+    );
+  }
+
+  private attachCustomerDbDatasetResult(
+    toolId: string,
+    requestInput: Record<string, unknown> | undefined,
+    wrapped: ToolExecuteResult,
+  ): ToolExecuteResult {
+    if (!isQueryResultDatasetTool(toolId)) return wrapped;
+    const raw = recordOrNull(wrapped.toolResponse.raw);
+    const dataset = recordOrNull(raw?.dataset);
+    const rows = rowsFromUnknown(raw?.rows);
+    const totalRows = finiteNonNegativeInteger(dataset?.total_rows);
+    const sql =
+      typeof requestInput?.sql === 'string'
+        ? requestInput.sql
+        : typeof requestInput?.query === 'string'
+          ? requestInput.query
+          : typeof raw?.sql === 'string'
+            ? raw.sql
+            : null;
+    if (!dataset || totalRows === null || !sql) {
+      return wrapped;
+    }
+    const originalRequestInput = requestInput as Record<string, unknown>;
+    const datasetLimit =
+      finitePositiveInteger(dataset.returned_limit) ?? totalRows;
+    const previewRows = rows.slice(0, Math.min(rows.length, 25));
+    const executionNonce =
+      typeof wrapped.job_id === 'string' && wrapped.job_id.trim()
+        ? wrapped.job_id.trim()
+        : `${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}`;
+    const datasetId = `tool-list:${sha256Hex(
+      `${toolId}:${this.currentRunId}:${sql}:${datasetLimit}:${executionNonce}`,
+    )}`;
+    const fetchPage = async (
+      offset: number,
+      limit: number,
+    ): Promise<Record<string, unknown>[]> => {
+      if (limit <= 0 || offset >= totalRows) return [];
+      const execution = await this.callToolExecutionAPI(
+        toolId,
+        originalRequestInput,
+        {
+          timeoutMs: resolveToolRuntimeTimeoutMs(toolId),
+          customerDbDataset: {
+            limit: datasetLimit,
+            offset,
+            pageSize: Math.min(limit, QUERY_RESULT_DATASET_PAGE_SIZE),
+            totalRows,
+          },
+        },
+      );
+      const pageRaw =
+        recordOrNull(execution.toolResponse)?.raw ??
+        recordOrNull(execution.result)?.data ??
+        execution.result;
+      return rowsFromUnknown(recordOrNull(pageRaw)?.rows);
+    };
+    const collectRows = async (
+      limit: number | undefined,
+    ): Promise<Record<string, unknown>[]> => {
+      const target = Math.min(limit ?? totalRows, totalRows, datasetLimit);
+      const collected: Record<string, unknown>[] = [];
+      for (
+        let offset = 0;
+        offset < target;
+        offset += QUERY_RESULT_DATASET_PAGE_SIZE
+      ) {
+        collected.push(
+          ...(await fetchPage(
+            offset,
+            Math.min(QUERY_RESULT_DATASET_PAGE_SIZE, target - offset),
+          )),
+        );
+      }
+      return collected.slice(0, target);
+    };
+    const playDataset = createDeferredPlayDataset({
+      datasetKind: 'csv',
+      datasetId,
+      count: Math.min(totalRows, datasetLimit),
+      previewRows,
+      sourceLabel: 'query result rows',
+      tableNamespace: null,
+      resolvers: {
+        count: async () => Math.min(totalRows, datasetLimit),
+        peek: async (limit) => collectRows(limit),
+        materialize: async (limit) => collectRows(limit),
+        iterate: () =>
+          ({
+            async *[Symbol.asyncIterator]() {
+              const count = Math.min(totalRows, datasetLimit);
+              for (
+                let offset = 0;
+                offset < count;
+                offset += QUERY_RESULT_DATASET_PAGE_SIZE
+              ) {
+                yield* await fetchPage(
+                  offset,
+                  Math.min(QUERY_RESULT_DATASET_PAGE_SIZE, count - offset),
+                );
+              }
+            },
+          }) as AsyncIterable<Record<string, unknown>>,
+      },
+    });
+    return attachToolResultListDataset(wrapped, {
+      name: 'rows',
+      path: 'toolResponse.raw.rows',
+      dataset: playDataset,
+      count: Math.min(totalRows, datasetLimit),
     });
   }
 
@@ -1738,6 +1904,7 @@ export class PlayContextImpl {
         cacheKey,
         receiptKey,
       }),
+      requestInput: request.input,
     });
     let completed: RuntimeStepReceipt | null | undefined = null;
     if (receiptKey) {
@@ -1791,6 +1958,7 @@ export class PlayContextImpl {
             cacheKey: entry.request.cacheKey,
             receiptKey: entry.request.receiptKey,
           }),
+          requestInput: entry.request.input,
         }),
       })),
     );
@@ -1851,6 +2019,7 @@ export class PlayContextImpl {
                 ? 'no_result'
                 : 'completed',
             result: this.runtimeReceiptOutput(completed),
+            requestInput: request.input,
             execution: toolExecutionMetadataForOutcome(
               completed.runId === this.currentRunId
                 ? {
@@ -1867,7 +2036,9 @@ export class PlayContextImpl {
             ),
           })
         : wrapped;
-    this.cacheToolResult(toolId, cacheKey, finalWrapped);
+    if (request.cacheable !== false) {
+      this.cacheToolResult(toolId, cacheKey, finalWrapped);
+    }
 
     const resolver = this.toolCallResolvers.get(request.callId);
     if (resolver) {
@@ -2211,8 +2382,7 @@ export class PlayContextImpl {
           playId: this.#options.playId,
           runId: this.#options.runId,
           staticPipeline: this.#options.staticPipeline,
-          forceRefresh:
-            this.#options.cachePolicy?.forceToolRefresh === true,
+          forceRefresh: this.#options.cachePolicy?.forceToolRefresh === true,
         },
       );
       resolvedTableNamespace = normalizeTableNamespace(
@@ -3486,11 +3656,13 @@ export class PlayContextImpl {
       toolId,
       requestInput: input,
     });
+    const cacheableToolResult = !isQueryResultDatasetReadRequest(toolId, input);
     const durableCacheKey = await this.durableToolCallCacheKey({
       toolId,
       requestInput: input,
       staleAfterSeconds: toolCachePolicy.staleAfterSeconds,
     });
+    const durableReceiptKey = cacheableToolResult ? durableCacheKey : null;
     const eventWaitHandler =
       (await this.#options.getIntegrationEventWaitHandler?.(toolId)) ?? null;
     const store = rowContext.getStore();
@@ -3518,15 +3690,18 @@ export class PlayContextImpl {
 
       if (!store) {
         const directCacheKey = durableCacheKey;
-        const cached = toolCachePolicy.force
+        const cached = !cacheableToolResult
           ? null
-          : this.getCachedToolResult(toolId, directCacheKey);
+          : toolCachePolicy.force
+            ? null
+            : this.getCachedToolResult(toolId, directCacheKey);
         if (cached?.done) {
           this.log(`Tool cache hit: ${toolId} exact payload`);
           return await this.wrapToolExecutionResult({
             toolId,
             status: cached.result == null ? 'no_result' : 'completed',
             result: cached.result,
+            requestInput: input,
             execution: toolExecutionMetadataForOutcome({
               kind: 'checkpoint',
               cacheKey: directCacheKey,
@@ -3548,19 +3723,22 @@ export class PlayContextImpl {
           result: execution.result,
           metadata: execution.metadata,
           meta: execution.meta,
+          requestInput: input,
           execution: toolExecutionMetadataForOutcome({
             kind: 'live',
             cacheKey: directCacheKey,
           }),
         });
-        this.checkpoint.completedToolBatches[toolId] = {
-          ...(this.checkpoint.completedToolBatches[toolId] ?? {}),
-          [directCacheKey]: {
-            done: true,
-            result: wrapped,
-          },
-        };
-        this.#options.onBatchComplete?.(this.checkpoint);
+        if (cacheableToolResult) {
+          this.checkpoint.completedToolBatches[toolId] = {
+            ...(this.checkpoint.completedToolBatches[toolId] ?? {}),
+            [directCacheKey]: {
+              done: true,
+              result: wrapped,
+            },
+          };
+          this.#options.onBatchComplete?.(this.checkpoint);
+        }
         return wrapped;
       }
 
@@ -3581,15 +3759,18 @@ export class PlayContextImpl {
       }
 
       const toolResultCacheKey = durableCacheKey;
-      const cached = toolCachePolicy.force
+      const cached = !cacheableToolResult
         ? null
-        : this.getCachedToolResult(toolId, toolResultCacheKey);
+        : toolCachePolicy.force
+          ? null
+          : this.getCachedToolResult(toolId, toolResultCacheKey);
       if (cached?.done) {
         this.log(`  Row ${rowId} ${toolId}: cache hit exact payload`);
         return await this.wrapToolExecutionResult({
           toolId,
           status: cached.result == null ? 'no_result' : 'completed',
           result: cached.result,
+          requestInput: input,
           execution: toolExecutionMetadataForOutcome({
             kind: 'checkpoint',
             cacheKey: toolResultCacheKey,
@@ -3624,7 +3805,8 @@ export class PlayContextImpl {
         this.toolCallQueue.push({
           callId,
           cacheKey: toolResultCacheKey,
-          receiptKey: durableCacheKey,
+          cacheable: cacheableToolResult,
+          receiptKey: durableReceiptKey,
           force: toolCachePolicy.force,
           rowId,
           fieldName,
@@ -3641,7 +3823,7 @@ export class PlayContextImpl {
       });
     };
 
-    if (store) {
+    if (store || !cacheableToolResult) {
       return await executeTool();
     }
 
@@ -3650,7 +3832,7 @@ export class PlayContextImpl {
       normalizedKey,
       this.currentRunId,
       {
-        receiptKey: durableCacheKey,
+        receiptKey: durableReceiptKey,
         semanticKey: toolRequestIdentity,
         force: toolCachePolicy.force,
         staleAfterSeconds: toolCachePolicy.staleAfterSeconds,
@@ -4347,17 +4529,21 @@ export class PlayContextImpl {
     await Promise.all(
       [...byTool.entries()].map(async ([toolId, requests]) => {
         this.log(`Executing tool batch ${toolId}: ${requests.length} calls`);
+        const liveStepResults = new Map<string, unknown>();
 
         const recordToolStep = (stepRequests: ToolCallRequest[]): void => {
           if (stepRequests.length === 0) return;
           const stepResults: PlayStepRowResult[] = stepRequests.map((req) => {
-            const cachedResult = this.getCachedToolResult(toolId, req.cacheKey);
+            const hasLiveResult = liveStepResults.has(req.callId);
+            const result = hasLiveResult
+              ? liveStepResults.get(req.callId)
+              : this.getCachedToolResult(toolId, req.cacheKey)?.result;
             return {
               rowId: req.rowId,
-              status: cachedResult?.result != null ? 'completed' : 'failed',
-              success: cachedResult?.result != null,
-              value: cachedResult?.result,
-              error: cachedResult?.result != null ? null : 'Tool call failed',
+              status: result != null ? 'completed' : 'failed',
+              success: result != null,
+              value: result,
+              error: result != null ? null : 'Tool call failed',
             };
           });
           const toolStep = {
@@ -4377,9 +4563,12 @@ export class PlayContextImpl {
         const pendingRequests: ToolCallRequest[] = [];
         const recoveredRequests: ToolCallRequest[] = [];
         for (const req of requests) {
-          const cached = req.force
-            ? undefined
-            : this.getCachedToolResult(toolId, req.cacheKey);
+          const cached =
+            req.cacheable === false
+              ? undefined
+              : req.force
+                ? undefined
+                : this.getCachedToolResult(toolId, req.cacheKey);
           if (cached?.done) {
             this.log(`  Row ${req.rowId} ${toolId}: recovered from checkpoint`);
             const resolver = this.toolCallResolvers.get(req.callId);
@@ -4411,7 +4600,9 @@ export class PlayContextImpl {
                   attachedToReceiptKey: followerReceiptKey,
                 })
               : result;
-            this.cacheToolResult(toolId, follower.cacheKey, followerResult);
+            if (follower.cacheable !== false) {
+              this.cacheToolResult(toolId, follower.cacheKey, followerResult);
+            }
             const resolver = this.toolCallResolvers.get(follower.callId);
             if (resolver) {
               resolver.resolve(followerResult);
@@ -4558,6 +4749,7 @@ export class PlayContextImpl {
                   ? 'no_result'
                   : 'completed',
               result: this.runtimeReceiptOutput(receipt),
+              requestInput: request.input,
               execution: toolExecutionMetadataForOutcome({
                 kind: source,
                 cacheKey: request.cacheKey,
@@ -4689,6 +4881,7 @@ export class PlayContextImpl {
                   execution?.jobId,
                   execution?.meta,
                 );
+                liveStepResults.set(owner.callId, result);
                 resolveLiveFollowers(owner, result);
                 recordToolStep([owner]);
                 this.#options.onBatchComplete?.(this.checkpoint);
@@ -4837,6 +5030,7 @@ export class PlayContextImpl {
                     index += 1
                   ) {
                     const request = entry.request.memberRequests[index]!;
+                    liveStepResults.set(request.callId, resolvedResults[index]);
                     resolveLiveFollowers(request, resolvedResults[index]);
                   }
                 }
@@ -4894,6 +5088,7 @@ export class PlayContextImpl {
                       execution?.jobId,
                       execution?.meta,
                     );
+                    liveStepResults.set(request.callId, result);
                     resolveLiveFollowers(request, result);
                   } catch (error) {
                     await rejectWithLiveFollowers(request, error);
@@ -4989,7 +5184,7 @@ export class PlayContextImpl {
   private async callToolAPI(
     toolId: string,
     input: Record<string, unknown>,
-    options?: { timeoutMs?: number },
+    options?: ToolExecutionApiOptions,
   ): Promise<unknown> {
     const execution = await this.callToolExecutionAPI(toolId, input, options);
     if (execution.toolResponse && 'raw' in execution.toolResponse) {
@@ -5006,7 +5201,7 @@ export class PlayContextImpl {
   private async callToolExecutionAPI(
     toolId: string,
     input: Record<string, unknown>,
-    options?: { timeoutMs?: number },
+    options?: ToolExecutionApiOptions,
   ): Promise<ParsedToolExecuteResponse> {
     if (!this.#options.executorToken || !this.#options.baseUrl) {
       throw new Error(
@@ -5064,6 +5259,7 @@ export class PlayContextImpl {
                   Authorization: `Bearer ${this.#options.executorToken}`,
                   [EXECUTE_RESPONSE_CONTRACT_HEADER]:
                     V2_EXECUTE_RESPONSE_CONTRACT,
+                  [EXECUTE_RESPONSE_INTENT_HEADER]: 'dataset',
                   [EXECUTE_TOOL_METADATA_HEADER]: 'true',
                   ...(this.#options.vercelProtectionBypassToken?.trim()
                     ? {
@@ -5074,7 +5270,30 @@ export class PlayContextImpl {
                 },
                 body: JSON.stringify({
                   payload: input,
-                  metadata: { parent_run_id: this.#options.runId },
+                  metadata: {
+                    parent_run_id: this.#options.runId,
+                    ...(options?.customerDbDataset
+                      ? {
+                          query_result_dataset: {
+                            limit: options.customerDbDataset.limit,
+                            offset: options.customerDbDataset.offset,
+                            page_size: options.customerDbDataset.pageSize,
+                            total_rows: options.customerDbDataset.totalRows,
+                          },
+                          ...(isCustomerDbDatasetTool(toolId)
+                            ? {
+                                customer_db_dataset: {
+                                  limit: options.customerDbDataset.limit,
+                                  offset: options.customerDbDataset.offset,
+                                  page_size: options.customerDbDataset.pageSize,
+                                  total_rows:
+                                    options.customerDbDataset.totalRows,
+                                },
+                              }
+                            : {}),
+                        }
+                      : {}),
+                  },
                   ...(this.#options.integrationMode
                     ? { integration_mode: this.#options.integrationMode }
                     : {}),

package/dist/bundling-sources/shared_libs/play-runtime/ctx-types.ts

@@ -22,6 +22,7 @@ export interface RowState {
 export interface ToolCallRequest {
   callId: string;
   cacheKey: string;
+  cacheable?: boolean;
   receiptKey?: string | null;
   force?: boolean;
   rowId: number;

package/dist/bundling-sources/shared_libs/play-runtime/query-result-dataset.ts

@@ -0,0 +1,120 @@
+export type QueryResultDatasetRequest = {
+  limit?: number;
+  offset?: number;
+  pageSize?: number;
+  totalRows?: number;
+};
+
+export const QUERY_RESULT_DATASET_PAGE_SIZE = 1_000;
+
+function maskSqlForDatasetClassifier(sql: string): string {
+  let output = '';
+  let index = 0;
+  while (index < sql.length) {
+    const char = sql[index];
+    const next = sql[index + 1];
+    if (char === '-' && next === '-') {
+      const newline = sql.indexOf('\n', index + 2);
+      output += ' ';
+      index = newline === -1 ? sql.length : newline + 1;
+      continue;
+    }
+    if (char === '/' && next === '*') {
+      index += 2;
+      let depth = 1;
+      while (index < sql.length && depth > 0) {
+        if (sql[index] === '/' && sql[index + 1] === '*') {
+          depth += 1;
+          index += 2;
+          continue;
+        }
+        if (sql[index] === '*' && sql[index + 1] === '/') {
+          depth -= 1;
+          index += 2;
+          continue;
+        }
+        index += 1;
+      }
+      output += ' ';
+      continue;
+    }
+    if (char === "'" || char === '"') {
+      const quote = char;
+      output += ' ';
+      index += 1;
+      while (index < sql.length) {
+        const quoted = sql[index];
+        if (quoted === quote && sql[index + 1] === quote) {
+          index += 2;
+          continue;
+        }
+        index += 1;
+        if (quoted === quote) break;
+      }
+      continue;
+    }
+    output += char;
+    index += 1;
+  }
+  return output;
+}
+
+export function isCustomerDbDatasetTool(toolId: string): boolean {
+  return /^(?:customer_db_)?query_customer_db$/.test(toolId);
+}
+
+export function isSnowflakeQueryDatasetTool(toolId: string): boolean {
+  return /^snowflake_(?:snowflake_)?run_(?:semantic_)?query$/.test(toolId);
+}
+
+function isSnowflakeSemanticQueryDatasetTool(toolId: string): boolean {
+  return /^snowflake_(?:snowflake_)?run_semantic_query$/.test(toolId);
+}
+
+export function isQueryResultDatasetTool(toolId: string): boolean {
+  return isCustomerDbDatasetTool(toolId) || isSnowflakeQueryDatasetTool(toolId);
+}
+
+export function isQueryResultDatasetOperation(input: {
+  provider?: string | null;
+  operation?: string | null;
+}): boolean {
+  return (
+    (input.provider === 'customer_db' &&
+      input.operation === 'query_customer_db') ||
+    (input.provider === 'snowflake' &&
+      (input.operation === 'snowflake_run_query' ||
+        input.operation === 'snowflake_run_semantic_query'))
+  );
+}
+
+function isReadOnlySelectOrWith(sql: string): boolean {
+  const normalized = maskSqlForDatasetClassifier(sql).trimStart().toLowerCase();
+  return (
+    (normalized.startsWith('select') &&
+      !/\bselect\b[\s\S]*?\binto\b/i.test(normalized)) ||
+    (normalized.startsWith('with') &&
+      !/\b(insert\s+into|update|delete\s+from|merge\s+into|select\b[\s\S]*?\binto|create|alter|drop|truncate)\b/i.test(
+        normalized,
+      ))
+  );
+}
+
+export function isQueryResultDatasetReadRequest(
+  toolId: string,
+  input: Record<string, unknown>,
+): boolean {
+  if (!isQueryResultDatasetTool(toolId)) return false;
+  if (isSnowflakeQueryDatasetTool(toolId) && input.write === true) return false;
+  if (isSnowflakeSemanticQueryDatasetTool(toolId)) return true;
+  const sql =
+    typeof input.sql === 'string'
+      ? input.sql
+      : typeof input.query === 'string'
+        ? input.query
+        : '';
+  return isReadOnlySelectOrWith(sql);
+}
+
+export const isCustomerDbDatasetReadRequest =
+  isQueryResultDatasetReadRequest;