@loadmill/droid-cua 2.8.0 → 2.10.0

package/build/src/device/openai.js

@@ -1,15 +1,10 @@
-import OpenAI from "openai";
-import { readContextOptimizationEnv } from "../core/context-optimization.js";
 import { buildTestRevisionSystemPrompt } from "../prompts/editor.js";
-import { buildAppContextCompactionInput } from "../prompts/execution.js";
+import { createOpenAiClient, resolveRuntimeConfig } from "../core/runtime-config.js";
+import { DEFAULT_AIM_ASSIST_MODE } from "../core/aim-assist.js";
+import { DEFAULT_CONTEXT_OPTIMIZATION_THRESHOLD } from "../core/context-optimization.js";
 import { logger } from "../utils/logger.js";
 import { CuaDebugTracer } from "../utils/cua-debug-tracer.js";
-let openai = null;
 const cuaDebugTracer = new CuaDebugTracer(logger);
-const VALID_CUA_MODELS = new Set(["gpt-5.4", "gpt-5.4-mini", "computer-use-preview"]);
-function getSelectedCuaModel() {
-    return VALID_CUA_MODELS.has(process.env.OPENAI_CUA_MODEL) ? process.env.OPENAI_CUA_MODEL : "gpt-5.4";
-}
 function normalizeMessageContent(content) {
     if (Array.isArray(content)) {
         return content.map((item) => {
@@ -90,37 +85,16 @@ export function buildCuaInput({ messages, screenshotBase64, previousResponseId,
     }
     return input;
 }
-function buildContextManagement(cuaModel) {
-    if (cuaModel === "computer-use-preview") {
-        return undefined;
-    }
-    const { enabled, threshold } = readContextOptimizationEnv();
-    if (!enabled) {
+function buildContextManagement(contextOptimizationEnabled, contextOptimizationThreshold) {
+    if (!contextOptimizationEnabled) {
         return undefined;
     }
-    return [{ type: "compaction", compact_threshold: threshold }];
+    return [{ type: "compaction", compact_threshold: contextOptimizationThreshold }];
 }
-export function buildCuaRequestParams({ cuaModel, previousResponseId, deviceInfo, input }) {
-    const contextManagement = buildContextManagement(cuaModel);
-    if (cuaModel === "computer-use-preview") {
-        return {
-            model: cuaModel,
-            previous_response_id: previousResponseId || undefined,
-            input,
-            store: true,
-            truncation: "auto",
-            tools: [
-                {
-                    type: "computer_use_preview",
-                    display_width: deviceInfo.scaled_width,
-                    display_height: deviceInfo.scaled_height,
-                    environment: "browser",
-                },
-                previewClickTool,
-            ],
-            reasoning: { summary: "concise" },
-        };
-    }
+export function buildCuaRequestParams({ cuaModel, aimAssistMode = DEFAULT_AIM_ASSIST_MODE, previousResponseId, input, contextOptimizationEnabled = true, contextOptimizationThreshold = DEFAULT_CONTEXT_OPTIMIZATION_THRESHOLD, }) {
+    const contextManagement = buildContextManagement(contextOptimizationEnabled, contextOptimizationThreshold);
+    // Keep Responses API storage enabled on every turn so execution can chain
+    // reliably through previous_response_id across the whole run.
     return {
         model: cuaModel,
         reasoning: { effort: "medium" },
@@ -129,98 +103,26 @@ export function buildCuaRequestParams({ cuaModel, previousResponseId, deviceInfo
         input,
         store: true,
         truncation: "auto",
-        tools: [
-            {
-                type: "computer",
-            },
-            previewClickTool,
-        ],
+        tools: aimAssistMode === "off"
+            ? [{ type: "computer" }]
+            : [
+                { type: "computer" },
+                previewClickTool,
+            ],
     };
 }
-function shouldMapPreviewAccessError(err) {
-    const status = err?.status;
-    const code = typeof err?.code === "string" ? err.code.toLowerCase() : "";
-    const type = typeof err?.type === "string" ? err.type.toLowerCase() : "";
-    const message = typeof err?.message === "string" ? err.message.toLowerCase() : "";
-    if (status === 403 || status === 404)
-        return true;
-    if (code.includes("model_not_found") || code.includes("permission"))
-        return true;
-    if (type.includes("permission"))
-        return true;
-    return (message.includes("computer-use-preview") &&
-        (message.includes("access") || message.includes("permission") || message.includes("not found") || message.includes("unsupported")));
-}
-export function shouldMapMiniCompatibilityError(err) {
-    const status = Number(err?.status);
-    const code = typeof err?.code === "string" ? err.code.toLowerCase() : "";
-    const type = typeof err?.type === "string" ? err.type.toLowerCase() : "";
-    const message = typeof err?.message === "string" ? err.message.toLowerCase() : "";
-    const mentionsMiniModel = message.includes("gpt-5.4-mini");
-    const mentionsComputerTool = message.includes("computer") || message.includes("tool");
-    const mentionsUnsupported = (message.includes("unsupported") ||
-        message.includes("not supported") ||
-        message.includes("does not support"));
-    if (code.includes("model_not_found")) {
-        return true;
-    }
-    if (status === 404) {
-        return mentionsMiniModel && (message.includes("not found") || mentionsUnsupported);
-    }
-    if (code.includes("unsupported")) {
-        return mentionsMiniModel || mentionsComputerTool;
-    }
-    if (!type.includes("invalid")) {
-        return false;
-    }
-    return mentionsUnsupported && (mentionsMiniModel || mentionsComputerTool);
-}
-function mapCuaError(err, cuaModel) {
-    if (cuaModel === "computer-use-preview" && shouldMapPreviewAccessError(err)) {
-        const mapped = new Error("OpenAI API key does not have access to computer-use-preview. Switch to gpt-5.4 in Settings > CUA Model.");
-        mapped.cause = err;
-        return mapped;
-    }
-    if (cuaModel === "gpt-5.4-mini" && shouldMapMiniCompatibilityError(err)) {
-        const mapped = new Error("OpenAI rejected gpt-5.4-mini for computer-use execution. Switch to gpt-5.4 or computer-use-preview in Settings > CUA Model.");
-        mapped.cause = err;
-        return mapped;
-    }
-    return err;
-}
-export function isNonRetryableCuaError(err) {
-    const status = Number(err?.status);
-    const code = typeof err?.code === "string" ? err.code.toLowerCase() : "";
-    const type = typeof err?.type === "string" ? err.type.toLowerCase() : "";
-    const message = typeof err?.message === "string" ? err.message.toLowerCase() : "";
-    if ([401, 403, 404].includes(status)) {
-        return true;
-    }
-    return (code.includes("model_not_found") ||
-        code.includes("permission") ||
-        type.includes("permission") ||
-        message.includes("does not have access to computer-use-preview") ||
-        message.includes("switch to gpt-5.4 in settings > cua model") ||
-        message.includes("openai rejected gpt-5.4-mini for computer-use execution") ||
-        message.includes("switch to gpt-5.4 or computer-use-preview in settings > cua model"));
-}
-function getOpenAI() {
-    if (!openai) {
-        openai = new OpenAI({
-            apiKey: process.env.OPENAI_API_KEY,
-        });
-    }
-    return openai;
-}
 /**
  * Revise a test script based on user feedback using simple chat completion
- * @param {string} originalScript - The original test script
- * @param {string} revisionRequest - User's requested changes
+ * @param {{
+ *   client: import("openai").default,
+ *   originalScript: string,
+ *   revisionRequest: string,
+ * }} options
  * @returns {Promise<string>} - The revised test script
  */
-export async function reviseTestScript(originalScript, revisionRequest) {
-    const response = await getOpenAI().chat.completions.create({
-        model: "gpt-4o",
+export async function reviseTestScript({ client, originalScript, revisionRequest }) {
+    const response = await client.chat.completions.create({
+        model: "gpt-5.4",
         messages: [{
                 role: "system",
                 content: buildTestRevisionSystemPrompt(originalScript, revisionRequest)
@@ -228,29 +130,29 @@ export async function reviseTestScript(originalScript, revisionRequest) {
     });
     return response.choices[0].message.content.trim();
 }
-export async function compactAppContext({ contextDocument, taskDescription, tokenBudget }) {
-    const response = await getOpenAI().responses.create({
-        model: "gpt-5.4",
-        temperature: 0,
-        input: buildAppContextCompactionInput({
-            contextDocument,
-            taskDescription,
-            tokenBudget,
-        })
-    });
-    return {
-        briefing: typeof response.output_text === "string" ? response.output_text.trim() : "",
-        outputTokens: typeof response.usage?.output_tokens === "number" ? response.usage.output_tokens : null,
-    };
-}
-export async function sendCUARequest({ messages, screenshotBase64, previousResponseId, callId, deviceInfo, debugContext, }) {
-    const cuaModel = getSelectedCuaModel();
+/**
+ * @param {{
+ *   client: import("openai").default,
+ *   cuaModel: string,
+ *   contextOptimizationEnabled?: boolean,
+ *   contextOptimizationThreshold?: number,
+ *   messages: Array<unknown>,
+ *   screenshotBase64?: string,
+ *   previousResponseId?: string | null,
+ *   callId?: string | null,
+ *   deviceInfo?: unknown,
+ *   debugContext?: unknown,
+ * }} options
+ */
+export async function sendCUARequest({ client, cuaModel, aimAssistMode = DEFAULT_AIM_ASSIST_MODE, contextOptimizationEnabled = true, contextOptimizationThreshold = DEFAULT_CONTEXT_OPTIMIZATION_THRESHOLD, messages, screenshotBase64, previousResponseId, callId, deviceInfo, debugContext, }) {
     const input = buildCuaInput({ messages, screenshotBase64, previousResponseId, callId });
     const requestParams = buildCuaRequestParams({
         cuaModel,
+        aimAssistMode,
         previousResponseId,
-        deviceInfo,
         input,
+        contextOptimizationEnabled,
+        contextOptimizationThreshold,
     });
     const trace = cuaDebugTracer.startTurn({
         requestParams,
@@ -262,12 +164,44 @@ export async function sendCUARequest({ messages, screenshotBase64, previousRespo
     });
     logger.debug("CUA Request:", trace.requestLog);
     try {
-        const response = await getOpenAI().responses.create(requestParams);
+        const response = await client.responses.create(requestParams);
         cuaDebugTracer.onResponse(trace, response);
         return response;
     }
     catch (err) {
         cuaDebugTracer.onError(trace, err);
-        throw mapCuaError(err, cuaModel);
+        throw err;
     }
 }
+/**
+ * @param {{
+ *   client?: import("openai").default | null,
+ *   config?: Record<string, unknown>,
+ * }} options
+ */
+export function createOpenAiExecutor({ client, config }) {
+    const runtimeConfig = resolveRuntimeConfig(config);
+    const openAiClient = client || createOpenAiClient(null);
+    return {
+        client: openAiClient,
+        config: runtimeConfig,
+        async sendCUARequest(args) {
+            // Bind model and request-shaping config to the active run so execution
+            // stays deterministic even if settings change after the run starts.
+            return await sendCUARequest({
+                client: openAiClient,
+                cuaModel: runtimeConfig.cuaModel,
+                contextOptimizationEnabled: runtimeConfig.contextOptimizationEnabled,
+                contextOptimizationThreshold: runtimeConfig.contextOptimizationThreshold,
+                ...args,
+            });
+        },
+        async reviseTestScript(originalScript, revisionRequest) {
+            return await reviseTestScript({
+                client: openAiClient,
+                originalScript,
+                revisionRequest,
+            });
+        },
+    };
+}

package/build/src/integrations/loadmill/client.js

@@ -20,153 +20,171 @@ function normalizeApiBaseUrl(rawBaseUrl) {
     }
 }
 /**
- * Get Loadmill API token from environment
- * @returns {string|null}
+ * @param {{
+ *   token?: string | null,
+ *   baseUrl?: string | null,
+ *   onUnauthorized?: (() => Promise<unknown> | unknown) | null,
+ * }} [options]
  */
-export function getApiToken() {
-    return process.env.LOADMILL_API_TOKEN || null;
-}
-/**
- * Get Loadmill base URL from environment
- * @returns {string}
- */
-export function getBaseUrl() {
-    return normalizeApiBaseUrl(process.env.LOADMILL_BASE_URL);
-}
-/**
- * Make an authenticated request to Loadmill API
- * @param {string} endpoint - API endpoint (without base URL)
- * @param {Object} options - Fetch options
- * @returns {Promise<Object>}
- */
-async function apiRequest(endpoint, options = {}) {
-    const token = getApiToken();
-    if (!token) {
-        throw new Error("LOADMILL_API_TOKEN environment variable is not set");
+export function createLoadmillApiClient({ token = null, baseUrl = null, onUnauthorized = null, } = {}) {
+    function getApiToken() {
+        return typeof token === "string" && token.trim() ? token.trim() : null;
+    }
+    function getBaseUrl() {
+        return normalizeApiBaseUrl(baseUrl);
     }
-    const baseUrl = getBaseUrl();
-    const url = `${baseUrl}${endpoint}`;
-    const response = await fetch(url, {
-        ...options,
-        headers: {
-            "Authorization": `Bearer ${token}`,
-            "Content-Type": "application/json",
-            ...options.headers,
-        },
-    });
-    if (response.status === 401 || response.status === 403) {
-        const unauthorizedHandler = globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__;
-        if (typeof unauthorizedHandler === "function") {
-            await unauthorizedHandler();
+    async function apiRequest(endpoint, options = {}) {
+        const resolvedToken = getApiToken();
+        if (!resolvedToken) {
+            throw new Error("LOADMILL_API_TOKEN environment variable is not set");
+        }
+        const resolvedBaseUrl = getBaseUrl();
+        const url = `${resolvedBaseUrl}${endpoint}`;
+        const response = await fetch(url, {
+            ...options,
+            headers: {
+                "Authorization": `Bearer ${resolvedToken}`,
+                "Content-Type": "application/json",
+                ...options.headers,
+            },
+        });
+        if ((response.status === 401 || response.status === 403) && typeof onUnauthorized === "function") {
+            await onUnauthorized();
+        }
+        if (!response.ok) {
+            const errorText = await response.text();
+            throw new Error(`Loadmill API error (${response.status}): ${errorText}`);
         }
+        return response.json();
     }
-    if (!response.ok) {
-        const errorText = await response.text();
-        throw new Error(`Loadmill API error (${response.status}): ${errorText}`);
+    async function searchFlows(searchQuery, options = {}) {
+        const { limit = 10 } = options;
+        const encodedQuery = encodeURIComponent(searchQuery);
+        const result = await apiRequest(`/test-flows?search=${encodedQuery}&limit=${limit}`);
+        if (Array.isArray(result)) {
+            return result;
+        }
+        if (result && Array.isArray(result.flows)) {
+            return result.flows;
+        }
+        if (result && Array.isArray(result.data)) {
+            return result.data;
+        }
+        if (result && Array.isArray(result.items)) {
+            return result.items;
+        }
+        if (result && Array.isArray(result.testFlows)) {
+            return result.testFlows;
+        }
+        console.error("[Loadmill] Unexpected API response format:", JSON.stringify(result, null, 2));
+        return [];
     }
-    return response.json();
-}
-/**
- * Search for test flows by query
- * @param {string} searchQuery - Search query string
- * @param {Object} options - Additional options
- * @param {number} options.limit - Maximum number of results (default: 10)
- * @returns {Promise<Array>} - Array of matching flows
- */
-export async function searchFlows(searchQuery, options = {}) {
-    const { limit = 10 } = options;
-    const encodedQuery = encodeURIComponent(searchQuery);
-    const result = await apiRequest(`/test-flows?search=${encodedQuery}&limit=${limit}`);
-    // Handle different response formats
-    if (Array.isArray(result)) {
-        return result;
+    async function runTestSuite(suiteId, { flowIds = [], parameters = {} } = {}) {
+        const body = {};
+        if (flowIds.length > 0) {
+            body.flowIds = flowIds;
+        }
+        if (Object.keys(parameters).length > 0) {
+            body.parameters = parameters;
+        }
+        return await apiRequest(`/test-suites/${suiteId}/run`, {
+            method: "POST",
+            body: JSON.stringify(body),
+        });
     }
-    if (result && Array.isArray(result.flows)) {
-        return result.flows;
+    async function runTestFlow(flowId, suiteId, { parameters = {} } = {}) {
+        const body = {
+            flows: [flowId],
+            inlineParameterOverride: true,
+            sharedFlowVersionOverrides: [],
+        };
+        if (Object.keys(parameters).length > 0) {
+            body.overrideParameters = parameters;
+        }
+        return await apiRequest(`/test-suites/${suiteId}/run?ui=true`, {
+            method: "POST",
+            body: JSON.stringify(body),
+        });
     }
-    if (result && Array.isArray(result.data)) {
-        return result.data;
+    async function getTestRunStatus(testRunId) {
+        return await apiRequest(`/test-suites-runs/${testRunId}`);
     }
-    if (result && Array.isArray(result.items)) {
-        return result.items;
+    async function getFlowRunStatus(flowRunId) {
+        return await apiRequest(`/test-flows-runs/${flowRunId}`);
     }
-    if (result && Array.isArray(result.testFlows)) {
-        return result.testFlows;
+    async function explainFailure(testRunId) {
+        return await apiRequest(`/explain-failures`, {
+            method: "POST",
+            body: JSON.stringify({
+                testRunId,
+                testRunType: "flowRun"
+            }),
+        });
     }
-    // Log unexpected response format for debugging
-    console.error("[Loadmill] Unexpected API response format:", JSON.stringify(result, null, 2));
-    return [];
+    return {
+        getApiToken,
+        getBaseUrl,
+        searchFlows,
+        runTestSuite,
+        runTestFlow,
+        getTestRunStatus,
+        getFlowRunStatus,
+        explainFailure,
+    };
 }
-/**
- * Run a test suite with specified flows
- * @param {string} suiteId - Test suite ID
- * @param {Object} options - Run options
- * @param {string[]} options.flowIds - Array of flow IDs to run
- * @param {Object} options.parameters - Parameters to pass to the flows
- * @returns {Promise<Object>} - Run result with runId
- */
-export async function runTestSuite(suiteId, { flowIds = [], parameters = {} } = {}) {
-    const body = {};
-    if (flowIds.length > 0) {
-        body.flowIds = flowIds;
-    }
-    if (Object.keys(parameters).length > 0) {
-        body.parameters = parameters;
-    }
-    return apiRequest(`/test-suites/${suiteId}/run`, {
-        method: "POST",
-        body: JSON.stringify(body),
-    });
+export function getApiToken() {
+    return createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+    }).getApiToken();
 }
-/**
- * Run a single test flow
- * @param {string} flowId - Flow ID to run
- * @param {string} suiteId - Test suite ID containing the flow
- * @param {Object} options - Run options
- * @param {Object} options.parameters - Parameters to pass to the flow
- * @returns {Promise<Object>} - Run result with testSuiteRunId
- */
-export async function runTestFlow(flowId, suiteId, { parameters = {} } = {}) {
-    const body = {
-        flows: [flowId],
-        inlineParameterOverride: true,
-        sharedFlowVersionOverrides: [],
-    };
-    if (Object.keys(parameters).length > 0) {
-        body.overrideParameters = parameters;
-    }
-    return apiRequest(`/test-suites/${suiteId}/run?ui=true`, {
-        method: "POST",
-        body: JSON.stringify(body),
-    });
+export function getBaseUrl() {
+    return createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+    }).getBaseUrl();
+}
+export async function searchFlows(searchQuery, options = {}) {
+    // Legacy env-based wrapper kept for older call sites. New execution flows
+    // should pass an explicit client created from resolved run credentials.
+    return await createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+        onUnauthorized: globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__,
+    }).searchFlows(searchQuery, options);
+}
+export async function runTestSuite(suiteId, options = {}) {
+    return await createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+        onUnauthorized: globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__,
+    }).runTestSuite(suiteId, options);
+}
+export async function runTestFlow(flowId, suiteId, options = {}) {
+    return await createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+        onUnauthorized: globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__,
+    }).runTestFlow(flowId, suiteId, options);
 }
-/**
- * Get the status of a test run
- * @param {string} testRunId - Test run ID
- * @returns {Promise<Object>} - Test run status
- */
 export async function getTestRunStatus(testRunId) {
-    return apiRequest(`/test-suites-runs/${testRunId}`);
+    return await createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+        onUnauthorized: globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__,
+    }).getTestRunStatus(testRunId);
 }
-/**
- * Get the status of a flow run
- * @param {string} flowRunId - Flow run ID
- * @returns {Promise<Object>} - Flow run status
- */
 export async function getFlowRunStatus(flowRunId) {
-    return apiRequest(`/test-flows-runs/${flowRunId}`);
+    return await createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+        onUnauthorized: globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__,
+    }).getFlowRunStatus(flowRunId);
 }
-/**
- * Get AI-generated explanation for a failed test run
- * @param {string} testRunId - The flow run ID (from testSuiteFlowRuns[].id)
- * @returns {Promise<Object>} - Explanation with reason, suggestion, etc.
- */
 export async function explainFailure(testRunId) {
-    return apiRequest(`/explain-failures`, {
-        method: "POST",
-        body: JSON.stringify({
-            testRunId,
-            testRunType: "flowRun"
-        }),
-    });
+    return await createLoadmillApiClient({
+        token: process.env.LOADMILL_API_TOKEN,
+        baseUrl: process.env.LOADMILL_BASE_URL,
+        onUnauthorized: globalThis.__DROID_CUA_HANDLE_LOADMILL_UNAUTHORIZED__,
+    }).explainFailure(testRunId);
 }

package/build/src/integrations/loadmill/executor.js

@@ -1,7 +1,7 @@
 /**
  * Orchestrates Loadmill flow execution
  */
-import { searchFlows, runTestFlow, getTestRunStatus, getApiToken, explainFailure } from "./client.js";
+import { createLoadmillApiClient } from "./client.js";
 import { interpretLoadmillCommand, selectBestFlow } from "./interpreter.js";
 import { logger } from "../../utils/logger.js";
 const POLL_INTERVAL_MS = 5000; // 5 seconds
@@ -20,11 +20,11 @@ function sleep(ms) {
  * @param {Function} onStatusUpdate - Callback for status updates
  * @returns {Promise<Object>} - Final run status
  */
-async function pollForCompletion(runId, onStatusUpdate = () => { }) {
+async function pollForCompletion(runId, loadmillClient, onStatusUpdate = () => { }) {
     const startTime = Date.now();
     while (Date.now() - startTime < TIMEOUT_MS) {
         try {
-            const status = await getTestRunStatus(runId);
+            const status = await loadmillClient.getTestRunStatus(runId);
             onStatusUpdate(status);
             // Check if run is complete
             if (status.status === "PASSED" || status.status === "passed") {
@@ -36,7 +36,7 @@ async function pollForCompletion(runId, onStatusUpdate = () => { }) {
                 try {
                     const failedFlowRun = status.testSuiteFlowRuns?.find(f => f.status === "FAILED");
                     if (failedFlowRun) {
-                        const explanation = await explainFailure(failedFlowRun.id);
+                        const explanation = await loadmillClient.explainFailure(failedFlowRun.id);
                         if (explanation?.reason) {
                             failureReason = explanation.reason;
                         }
@@ -69,22 +69,28 @@ async function pollForCompletion(runId, onStatusUpdate = () => { }) {
  * @returns {Promise<{success: boolean, error?: string, result?: Object}>}
  */
 export async function executeLoadmillCommand(userInput, options = {}) {
-    const { onProgress = () => { } } = options;
+    const { onProgress = () => { }, loadmillClient = createLoadmillApiClient(), openAiClient = null, } = options;
     // Check for API token
-    if (!getApiToken()) {
+    if (!loadmillClient.getApiToken()) {
         return {
             success: false,
             error: "LOADMILL_API_TOKEN environment variable is not set. Please set it in your .env file."
         };
     }
+    if (!openAiClient) {
+        return {
+            success: false,
+            error: "OpenAI API key is not configured."
+        };
+    }
     try {
         // Step 1: Interpret the command
         onProgress({ step: "interpreting", message: "Interpreting command..." });
-        const interpreted = await interpretLoadmillCommand(userInput);
+        const interpreted = await interpretLoadmillCommand(userInput, { client: openAiClient });
         logger.debug("Loadmill command interpreted", interpreted);
         // Step 2: Search for flows
         onProgress({ step: "searching", message: `Searching for flows matching "${interpreted.searchQuery}"...` });
-        const flows = await searchFlows(interpreted.searchQuery);
+        const flows = await loadmillClient.searchFlows(interpreted.searchQuery);
         // Ensure flows is an array
         if (!Array.isArray(flows) || flows.length === 0) {
             return {
@@ -94,7 +100,7 @@ export async function executeLoadmillCommand(userInput, options = {}) {
         }
         // Step 3: Select best match
         onProgress({ step: "selecting", message: `Found ${flows.length} flow(s). Selecting best match...` });
-        const { selectedFlow, confidence } = await selectBestFlow(flows, interpreted.searchQuery);
+        const { selectedFlow, confidence } = await selectBestFlow(flows, interpreted.searchQuery, { client: openAiClient });
         if (!selectedFlow) {
             return {
                 success: false,
@@ -119,7 +125,7 @@ export async function executeLoadmillCommand(userInput, options = {}) {
             step: "running",
             message: `Running flow "${flowName}" (confidence: ${(confidence * 100).toFixed(0)}%)...`
         });
-        const runResult = await runTestFlow(selectedFlow.id, selectedFlow.testSuiteId, {
+        const runResult = await loadmillClient.runTestFlow(selectedFlow.id, selectedFlow.testSuiteId, {
             parameters: interpreted.parameters
         });
         const runId = runResult.testSuiteRunId || runResult.id || runResult.runId;
@@ -131,7 +137,7 @@ export async function executeLoadmillCommand(userInput, options = {}) {
         }
         // Step 5: Poll for completion
         onProgress({ step: "polling", message: `Test started (ID: ${runId}). Waiting for completion...`, runId });
-        const finalResult = await pollForCompletion(runId, (status) => {
+        const finalResult = await pollForCompletion(runId, loadmillClient, (status) => {
             onProgress({ step: "polling", message: `Status: ${status.status}...`, runId });
         });
         return {