fix(askui-nodejs): only report a screenshot for the final waitUntil failure

packages/askui-nodejs/src/execution/dsl.ts

@@ -59,6 +59,11 @@ export interface ExecOptions {
   modelComposition?: ModelCompositionBranch[];
   skipCache?: boolean;
   retryStrategy?: RetryStrategy;
+  /**
+   * Prepended to the reported instruction/step title, e.g., to mark a command
+   * as having been run as part of a `waitUntil` call.
+   */
+  instructionPrefix?: string;
 }
 
 abstract class FluentBase {
@@ -85,6 +90,7 @@ abstract class FluentBase {
     modelComposition: ModelCompositionBranch[] = [],
     skipCache = false,
     retryStrategy?: RetryStrategy,
+    instructionPrefix?: string,
     currentInstruction = '',
     paramsList: Map<string, unknown[]> = new Map<string, unknown[]>(),
   ): Promise<void> {
@@ -95,7 +101,7 @@ abstract class FluentBase {
       const customElements = newParamsList.has('customElement') ? newParamsList.get('customElement') as CustomElementJson[] : [];
       const aiElementNames = newParamsList.has('aiElementName') ? newParamsList.get('aiElementName') as string[] : [];
       return fluentCommand.fluentCommandExecutor(
-        newCurrentInstruction.trim(),
+        instructionPrefix ? `${instructionPrefix}${newCurrentInstruction.trim()}` : newCurrentInstruction.trim(),
         modelComposition,
         {
           customElementsJson: customElements,
@@ -112,6 +118,7 @@ abstract class FluentBase {
       modelComposition,
       skipCache,
       retryStrategy,
+      instructionPrefix,
       newCurrentInstruction,
       newParamsList,
     );
@@ -157,7 +164,7 @@ export class Exec extends FluentBase implements Executable {
   exec(execOptions?: ExecOptions): Promise<void> {
     const originStacktrace = { stack: '' };
     Error.captureStackTrace(originStacktrace, this.exec);
-    return this.fluentCommandStringBuilder(execOptions?.modelComposition, execOptions?.skipCache, execOptions?.retryStrategy).catch((err: Error) => Promise.reject(rewriteStackTraceForError(err, originStacktrace.stack)));
+    return this.fluentCommandStringBuilder(execOptions?.modelComposition, execOptions?.skipCache, execOptions?.retryStrategy, execOptions?.instructionPrefix).catch((err: Error) => Promise.reject(rewriteStackTraceForError(err, originStacktrace.stack)));
   }
 }
 
@@ -1242,7 +1249,7 @@ export class FluentFiltersOrRelations extends FluentFilters {
   exec(execOptions?: ExecOptions): Promise<void> {
     const originStacktrace = { stack: '' };
     Error.captureStackTrace(originStacktrace, this.exec);
-    return this.fluentCommandStringBuilder(execOptions?.modelComposition, execOptions?.skipCache, execOptions?.retryStrategy).catch((err: Error) => Promise.reject(rewriteStackTraceForError(err, originStacktrace.stack)));
+    return this.fluentCommandStringBuilder(execOptions?.modelComposition, execOptions?.skipCache, execOptions?.retryStrategy, execOptions?.instructionPrefix).catch((err: Error) => Promise.reject(rewriteStackTraceForError(err, originStacktrace.stack)));
   }
 }
 

packages/askui-nodejs/src/execution/ui-control-client.spec.ts

@@ -0,0 +1,135 @@
+import { UiControlClient } from './ui-control-client';
+import { ExecutionRuntime } from './execution-runtime';
+import { StepReporter, Reporter, Step } from '../core/reporting';
+import { DeviceClient } from './device-client';
+import { InferenceClient } from './inference-client';
+import { ControlCommand } from '../core/ui-control-commands/control-command';
+import { ControlCommandCode } from '../core/ui-control-commands/control-command-code';
+import { UiControllerClientConnectionState } from './ui-controller-client-connection-state';
+import { NoRetryStrategy } from './retry-strategies';
+import { AskUIAgent } from '../core/models/anthropic';
+import { ControlCommandError } from './control-command-error';
+import { AIElementArgs } from '../core/ai-element/ai-elements-args';
+import { Annotation } from '../core/annotation/annotation';
+
+jest.mock('../lib/logger');
+
+type UiControlClientConstructor = new (
+  workspaceId: string | undefined,
+  executionRuntime: ExecutionRuntime,
+  stepReporter: StepReporter,
+  aiElementArgs: AIElementArgs,
+  agent: AskUIAgent,
+) => UiControlClient;
+
+function buildDeviceClient(): jest.Mocked<DeviceClient> {
+  return {
+    connect: jest.fn(),
+    connectionState: UiControllerClientConnectionState.CONNECTED,
+    disconnect: jest.fn(),
+    getStartingArguments: jest.fn().mockResolvedValue({}),
+    requestControl: jest.fn().mockResolvedValue(undefined),
+    requestScreenshot: jest.fn().mockResolvedValue('base64-screenshot'),
+    setActiveDisplay: jest.fn().mockResolvedValue(undefined),
+  };
+}
+
+function buildInferenceClient(predictControlCommand: jest.Mock): InferenceClient {
+  return {
+    cacheManager: { loadFromFile: jest.fn(), saveToFile: jest.fn() },
+    isImageRequired: jest.fn().mockResolvedValue(false),
+    predictControlCommand,
+    predictImageAnnotation: jest.fn().mockResolvedValue(new Annotation('base64-annotated-image')),
+  } as unknown as InferenceClient;
+}
+
+function buildReporter(): Reporter {
+  return {
+    config: {},
+    onStepBegin: jest.fn().mockResolvedValue(undefined),
+    onStepEnd: jest.fn().mockResolvedValue(undefined),
+    onStepRetry: jest.fn().mockResolvedValue(undefined),
+  };
+}
+
+function buildClient(predictControlCommand: jest.Mock) {
+  const deviceClient = buildDeviceClient();
+  const inferenceClient = buildInferenceClient(predictControlCommand);
+  const reporter = buildReporter();
+  const stepReporter = new StepReporter(reporter);
+  const executionRuntime = new ExecutionRuntime(
+    deviceClient,
+    inferenceClient,
+    stepReporter,
+    new NoRetryStrategy(),
+  );
+  const agent = new AskUIAgent(executionRuntime);
+  const aiElementArgs: AIElementArgs = { additionalLocations: [], onLocationNotExist: 'error' };
+  const Ctor = UiControlClient as unknown as UiControlClientConstructor;
+  const client = new Ctor(undefined, executionRuntime, stepReporter, aiElementArgs, agent);
+  return { client, deviceClient, reporter };
+}
+
+const okCommand = () => new ControlCommand(ControlCommandCode.OK, []);
+const errorCommand = () => new ControlCommand(ControlCommandCode.ERROR, [], false);
+
+describe('UiControlClient.waitUntil', () => {
+  it('reports a single step with a single onStepEnd when the command succeeds immediately', async () => {
+    const predictControlCommand = jest.fn().mockResolvedValue(okCommand());
+    const { client, reporter } = buildClient(predictControlCommand);
+
+    await client.waitUntil(client.click().button(), 5, 1);
+
+    expect(reporter.onStepBegin).toHaveBeenCalledTimes(1);
+    expect(reporter.onStepRetry).not.toHaveBeenCalled();
+    expect(reporter.onStepEnd).toHaveBeenCalledTimes(1);
+    const [step] = (reporter.onStepEnd as jest.Mock).mock.calls[0] as [Step];
+    expect(step.error).toBeUndefined();
+  });
+
+  it('prefixes the reported instruction so waitUntil steps are recognizable in the report', async () => {
+    const predictControlCommand = jest.fn().mockResolvedValue(okCommand());
+    const { client, reporter } = buildClient(predictControlCommand);
+
+    await client.waitUntil(client.click().button(), 5, 1);
+
+    const [step] = (reporter.onStepEnd as jest.Mock).mock.calls[0] as [Step];
+    expect(step.instruction.value).toBe('waitUntil: Click on button');
+  });
+
+  it('reports failed attempts as retries of the same step, not as separate steps', async () => {
+    const predictControlCommand = jest.fn()
+      .mockResolvedValueOnce(errorCommand())
+      .mockResolvedValueOnce(errorCommand())
+      .mockResolvedValueOnce(okCommand());
+    const { client, reporter } = buildClient(predictControlCommand);
+
+    await client.waitUntil(client.click().button(), 5, 1);
+
+    expect(reporter.onStepBegin).toHaveBeenCalledTimes(1);
+    expect(reporter.onStepRetry).toHaveBeenCalledTimes(2);
+    expect(reporter.onStepEnd).toHaveBeenCalledTimes(1);
+    const [step] = (reporter.onStepEnd as jest.Mock).mock.calls[0] as [Step];
+    expect(step.error).toBeUndefined();
+    expect(step.retryCount).toBe(2);
+  });
+
+  it('only reports (and saves) a screenshot for the last, truly final failure, not for every failed attempt', async () => {
+    const predictControlCommand = jest.fn().mockResolvedValue(errorCommand());
+    const { client, reporter } = buildClient(predictControlCommand);
+
+    await expect(client.waitUntil(client.click().button(), 3, 1))
+      .rejects.toBeInstanceOf(ControlCommandError);
+
+    // exactly one step is reported as begun and ended -- not one per failed attempt
+    expect(reporter.onStepBegin).toHaveBeenCalledTimes(1);
+    expect(reporter.onStepEnd).toHaveBeenCalledTimes(1);
+    // 3 tries total => 2 intermediate retries before the final, non-retryable failure
+    expect(reporter.onStepRetry).toHaveBeenCalledTimes(2);
+
+    const [step] = (reporter.onStepEnd as jest.Mock).mock.calls[0] as [Step];
+    expect(step.error).toBeInstanceOf(ControlCommandError);
+    // the final, saved failure carries a screenshot (default withScreenshots: 'onFailure')
+    expect(step.end?.screenshot).toBeDefined();
+  });
+});

packages/askui-nodejs/src/execution/ui-control-client.ts

@@ -25,8 +25,7 @@ import { Instruction, StepReporter } from '../core/reporting';
 import { AIElementCollection } from '../core/ai-element/ai-element-collection';
 import { ModelCompositionBranch } from './model-composition-branch';
 import { AIElementArgs } from '../core/ai-element/ai-elements-args';
-import { NoRetryStrategy, RetryStrategy } from './retry-strategies';
-import { ControlCommandError } from './control-command-error';
+import { FixedRetryStrategy, RetryStrategy } from './retry-strategies';
 import { AskUIAgent, AgentHistory, ActOptions } from '../core/models/anthropic';
 import { AskUIGetAskUIElementTool, AskUIListAIElementTool } from '../core/models/anthropic/tools/askui-api-tools';
 
@@ -509,37 +508,13 @@ export class UiControlClient extends ApiCommands {
    * @param {number} maxTry - Number of maximum retries
    * @param {number} waitTime - Time in milliseconds
    */
-  async waitUntil(AskUICommand: Executable, maxTry = 5, waitTime = 2000) {
-    logger.debug(`waitUntil: Starting with maxTry=${maxTry}, waitTime=${waitTime}ms, retryStrategy=${this.executionRuntime.retryStrategy.constructor.name}`);
-
-    const userDefinedStrategy = this.executionRuntime.retryStrategy;
-    this.executionRuntime.retryStrategy = new NoRetryStrategy();
-
-    const attempt = async (retriesLeft: number): Promise<void> => {
-      const attemptNumber = maxTry - retriesLeft;
-      logger.debug(`waitUntil: Attempt ${attemptNumber}/${maxTry} (${retriesLeft} retries remaining)`);
-      try {
-        await AskUICommand.exec();
-        logger.debug(`waitUntil: Command succeeded on attempt ${attemptNumber}/${maxTry}`);
-        return;
-      } catch (error: unknown) {
-        if (error instanceof ControlCommandError && retriesLeft > 0) {
-          logger.debug(`waitUntil: ControlCommandError on attempt ${attemptNumber}/${maxTry}, waiting ${waitTime}ms before retry.`, error);
-          await this.waitFor(waitTime).exec();
-          await attempt(retriesLeft - 1);
-          return;
-        }
-        const errorName = error instanceof Error ? error.name : 'Error';
-        logger.debug(`waitUntil: ${errorName} on attempt ${attemptNumber}/${maxTry}, no retries remaining.`, error);
-        throw error;
-      }
-    };
-
-    try {
-      await attempt(maxTry - 1);
-    } finally {
-      this.executionRuntime.retryStrategy = userDefinedStrategy;
-    }
+  // eslint-disable-next-line class-methods-use-this
+  async waitUntil(AskUICommand: Executable, maxTry = 5, waitTime = 2000): Promise<void> {
+    logger.debug(`waitUntil: Starting with maxTry=${maxTry}, waitTime=${waitTime}ms`);
+    await AskUICommand.exec({
+      instructionPrefix: 'waitUntil: ',
+      retryStrategy: new FixedRetryStrategy(waitTime, maxTry - 1),
+    });
   }
 
   private evaluateRelation(