fix: preserve selector get and screenshot cleanup

Author: thymikee

src/__tests__/runtime-public.test.ts

@@ -113,6 +113,35 @@ test('runtime screenshot command reserves output and calls backend primitive', a
   });
 });
 
+test('runtime screenshot command cleans reserved output when publish fails', async () => {
+  let cleanupCalled = false;
+  const device = createAgentDevice({
+    backend,
+    artifacts: {
+      ...artifacts,
+      reserveOutput: async (ref: FileOutputRef | undefined, options) => ({
+        path: ref?.kind === 'path' ? ref.path : `/tmp/${options.field}${options.ext}`,
+        visibility: options.visibility ?? 'client-visible',
+        publish: async () => {
+          throw new Error('publish failed');
+        },
+        cleanup: async () => {
+          cleanupCalled = true;
+        },
+      }),
+    },
+    sessions,
+    policy: localCommandPolicy(),
+  });
+
+  await assert.rejects(
+    () => device.capture.screenshot({ out: { kind: 'path', path: '/tmp/screen.png' } }),
+    /publish failed/,
+  );
+
+  assert.equal(cleanupCalled, true);
+});
+
 test('public runtime policy helpers expose local and restricted defaults', async () => {
   assert.equal(typeof createLocalArtifactAdapter, 'function');
   assert.equal(rootCommands.capture.screenshot, commands.capture.screenshot);

src/__tests__/runtime-selector-read.test.ts

@@ -35,6 +35,34 @@ test('runtime get reads text from a selector target', async () => {
   ]);
 });
 
+test('runtime get selector target captures fresh snapshot without a stored session snapshot', async () => {
+  const snapshot = selectorSnapshot();
+  const sessions = createMemorySessionStore([{ name: 'default' }]);
+  let captures = 0;
+  const device = createAgentDevice({
+    backend: {
+      platform: 'ios',
+      captureSnapshot: async () => {
+        captures += 1;
+        return { snapshot };
+      },
+      readText: async () => ({ text: 'Fresh text' }),
+    } satisfies AgentDeviceBackend,
+    artifacts: createLocalArtifactAdapter(),
+    sessions,
+    policy: localCommandPolicy(),
+  });
+
+  const result = await device.selectors.getText(selector('label=Continue'), {
+    session: 'default',
+  });
+
+  assert.equal(result.kind, 'text');
+  assert.equal(result.text, 'Fresh text');
+  assert.equal(captures, 1);
+  assert.equal((await sessions.get('default'))?.snapshot?.nodes[0]?.label, 'Continue');
+});
+
 test('runtime get returns attrs for a ref target without recapturing', async () => {
   const snapshot = selectorSnapshot();
   let captures = 0;

src/commands/capture-screenshot.ts

@@ -23,6 +23,7 @@ export const screenshotCommand: RuntimeCommand<
     ext: '.png',
   });
 
+  let artifact: ArtifactDescriptor | undefined;
   try {
     await runtime.backend.captureScreenshot(
       {
@@ -40,12 +41,12 @@ export const screenshotCommand: RuntimeCommand<
         surface: options.surface,
       },
     );
+    artifact = await reserved.publish();
   } catch (error) {
     await reserved.cleanup?.();
     throw error;
   }
 
