@loadmill/droid-cua 2.8.0 → 2.10.0

package/build/src/core/execution-engine.js

@@ -9,7 +9,6 @@ import { getConfiguredStepDelayMs } from "../utils/step-delay.js";
 import { printCliOutput } from "../utils/console-output.js";
 import { annotateScreenshotArtifact, createAnnotatedScreenshotFromBase64 } from "../utils/run-screenshot-recorder.js";
 import { CAPTURE_SOURCES } from "./capture-sources.js";
-const VALID_CUA_MODELS = new Set(["gpt-5.4", "gpt-5.4-mini", "computer-use-preview"]);
 function extractActionableOutputItems(items) {
     const entries = [];
     for (const item of items) {
@@ -102,55 +101,20 @@ function buildPreviewOutputContent({ text, annotatedScreenshotBase64 = null }) {
     }
     return content;
 }
-function getScopeAndIds(context = null, stepContext = null) {
-    const scope = context?.sessionId ? "design" : "execution";
-    const ids = scope === "design"
-        ? {
-            sessionId: context?.sessionId,
-            stepId: stepContext?.stepId,
-            instructionIndex: stepContext?.instructionIndex
-        }
-        : {
-            runId: context?.runId,
-            stepId: stepContext?.stepId,
-            instructionIndex: stepContext?.instructionIndex
-        };
-    return { scope, ids };
-}
-function buildStrictModePlan(actions = [], strictMode = false) {
-    if (!strictMode || actions.length === 0) {
-        return {
-            actionsToExecute: actions,
-            droppedActions: [],
-            truncationReason: null,
-            runtimeNote: null
-        };
-    }
-    const actionsToExecute = [actions[0]];
-    const droppedActions = actions.slice(1);
-    if (droppedActions.length === 0) {
-        return {
-            actionsToExecute,
-            droppedActions,
-            truncationReason: null,
-            runtimeNote: null
-        };
-    }
-    const leadingActionType = actionsToExecute[0]?.type;
-    const truncationReason = leadingActionType === "screenshot"
-        ? "leading_screenshot_reobserve"
-        : "post_first_action_reobserve";
-    const droppedActionTypes = droppedActions
-        .map((action) => action?.type)
-        .filter(Boolean);
-    const runtimeNote = `Strict Mode: I executed only the first ${leadingActionType === "screenshot" ? "screenshot request" : "action"} from your previous chain and intentionally skipped the remaining ${droppedActions.length} action${droppedActions.length === 1 ? "" : "s"} (${droppedActionTypes.join(", ")}) so I could re-observe the device before continuing. Base your next step only on what is visible now.`;
+function buildDebugContext(context = null, stepContext = null) {
     return {
-        actionsToExecute,
-        droppedActions,
-        truncationReason,
-        runtimeNote
+        scope: context?.sessionId ? "design" : "execution",
+        runId: context?.runId,
+        sessionId: context?.sessionId,
+        stepId: stepContext?.stepId,
+        instructionIndex: stepContext?.instructionIndex
     };
 }
+function collectPendingSafetyChecks(items = []) {
+    return items
+        .filter((item) => item.type === "pending_safety_check")
+        .map((item) => ({ id: item.id }));
+}
 function normalizeArtifactPaths(result) {
     if (result && typeof result === "object" && !Array.isArray(result)) {
         if (Array.isArray(result.annotationSourcePaths)) {
@@ -188,14 +152,14 @@ export class ExecutionEngine {
         this.recordScreenshots = options.recordScreenshots || false;
         this.screenshotDir = options.screenshotDir || null;
         this.screenshotRecorder = options.screenshotRecorder || null;
-        this.strictMode = options.strictMode === true;
-        this.stepDelayMs = getConfiguredStepDelayMs();
+        this.stepDelayMs = options.stepDelayMs ?? getConfiguredStepDelayMs();
         this.reportedScreenshotWriteError = false;
         this.reportedScreenshotAnnotateError = false;
         this.lastRecordedScreenshotPaths = [];
         this.getScreenshotAsBase64 = options.getScreenshotAsBase64 || getScreenshotAsBase64;
         this.handleModelAction = options.handleModelAction || handleModelAction;
         this.sendCUARequest = options.sendCUARequest || sendCUARequest;
+        this.reviseTestScript = options.reviseTestScript || null;
         this.getCurrentPlatform = options.getCurrentPlatform || getCurrentPlatform;
     }
     async saveScreenshotArtifact(screenshotBase64, metadata = {}, options = {}) {
@@ -374,6 +338,108 @@ export class ExecutionEngine {
             }));
         }
     }
+    async requestNextTurn(messages, previousResponseId, context = null, stepContext = null) {
+        return await this.sendCUARequest({
+            messages,
+            previousResponseId,
+            deviceInfo: this.session.deviceInfo,
+            debugContext: buildDebugContext(context, stepContext)
+        });
+    }
+    async executeActionBatch(actions, shouldStop, context = null, stepContext = null, trackAction = null) {
+        let sawExplicitScreenshotAction = false;
+        for (const action of actions) {
+            const actionEventId = randomUUID();
+            if (shouldStop()) {
+                return { stopped: true, sawExplicitScreenshotAction };
+            }
+            if (action.type === "screenshot") {
+                sawExplicitScreenshotAction = true;
+                continue;
+            }
+            const annotatedScreenshotResult = await this.annotateLatestScreenshot(action, {
+                runId: context?.runId,
+                sessionId: context?.sessionId,
+                stepId: stepContext?.stepId,
+                instructionIndex: stepContext?.instructionIndex,
+                eventId: actionEventId,
+                captureSource: CAPTURE_SOURCES.ACTION_CLICK_ANNOTATION
+            });
+            const annotatedScreenshotPath = screenshotPathFromRecorderResult(annotatedScreenshotResult);
+            const annotatedScreenshotDataUrl = screenshotDataUrlFromRecorderResult(annotatedScreenshotResult);
+            await this.handleModelAction(this.session.deviceId, action, this.session.deviceInfo.scale, {
+                ...context,
+                shouldStop,
+                stepId: stepContext?.stepId,
+                instructionIndex: stepContext?.instructionIndex,
+                eventId: actionEventId,
+                captureSource: CAPTURE_SOURCES.ACTION_CLICK_ANNOTATION,
+                ...(annotatedScreenshotPath ? { screenshotPath: annotatedScreenshotPath } : {}),
+                ...(context?.includeScreenshotDataUrl && annotatedScreenshotDataUrl ? { screenshotDataUrl: annotatedScreenshotDataUrl } : {})
+            });
+            if (trackAction) {
+                trackAction(action);
+            }
+            if (shouldStop()) {
+                return { stopped: true, sawExplicitScreenshotAction };
+            }
+            if (action.type !== "wait") {
+                await new Promise((resolve) => setTimeout(resolve, this.stepDelayMs));
+            }
+        }
+        return { stopped: false, sawExplicitScreenshotAction };
+    }
+    async buildActionFollowUp(callId, pendingSafetyChecks, sawExplicitScreenshotAction, addOutput, eventMeta, context = null, stepContext = null) {
+        const screenshotBase64 = await this.getScreenshotAsBase64(this.session.deviceId, this.session.deviceInfo);
+        emitDesktopDebug("device.screenshot", "device", {
+            runId: context?.runId,
+            stepId: stepContext?.stepId,
+            instructionIndex: stepContext?.instructionIndex
+        }, {
+            source: sawExplicitScreenshotAction ? "explicit_action" : "post_action",
+            callId,
+            width: this.session.deviceInfo?.scaled_width,
+            height: this.session.deviceInfo?.scaled_height,
+            base64Length: screenshotBase64.length
+        });
+        const screenshotEventId = randomUUID();
+        const captureSource = sawExplicitScreenshotAction ? CAPTURE_SOURCES.CALL_OUTPUT_EXPLICIT_ACTION : CAPTURE_SOURCES.CALL_OUTPUT_POST_ACTION;
+        const screenshotResult = await this.recordScreenshot(screenshotBase64, {
+            runId: context?.runId,
+            sessionId: context?.sessionId,
+            stepId: stepContext?.stepId,
+            instructionIndex: stepContext?.instructionIndex,
+            eventId: screenshotEventId,
+            captureSource
+        });
+        const screenshotPath = screenshotPathFromRecorderResult(screenshotResult);
+        addOutput({
+            type: "info",
+            text: "Capturing screen",
+            eventType: "screenshot_captured",
+            actionType: "screenshot",
+            ...eventMeta({
+                payload: {
+                    eventId: screenshotEventId,
+                    captureSource,
+                    source: sawExplicitScreenshotAction ? "explicit_action" : "post_action",
+                    ...(typeof screenshotPath === "string" && screenshotPath ? { screenshotPath } : {}),
+                    ...(context?.includeScreenshotDataUrl
+                        ? { screenshotDataUrl: `data:image/png;base64,${screenshotBase64}` }
+                        : {})
+                }
+            })
+        });
+        return [{
+                type: "computer_call_output",
+                call_id: callId,
+                output: {
+                    type: "computer_screenshot",
+                    image_url: `data:image/png;base64,${screenshotBase64}`,
+                },
+                ...(pendingSafetyChecks.length > 0 ? { acknowledged_safety_checks: pendingSafetyChecks } : {})
+            }];
+    }
     /**
      * Run a full turn with the CUA model
      * Executes actions until the model stops requesting actions
@@ -386,7 +452,6 @@ export class ExecutionEngine {
         const addOutput = context?.addOutput || printCliOutput;
         let newResponseId = response.id;
         const shouldStop = () => Boolean(trackAction?.());
-        const { scope, ids } = getScopeAndIds(context, stepContext);
         const eventMeta = (extra = {}) => ({
             runId: context?.runId,
             stepId: stepContext?.stepId,
@@ -400,10 +465,7 @@ export class ExecutionEngine {
             }
             const items = response.output || [];
             const actionableItems = extractActionableOutputItems(items);
-            // ── Collect pending safety checks ──
-            const pendingSafetyChecks = items
-                .filter(item => item.type === "pending_safety_check")
-                .map(item => ({ id: item.id }));
+            const pendingSafetyChecks = collectPendingSafetyChecks(items);
             // ── Print assistant output ──
             for (const item of items) {
                 if (item.type === "reasoning") {
@@ -461,160 +523,20 @@ export class ExecutionEngine {
                 if (shouldStop()) {
                     return newResponseId;
                 }
-                response = await this.sendCUARequest({
-                    messages: input,
-                    previousResponseId: newResponseId,
-                    deviceInfo: this.session.deviceInfo,
-                    debugContext: {
-                        scope,
-                        runId: context?.runId,
-                        sessionId: context?.sessionId,
-                        stepId: stepContext?.stepId,
-                        instructionIndex: stepContext?.instructionIndex
-                    }
-                });
+                response = await this.requestNextTurn(input, newResponseId, context, stepContext);
                 newResponseId = response.id;
                 continue;
             }
             const { call_id, actions } = actionableItem;
-            const { actionsToExecute, droppedActions, truncationReason, runtimeNote } = buildStrictModePlan(actions, this.strictMode);
-            let sawExplicitScreenshotAction = false;
-            for (const action of actionsToExecute) {
-                const actionEventId = randomUUID();
-                if (shouldStop()) {
-                    return newResponseId;
-                }
-                if (action.type === "screenshot") {
-                    sawExplicitScreenshotAction = true;
-                }
-                else {
-                    const annotatedScreenshotResult = await this.annotateLatestScreenshot(action, {
-                        runId: context?.runId,
-                        sessionId: context?.sessionId,
-                        stepId: stepContext?.stepId,
-                        instructionIndex: stepContext?.instructionIndex,
-                        eventId: actionEventId,
-                        captureSource: CAPTURE_SOURCES.ACTION_CLICK_ANNOTATION
-                    });
-                    const annotatedScreenshotPath = screenshotPathFromRecorderResult(annotatedScreenshotResult);
-                    const annotatedScreenshotDataUrl = screenshotDataUrlFromRecorderResult(annotatedScreenshotResult);
-                    await this.handleModelAction(this.session.deviceId, action, this.session.deviceInfo.scale, {
-                        ...context,
-                        shouldStop,
-                        stepId: stepContext?.stepId,
-                        instructionIndex: stepContext?.instructionIndex,
-                        eventId: actionEventId,
-                        captureSource: CAPTURE_SOURCES.ACTION_CLICK_ANNOTATION,
-                        ...(annotatedScreenshotPath ? { screenshotPath: annotatedScreenshotPath } : {}),
-                        ...(context?.includeScreenshotDataUrl && annotatedScreenshotDataUrl ? { screenshotDataUrl: annotatedScreenshotDataUrl } : {})
-                    });
-                    if (trackAction) {
-                        trackAction(action);
-                    }
-                    if (shouldStop()) {
-                        return newResponseId;
-                    }
-                    if (action.type !== "wait") {
-                        await new Promise(resolve => setTimeout(resolve, this.stepDelayMs));
-                    }
-                }
-            }
-            if (shouldStop()) {
+            const actionBatchResult = await this.executeActionBatch(actions, shouldStop, context, stepContext, trackAction);
+            if (actionBatchResult.stopped) {
                 return newResponseId;
             }
-            const screenshotBase64 = await this.getScreenshotAsBase64(this.session.deviceId, this.session.deviceInfo);
-            emitDesktopDebug("device.screenshot", "device", {
-                runId: context?.runId,
-                stepId: stepContext?.stepId,
-                instructionIndex: stepContext?.instructionIndex
-            }, {
-                source: sawExplicitScreenshotAction ? "explicit_action" : "post_action",
-                callId: call_id,
-                width: this.session.deviceInfo?.scaled_width,
-                height: this.session.deviceInfo?.scaled_height,
-                base64Length: screenshotBase64.length
-            });
-            const screenshotEventId = randomUUID();
-            const captureSource = sawExplicitScreenshotAction ? CAPTURE_SOURCES.CALL_OUTPUT_EXPLICIT_ACTION : CAPTURE_SOURCES.CALL_OUTPUT_POST_ACTION;
-            const screenshotResult = await this.recordScreenshot(screenshotBase64, {
-                runId: context?.runId,
-                sessionId: context?.sessionId,
-                stepId: stepContext?.stepId,
-                instructionIndex: stepContext?.instructionIndex,
-                eventId: screenshotEventId,
-                captureSource
-            });
-            const screenshotPath = screenshotPathFromRecorderResult(screenshotResult);
-            addOutput({
-                type: 'info',
-                text: 'Capturing screen',
-                eventType: 'screenshot_captured',
-                actionType: 'screenshot',
-                ...eventMeta({
-                    payload: {
-                        eventId: screenshotEventId,
-                        captureSource,
-                        source: sawExplicitScreenshotAction ? "explicit_action" : "post_action",
-                        ...(typeof screenshotPath === "string" && screenshotPath ? { screenshotPath } : {}),
-                        ...(context?.includeScreenshotDataUrl
-                            ? { screenshotDataUrl: `data:image/png;base64,${screenshotBase64}` }
-                            : {})
-                    }
-                })
-            });
-            if (runtimeNote) {
-                const executedActionTypes = actionsToExecute
-                    .map((action) => action?.type)
-                    .filter(Boolean);
-                const droppedActionTypes = droppedActions
-                    .map((action) => action?.type)
-                    .filter(Boolean);
-                emitDesktopDebug("cua.strict_mode.truncation", scope, ids, {
-                    callId: call_id,
-                    executedCount: actionsToExecute.length,
-                    droppedCount: droppedActions.length,
-                    executedActionTypes,
-                    droppedActionTypes,
-                    reason: truncationReason,
-                    note: runtimeNote
-                });
-            }
-            const selectedCuaModel = VALID_CUA_MODELS.has(process.env.OPENAI_CUA_MODEL)
-                ? process.env.OPENAI_CUA_MODEL
-                : "gpt-5.4";
-            const input = [{
-                    type: "computer_call_output",
-                    call_id,
-                    output: {
-                        type: "computer_screenshot",
-                        image_url: `data:image/png;base64,${screenshotBase64}`,
-                    },
-                    ...(selectedCuaModel === "computer-use-preview"
-                        ? { current_url: this.getCurrentPlatform() === "ios" ? "ios://simulator" : "android://device" }
-                        : {}),
-                    ...(pendingSafetyChecks.length > 0 ? { acknowledged_safety_checks: pendingSafetyChecks } : {})
-                }];
-            if (runtimeNote) {
-                input.push({
-                    role: "user",
-                    content: runtimeNote
-                });
-            }
+            const input = await this.buildActionFollowUp(call_id, pendingSafetyChecks, actionBatchResult.sawExplicitScreenshotAction, addOutput, eventMeta, context, stepContext);
             if (shouldStop()) {
                 return newResponseId;
             }
-            response = await this.sendCUARequest({
-                messages: input,
-                previousResponseId: newResponseId,
-                deviceInfo: this.session.deviceInfo,
-                debugContext: {
-                    scope,
-                    runId: context?.runId,
-                    sessionId: context?.sessionId,
-                    stepId: stepContext?.stepId,
-                    instructionIndex: stepContext?.instructionIndex
-                }
-            });
+            response = await this.requestNextTurn(input, newResponseId, context, stepContext);
             newResponseId = response.id;
         }
         // ── At end, if last output was only reasoning ──

package/build/src/core/runtime-config.js

@@ -0,0 +1,95 @@
+import OpenAI from "openai";
+import { readContextOptimizationEnv, normalizeContextOptimizationEnabled, normalizeContextOptimizationThreshold, } from "./context-optimization.js";
+import { resolveCuaModel } from "./cua-models.js";
+import { normalizeStepDelayMs } from "../utils/step-delay.js";
+export function createEmptyPromptCustomizations() {
+    return {
+        basePromptInstructions: "",
+        designModeInstructions: "",
+        executionModeInstructions: "",
+    };
+}
+export function resolvePromptCustomizations(rawValue = null) {
+    const resolved = createEmptyPromptCustomizations();
+    if (!rawValue || typeof rawValue !== "object" || Array.isArray(rawValue)) {
+        return resolved;
+    }
+    for (const key of Object.keys(resolved)) {
+        if (typeof rawValue[key] === "string") {
+            resolved[key] = rawValue[key];
+        }
+    }
+    return resolved;
+}
+function trimNullableString(value) {
+    if (typeof value !== "string") {
+        return null;
+    }
+    const trimmed = value.trim();
+    return trimmed || null;
+}
+export function resolveRuntimeCredentials({ openAiApiKey = null, loadmillApiToken = null, loadmillBaseUrl = null, } = {}) {
+    return {
+        openAiApiKey: trimNullableString(openAiApiKey),
+        loadmillApiToken: trimNullableString(loadmillApiToken),
+        loadmillBaseUrl: trimNullableString(loadmillBaseUrl),
+    };
+}
+export function resolveRuntimeConfig({ cuaModel = null, stepDelayMs = undefined, contextOptimizationEnabled = undefined, contextOptimizationThreshold = undefined, promptCustomizations = null, appContext = null, } = {}) {
+    return {
+        cuaModel: resolveCuaModel(cuaModel),
+        stepDelayMs: normalizeStepDelayMs(Number(stepDelayMs)),
+        contextOptimizationEnabled: normalizeContextOptimizationEnabled(contextOptimizationEnabled),
+        contextOptimizationThreshold: normalizeContextOptimizationThreshold(contextOptimizationThreshold),
+        promptCustomizations: resolvePromptCustomizations(promptCustomizations),
+        appContext: appContext && typeof appContext === "object"
+            ? {
+                enabled: appContext.enabled !== false,
+                path: trimNullableString(appContext.path ?? appContext.contextPath),
+                status: typeof appContext.status === "string" ? appContext.status : null,
+                text: typeof appContext.text === "string" ? appContext.text : "",
+                message: trimNullableString(appContext.message),
+            }
+            : {
+                enabled: true,
+                path: null,
+                status: null,
+                text: "",
+                message: null,
+            },
+    };
+}
+export function resolveRuntimeConfigFromEnv(env = process.env, overrides = {}) {
+    const contextOptimization = readContextOptimizationEnv(env);
+    return resolveRuntimeConfig({
+        cuaModel: env.OPENAI_CUA_MODEL,
+        stepDelayMs: env.DROID_CUA_STEP_DELAY_MS,
+        contextOptimizationEnabled: contextOptimization.enabled,
+        contextOptimizationThreshold: contextOptimization.threshold,
+        ...overrides,
+    });
+}
+export function resolveRuntimeCredentialsFromEnv(env = process.env, overrides = {}) {
+    return resolveRuntimeCredentials({
+        openAiApiKey: env.OPENAI_API_KEY,
+        loadmillApiToken: env.LOADMILL_API_TOKEN,
+        loadmillBaseUrl: env.LOADMILL_BASE_URL,
+        ...overrides,
+    });
+}
+export function createOpenAiClient(apiKey) {
+    return new OpenAI({
+        apiKey: trimNullableString(apiKey) || undefined,
+    });
+}
+export function createRunContext({ config, credentials }) {
+    const resolvedConfig = resolveRuntimeConfig(config);
+    const resolvedCredentials = resolveRuntimeCredentials(credentials);
+    return {
+        config: resolvedConfig,
+        credentials: resolvedCredentials,
+        clients: {
+            openai: createOpenAiClient(resolvedCredentials.openAiApiKey),
+        },
+    };
+}

package/build/src/device/loadmill.js

@@ -3,8 +3,8 @@
  */
 import { executeLoadmillCommand } from "../integrations/loadmill/index.js";
 import { printCliOutput } from "../utils/console-output.js";
-function getLoadmillSiteBaseUrl() {
-    const rawBaseUrl = process.env.LOADMILL_BASE_URL || "https://app.loadmill.com/api";
+function getLoadmillSiteBaseUrl(loadmillClient = null) {
+    const rawBaseUrl = loadmillClient?.getBaseUrl?.() || "https://app.loadmill.com/api";
     return rawBaseUrl.replace(/\/api\/?$/, "");
 }
 /**
@@ -45,6 +45,8 @@ export function extractLoadmillCommand(userInput) {
  */
 export async function executeLoadmillInstruction(command, isHeadlessMode, context, stepContext = null) {
     const addOutput = context?.addOutput || printCliOutput;
+    const loadmillClient = context?.loadmillClient || null;
+    const openAiClient = context?.openAiClient || null;
     const meta = {
         runId: context?.runId,
         stepId: stepContext?.stepId,
@@ -54,14 +56,16 @@ export async function executeLoadmillInstruction(command, isHeadlessMode, contex
         type: 'info',
         text: `[Loadmill] Executing: ${command}`,
         eventType: 'system_message',
-        payload: { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() },
+        payload: { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) },
         ...meta
     });
     const result = await executeLoadmillCommand(command, {
+        loadmillClient,
+        openAiClient,
         onProgress: ({ message, runId }) => {
             const payload = runId
-                ? { loadmillSuiteRunId: runId, loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() }
-                : { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() };
+                ? { loadmillSuiteRunId: runId, loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) }
+                : { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) };
             addOutput({ type: 'info', text: `[Loadmill] ${message}`, eventType: 'system_message', payload, ...meta });
         }
     });
@@ -84,14 +88,15 @@ export async function executeLoadmillInstruction(command, isHeadlessMode, contex
  */
 export async function handleLoadmillFailure(command, error, isHeadlessMode, context, stepContext = null, suiteRunId = null) {
     const addOutput = context?.addOutput || printCliOutput;
+    const loadmillClient = context?.loadmillClient || null;
     const meta = {
         runId: context?.runId,
         stepId: stepContext?.stepId,
         instructionIndex: stepContext?.instructionIndex
     };
     const payload = suiteRunId
-        ? { loadmillSuiteRunId: suiteRunId, loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() }
-        : { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() };
+        ? { loadmillSuiteRunId: suiteRunId, loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) }
+        : { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) };
     addOutput({ type: 'error', text: '[Loadmill] FAILED', eventType: 'error', payload, ...meta });
     addOutput({ type: 'error', text: `Command: ${command}`, eventType: 'error', payload, ...meta });
     addOutput({ type: 'error', text: `Error: ${error}`, eventType: 'error', payload, ...meta });
@@ -148,14 +153,15 @@ export async function handleLoadmillFailure(command, error, isHeadlessMode, cont
  */
 export function handleLoadmillSuccess(command, result, context, stepContext = null) {
     const addOutput = context?.addOutput || printCliOutput;
+    const loadmillClient = context?.loadmillClient || null;
     const meta = {
         runId: context?.runId,
         stepId: stepContext?.stepId,
         instructionIndex: stepContext?.instructionIndex
     };
     const payload = result.runId
-        ? { loadmillSuiteRunId: result.runId, loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() }
-        : { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl() };
+        ? { loadmillSuiteRunId: result.runId, loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) }
+        : { loadmillCommand: command, loadmillBaseUrl: getLoadmillSiteBaseUrl(loadmillClient) };
     if (result.action === "search") {
         addOutput({ type: 'success', text: `[Loadmill] Found ${result.result.flows.length} flow(s)`, eventType: 'system_message', payload, ...meta });
         result.result.flows.forEach((flow, i) => {