[NET-1186] feat: Add file-handling support in simulation workflow (#90)

Author: akash-vijay-kv

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "netra-sdk",
-  "version": "1.1.0",
+  "version": "1.2.0",
   "description": "A comprehensive TypeScript/JavaScript SDK for AI application observability built on top of OpenTelemetry and Traceloop",
   "type": "module",
   "main": "./dist/index.cjs",

package.json

@@ -111,6 +111,7 @@ export type {
   ConversationResponse,
   ConversationResult,
   CreateRunResult,
+  ProcessedFile,
   SimulationItem,
   SimulationOptions,
   SimulationResult,

src/index.ts

@@ -9,6 +9,7 @@ import { SpanWrapper } from "../span-wrapper";
 import { SimulationHttpClient } from "./client";
 import {
     ConversationResult,
+    FileData,
     SimulationItem,
     SimulationResult,
 } from "./models";
@@ -197,10 +198,11 @@ export class Simulation {
         runItem: SimulationItem,
         task: BaseTask,
     ): Promise<ConversationResult> {
-        const { runItemId, message: initialMessage, turnId: initialTurnId } = runItem;
+        const { runItemId, message: initialMessage, turnId: initialTurnId, files: initialFiles } = runItem;
         let message = initialMessage;
         let turnId = initialTurnId;
         let sessionId: string | null = null;
+        let rawFiles: FileData[] = initialFiles ?? [];
 
         while (true) {
             const span = new SpanWrapper(SPAN_NAME, {}, LOG_PREFIX);
@@ -209,11 +211,8 @@ export class Simulation {
             try {
                 const traceId = span.getCurrentSpan()?.spanContext().traceId ?? "";
 
-                // Run the user's task inside the active span context so child
-                // spans (LLM calls, HTTP, etc.) are parented correctly and
-                // outgoing HTTP headers carry the W3C trace context.
                 const [responseMessage, taskSessionId] = await span.withActive(
-                    () => executeTask(task, message, sessionId),
+                    () => executeTask(task, message, sessionId, rawFiles),
                 );
 
                 if (taskSessionId) {
@@ -252,9 +251,9 @@ export class Simulation {
                     };
                 }
 
-                // Continue to next turn
                 message = response.nextUserMessage!;
                 turnId = response.nextTurnId!;
+                rawFiles = response.nextFiles || [];
             } catch (error) {
                 const errorMsg = error instanceof Error ? error.message : String(error);
                 Logger.error(

src/simulation/api.ts

@@ -3,6 +3,7 @@ import { Config } from "../config";
 import { Logger } from "../logger";
 import { injectTraceContextHeaders } from "../utils/context-propagation";
 import { ConversationResponse, SimulationItem } from "./models";
+import { parseFiles } from "./utils";
 
 const LOG_PREFIX = "netra.simulation";
 const DEFAULT_TIMEOUT = 10000; // 10 seconds in milliseconds
@@ -144,6 +145,7 @@ export class SimulationHttpClient {
                     runItemId: msg.testRunItemId || "",
                     message: msg.userMessage || "",
                     turnId: msg.turnId || "",
+                    files: parseFiles(msg.attachments),
                 }),
             );
 
@@ -206,6 +208,7 @@ export class SimulationHttpClient {
                 nextTurnId: nextMsg.turnId || "",
                 nextUserMessage: nextMsg.userMessage || "",
                 nextRunItemId: nextMsg.testRunItemId || "",
+                nextFiles: parseFiles(nextMsg.attachments),
             };
         } catch (error) {
             const errorMsg = this._extractErrorMessage(error);

src/simulation/client.ts

@@ -9,6 +9,7 @@ export {
     ConversationStatus,
     type ConversationResponse,
     type ConversationResult,
+    type ProcessedFile,
     type SimulationItem,
     type SimulationResult,
     type TaskResult,

src/simulation/index.ts

@@ -6,13 +6,42 @@ export enum ConversationStatus {
   STOP = "stop",
 }
 
+/**
+ * Raw file metadata received from the backend.
+ *
+ * Contains a pre-signed download URL that the SDK uses to fetch the actual
+ * file content. Instances are produced by parsing the `attachments` array
+ * returned on each user message from the simulation API.
+ */
+export interface FileData {
+  fileName: string;
+  contentType: string;
+  description?: string;
+  downloadUrl: string;
+}
+
+/**
+ * File after download and base64 encoding, delivered to the user task.
+ *
+ * The `data` field contains the raw file bytes encoded as a base64 ASCII
+ * string.  Consumers should decode with `Buffer.from(data, "base64")` before
+ * passing to an LLM or other downstream system.
+ */
+export interface ProcessedFile {
+  fileName: string;
+  contentType: string;
+  description?: string;
+  data: string;
+}
+
 /**
  * Represents a single item in a simulation run.
  */
 export interface SimulationItem {
   runItemId: string;
   message: string;
   turnId: string;
+  files?: FileData[];
 }
 
 /**
@@ -24,6 +53,7 @@ export interface ConversationResponse {
   nextTurnId?: string;
   nextUserMessage?: string;
   nextRunItemId?: string;
+  nextFiles?: FileData[];
 }
 
 /**

src/simulation/models.ts

@@ -1,16 +1,21 @@
-import { TaskResult } from "./models";
+import { ProcessedFile, TaskResult } from "./models";
 
 
 export abstract class BaseTask {
     /**
+     * Process a simulation turn and return the agent's response.
+     *
      * @param message - The input message from the simulation.
-     * @param sessionId - The Session identifier.
+     * @param sessionId - The session identifier.
+     * @param files - Optional list of base64-encoded file attachments from the
+     *                dataset item.  Will be `null` when no files are attached.
      * @returns The task result containing:
      *            - message (string): The response message from the task.
      *            - sessionId (string): The session identifier.
      */
     abstract run(
         message: string,
         sessionId?: string | null,
+        files?: ProcessedFile[] | null,
     ): Promise<TaskResult> | TaskResult;
 }

src/simulation/task.ts

@@ -1,7 +1,16 @@
+import axios from "axios";
+import pLimit from "p-limit";
 import { Logger } from "../logger";
+import { FileData, ProcessedFile } from "./models";
 import { BaseTask } from "./task";
 
+type LimitFunction = ReturnType<typeof pLimit>;
+
 const LOG_PREFIX = "netra.simulation";
+const DEFAULT_FILE_DOWNLOAD_TIMEOUT_S = 30;
+const MAX_FILE_DOWNLOAD_WORKERS = 8;
+
+let _cachedFileDownloadTimeoutMs: number | null = null;
 
 /**
  * Format the trace ID as a 32-digit hexadecimal string.
@@ -35,26 +44,188 @@ export function validateSimulationInputs(
     return true;
 }
 
+/**
+ * Get the file download timeout in milliseconds.
+ *
+ * Reads from the `NETRA_SIMULATION_FILE_DOWNLOAD_TIMEOUT` environment
+ * variable (value in seconds) on first access and caches the result.
+ * Falls back to the default of 30 s.
+ *
+ * @returns Timeout in milliseconds
+ */
+function _getFileDownloadTimeout(): number {
+    if (_cachedFileDownloadTimeoutMs !== null) {
+        return _cachedFileDownloadTimeoutMs;
+    }
+
+    const envVal = process.env.NETRA_SIMULATION_FILE_DOWNLOAD_TIMEOUT;
+    if (envVal) {
+        const parsed = parseFloat(envVal);
+        if (!isNaN(parsed) && parsed > 0) {
+            _cachedFileDownloadTimeoutMs = parsed * 1000;
+            return _cachedFileDownloadTimeoutMs;
+        }
+        Logger.warn(
+            `${LOG_PREFIX}: Invalid file download timeout '${envVal}', using default ${DEFAULT_FILE_DOWNLOAD_TIMEOUT_S}s`,
+        );
+    }
+    _cachedFileDownloadTimeoutMs = DEFAULT_FILE_DOWNLOAD_TIMEOUT_S * 1000;
+    return _cachedFileDownloadTimeoutMs;
+}
+
+/**
+ * Download a single file from its pre-signed URL and base64-encode the content.
+ *
+ * @param fileData - Metadata with the download URL
+ * @param timeoutMs - Request timeout in milliseconds
+ * @param signal - AbortSignal to cancel the download if another file in the
+ *                 batch fails
+ * @returns ProcessedFile with base64-encoded data
+ * @throws Error if the download fails, with contextual information about which
+ *         file failed and the HTTP status (if available)
+ */
+async function _downloadSingleFile(
+    fileData: FileData,
+    timeoutMs: number,
+    signal?: AbortSignal,
+): Promise<ProcessedFile> {
+    try {
+        const response = await axios.get(fileData.downloadUrl, {
+            responseType: "arraybuffer",
+            timeout: timeoutMs,
+            signal,
+        });
+
+        const encoded = Buffer.from(response.data).toString("base64");
+        return {
+            fileName: fileData.fileName,
+            contentType: fileData.contentType,
+            description: fileData.description,
+            data: encoded,
+        };
+    } catch (error) {
+        if (axios.isCancel(error)) {
+            throw new Error(
+                `Download of '${fileData.fileName}' was cancelled`,
+            );
+        }
+
+        const status = axios.isAxiosError(error) ? error.response?.status : undefined;
+        const reason = status
+            ? `HTTP ${status}`
+            : (error instanceof Error ? error.message : String(error));
+
+        Logger.error(
+            `${LOG_PREFIX}: Failed to download file '${fileData.fileName}': ${reason}`,
+        );
+
+        throw new Error(
+            `Failed to download file '${fileData.fileName}': ${reason}`,
+        );
+    }
+}
+
+/**
+ * Parse raw attachment entries from the backend into typed FileData objects.
+ *
+ * Entries missing a `fileName` or `downloadUrl` are skipped with a warning.
+ *
+ * @param rawFiles - Array of attachment objects from the API response
+ * @returns Parsed FileData array (empty when input is null/undefined)
+ */
+export function parseFiles(
+    rawFiles: Array<Record<string, string>> | null | undefined,
+): FileData[] {
+    if (!rawFiles) {
+        return [];
+    }
+
+    const parsed: FileData[] = [];
+    for (const entry of rawFiles) {
+        const fileName = entry.fileName || "";
+        const downloadUrl = entry.downloadUrl || "";
+
+        if (!fileName || !downloadUrl) {
+            Logger.warn(
+                `${LOG_PREFIX}: Skipping malformed file attachment (missing fileName or downloadUrl)`,
+            );
+            continue;
+        }
+
+        parsed.push({
+            fileName,
+            contentType: entry.contentType || "",
+            description: entry.description || undefined,
+            downloadUrl,
+        });
+    }
+
+    return parsed;
+}
+
+/**
+ * Download and base64-encode a batch of files concurrently.
+ *
+ * Uses up to {@link MAX_FILE_DOWNLOAD_WORKERS} parallel downloads.  If any
+ * single download fails, outstanding downloads are cancelled via
+ * AbortController and the error propagates to the caller.
+ *
+ * @param files - Array of FileData objects to download
+ * @returns Array of ProcessedFile objects with base64-encoded data, or null
+ *          when the input array is empty / undefined
+ * @throws Error if any file download fails
+ */
+export async function processFiles(
+    files: FileData[] | null | undefined,
+): Promise<ProcessedFile[] | null> {
+    if (!files || files.length === 0) {
+        return null;
+    }
+
+    const timeoutMs = _getFileDownloadTimeout();
+    const limit: LimitFunction = pLimit(MAX_FILE_DOWNLOAD_WORKERS);
+    const controller = new AbortController();
+
+    const downloadPromises = files.map((file) =>
+        limit(() => _downloadSingleFile(file, timeoutMs, controller.signal)),
+    );
+
+    try {
+        return await Promise.all(downloadPromises);
+    } catch (error) {
+        controller.abort();
+        limit.clearQueue();
+        throw error;
+    }
+}
+
 /**
  * Execute a task function (sync or async) and extract message and session_id.
  *
+ * When `rawFiles` is provided, the files are downloaded and base64-encoded
+ * before being passed to the task's `run` method.
+ *
  * @param task - The task function to execute
  * @param message - The input message to pass to the task
  * @param sessionId - The current session identifier
+ * @param rawFiles - Optional array of file metadata to download and forward
  * @returns A tuple of [response_message, session_id]
- * @throws Error if the task returns an unsupported type
+ * @throws Error if the task returns an unsupported type or file download fails
  */
 export async function executeTask(
     task: BaseTask,
     message: string,
     sessionId: string | null,
+    rawFiles?: FileData[] | null,
 ): Promise<[string, string | null]> {
-    const result = task.run(message, sessionId);
+    const processedFiles = rawFiles && rawFiles.length > 0
+        ? await processFiles(rawFiles)
+        : null;
+
+    const result = task.run(message, sessionId, processedFiles);
 
-    // Check if result is a Promise (async function)
     const resolvedResult = result instanceof Promise ? await result : result;
 
-    // Validate that the result is a TaskResult
     if (
         typeof resolvedResult === "object" &&
         resolvedResult !== null &&