-  const artifact = await reserved.publish();
   return {
     path: reserved.path,
     ...(artifact ? { artifacts: [artifact] } : {}),

src/commands/selector-read.ts

@@ -14,6 +14,7 @@ import { buildSelectorChainForNode } from '../utils/selector-build.ts';
 import { evaluateIsPredicate, isSupportedPredicate } from '../utils/selector-is-predicates.ts';
 import type { RuntimeCommand } from './index.ts';
 import {
+  type CapturedSnapshot,
   captureSelectorSnapshot,
   readText,
   requireSnapshotSession,
@@ -177,8 +178,8 @@ export const getCommand: RuntimeCommand<GetCommandOptions, GetCommandResult> = a
   runtime,
   options,
 ): Promise<GetCommandResult> => {
-  const capture = await requireSnapshotSession(runtime, options.session);
   if (options.target.kind === 'ref') {
+    const capture = await requireSnapshotSession(runtime, options.session);
     const resolved = resolveRefNode(capture.snapshot.nodes, options.target.ref, {
       fallbackLabel: options.target.fallbackLabel ?? '',
       invalidRefMessage: 'get text requires a ref like @e2',
@@ -195,7 +196,7 @@ export const getCommand: RuntimeCommand<GetCommandOptions, GetCommandResult> = a
     return { kind: 'text', target, text, node: resolved.node, selectorChain };
   }
 
-  const resolved = await resolveSelectorNode(runtime, options, capture.sessionName, {
+  const resolved = await resolveSelectorNode(runtime, options, options.session ?? 'default', {
     selector: options.target.selector,
     disambiguateAmbiguous: options.property === 'text',
   });
@@ -213,7 +214,7 @@ export const getCommand: RuntimeCommand<GetCommandOptions, GetCommandResult> = a
     };
   }
 
-  const text = await readText(runtime, capture, resolved.node);
+  const text = await readText(runtime, resolved.capture, resolved.node);
   return {
     kind: 'text',
     target: { kind: 'selector', selector: resolved.selector },
@@ -455,7 +456,7 @@ async function resolveSelectorNode(
   options: GetCommandOptions,
   sessionName: string,
   params: { selector: string; disambiguateAmbiguous: boolean },
-): Promise<{ node: SnapshotNode; selector: string; ref: string }> {
+): Promise<{ capture: CapturedSnapshot; node: SnapshotNode; selector: string; ref: string }> {
   const capture = await captureSelectorSnapshot(
     runtime,
     { ...options, session: sessionName },
@@ -473,5 +474,10 @@ async function resolveSelectorNode(
   if (!resolved) {
     throw new AppError('COMMAND_FAILED', formatSelectorFailure(chain, [], { unique: true }));
   }
-  return { node: resolved.node, selector: resolved.selector.raw, ref: `@${resolved.node.ref}` };
+  return {
+    capture,
+    node: resolved.node,
+    selector: resolved.selector.raw,
+    ref: `@${resolved.node.ref}`,
+  };
 }

src/daemon/__tests__/request-router-screenshot.test.ts

@@ -10,6 +10,7 @@ vi.mock('../../core/dispatch.ts', async (importOriginal) => {
 
 import { dispatchCommand } from '../../core/dispatch.ts';
 import { createRequestHandler } from '../request-router.ts';
+import { dispatchScreenshotViaRuntime } from '../screenshot-runtime.ts';
 import type { SessionState } from '../types.ts';
 import { LeaseRegistry } from '../lease-registry.ts';
 import { attachRefs } from '../../utils/snapshot.ts';
@@ -94,6 +95,28 @@ test('screenshot resolves relative positional path against request cwd', async (
   expect(recordedAction?.positionals).toEqual([path.join(callerCwd, 'evidence/test.png')]);
 });
 
+test('default screenshot temp directory is cleaned when capture fails', async () => {
+  const session = makeSession('default');
+  let capturedPath: string | undefined;
+  mockDispatch.mockImplementation(async (_device, command, positionals) => {
+    if (command === 'screenshot') capturedPath = positionals[0];
+    throw new Error('capture failed');
+  });
+
+  await expect(
+    dispatchScreenshotViaRuntime({
+      session,
+      sessionName: session.name,
+      outputPlacement: 'default',
+      dispatchContext: {},
+    }),
+  ).rejects.toThrow(/capture failed/);
+
+  expect(capturedPath).toBeTruthy();
+  expect(path.basename(capturedPath!)).toBe('screenshot.png');
+  expect(fs.existsSync(path.dirname(capturedPath!))).toBe(false);
+});
+
 test('router serializes concurrent commands for the same device across sessions', async () => {
   const sessionStore = makeSessionStore('agent-device-router-screenshot-');
   sessionStore.set('session-a', makeSession('session-a'));

src/daemon/screenshot-runtime.ts

@@ -90,18 +90,26 @@ function createDaemonScreenshotArtifactAdapter(): ArtifactAdapter {
       throw new AppError('UNSUPPORTED_OPERATION', 'screenshot does not resolve input artifacts');
     },
     reserveOutput: async (ref) => {
-      const outputPath =
-        ref?.kind === 'path'
-          ? ref.path
-          : path.join(
-              await fs.mkdtemp(path.join(os.tmpdir(), 'agent-device-screenshot-')),
-              'screenshot.png',
-            );
+      let tempRoot: string | undefined;
+      let outputPath: string;
+      if (ref?.kind === 'path') {
+        outputPath = ref.path;
+      } else {
+        tempRoot = await fs.mkdtemp(path.join(os.tmpdir(), 'agent-device-screenshot-'));
+        outputPath = path.join(tempRoot, 'screenshot.png');
+      }
       await fs.mkdir(path.dirname(outputPath), { recursive: true });
       return {
         path: outputPath,
         visibility: 'client-visible',
         publish: async () => undefined,
+        ...(tempRoot
+          ? {
+              cleanup: async () => {
+                await fs.rm(tempRoot, { recursive: true, force: true });
+              },
+            }
+          : {}),
       };
     },
     createTempFile: async (options) => {