fix(desktop): accept press_key without a window and screenshot shorthand on actions

Mirrors the worker contract: press_key may omit windowId (the key goes to the
focused observed window, or a dialog above it), and mutations accept
includeScreenshot as shorthand for observeAfter.includeScreenshot.

Author: waleedlatif1

apps/desktop/native/computer-use/ComputerUse.swift

@@ -681,14 +681,20 @@ final class Driver {
             }
             if request.method == "input_sequence" { return ["kind": "action", "action": request.method, "bundleId": bundle, "dispatched": true, "verified": false, "sequence": ["completedSteps": completed, "totalSteps": steps.count]] }
         case "press_key":
-            let windowId = try required(p.windowId, "windowId"); let frame = try windowRect(windowId, pid: pid)
-            guard snapshot.windowFrames[windowId] == frame else { throw ComputerError("stale_window", "Observe the target window again.") }
+            // Without a windowId the key goes to whichever observed window has focus.
+            let frame = try p.windowId.map { windowId in
+                let frame = try windowRect(windowId, pid: pid)
+                guard snapshot.windowFrames[windowId] == frame else { throw ComputerError("stale_window", "Observe the target window again.") }
+                return frame
+            }
             let root = AXUIElementCreateApplication(pid)
             // Keys go to the app's key window: the observed window, or a dialog, panel or popover above it,
             // but never another document window the snapshot did not observe.
             guard let rawWindow = attribute(root, kAXFocusedWindowAttribute), CFGetTypeID(rawWindow) == AXUIElementGetTypeID() else { throw ComputerError("focus_failed", "The app has no focused window to receive keys; click a window or call activate_app, then observe again.") }
             let focusedWindow = unsafeBitCast(rawWindow, to: AXUIElement.self)
-            guard elementRect(focusedWindow) == frame || attribute(focusedWindow, kAXSubroleAttribute) as? String != kAXStandardWindowSubrole else { throw ComputerError("focus_failed", "Another document window has keyboard focus; observe that window and target it instead.") }
+            let focusedFrame = elementRect(focusedWindow)
+            let observed = frame.map { focusedFrame == $0 } ?? snapshot.windowFrames.values.contains { focusedFrame == $0 }
+            guard observed || attribute(focusedWindow, kAXSubroleAttribute) as? String != kAXStandardWindowSubrole else { throw ComputerError("focus_failed", "Another document window has keyboard focus; observe that window and target it instead.") }
             if let focused = attribute(root, kAXFocusedUIElementAttribute), CFGetTypeID(focused) == AXUIElementGetTypeID() { try requireNonSecure(unsafeBitCast(focused, to: AXUIElement.self)) }
             let (code, flags) = try parseKey(required(p.key, "key"))
             var selectedAll = false

apps/desktop/src/main/computer-use/service.ts

@@ -235,8 +235,17 @@ export class ComputerUseService {
             ? (owner?.elementWindows.get(input.elementId) ?? owner?.windowId)
             : undefined
       if ('snapshotId' in input) this.snapshots.delete(input.snapshotId)
-      const params =
-        'observeAfter' in input ? omit(input, ['action', 'observeAfter']) : omit(input, ['action'])
+      /** Mutations accept includeScreenshot as shorthand for observeAfter.includeScreenshot. */
+      const isMutation = input.action !== 'get_app_state'
+      const observeAfter = {
+        ...('observeAfter' in input ? input.observeAfter : {}),
+        ...(isMutation && 'includeScreenshot' in input && input.includeScreenshot !== undefined
+          ? { includeScreenshot: input.includeScreenshot }
+          : {}),
+      }
+      const params = isMutation
+        ? omit(input, ['action', 'observeAfter', 'includeScreenshot'])
+        : omit(input, ['action'])
       let result: ComputerUseResult
       try {
         result = await this.deps.native.request(input.action, params)
@@ -264,7 +273,7 @@ export class ComputerUseService {
           const observation = await this.deps.native.request('get_app_state', {
             bundleId: input.bundleId,
             ...(windowId ? { windowId } : {}),
-            ...('observeAfter' in input ? input.observeAfter : {}),
+            ...observeAfter,
           })
           this.check(work)
           if (observation.kind !== 'state')

apps/sim/lib/mothership/generated/computer-use.ts

@@ -25,6 +25,7 @@ const ObserveAfter = {
     .describe(
       "Every mutation returns fresh app state by default. Set includeScreenshot to true to include an image; otherwise no screenshot is captured. Observation failure never retries the action.",
     ),
+  includeScreenshot: z.boolean().optional().describe("Same as observeAfter.includeScreenshot."),
 };
 const InputSequenceStep = z.discriminatedUnion("action", [
   z.strictObject({ action: z.literal("type_text"), text: z.string().max(32_000) }),
@@ -75,7 +76,11 @@ export const ComputerUseSchema = z.union([
   z.strictObject({
     action: z.literal("press_key"),
     ...ObserveAfter,
-    ...WindowTarget,
+    bundleId: BundleId,
+    snapshotId: SnapshotId,
+    windowId: WindowId.optional().describe(
+      "Window reference from that snapshot. Omit to send the key to the app's focused window.",
+    ),
     key: z.string().min(1).max(128).describe("One key or chord, such as Enter, Tab, Escape, or Cmd+A."),
   }),
   z.strictObject({ action: z.literal("scroll"), ...ElementTarget, ...ObserveAfter, ...ScrollOffsets }),

packages/desktop-bridge/src/computer-use.generated.ts

@@ -25,6 +25,7 @@ const ObserveAfter = {
     .describe(
       "Every mutation returns fresh app state by default. Set includeScreenshot to true to include an image; otherwise no screenshot is captured. Observation failure never retries the action.",
     ),
+  includeScreenshot: z.boolean().optional().describe("Same as observeAfter.includeScreenshot."),
 };
 const InputSequenceStep = z.discriminatedUnion("action", [
   z.strictObject({ action: z.literal("type_text"), text: z.string().max(32_000) }),
@@ -75,7 +76,11 @@ export const ComputerUseSchema = z.union([
   z.strictObject({
     action: z.literal("press_key"),
     ...ObserveAfter,
-    ...WindowTarget,
+    bundleId: BundleId,
+    snapshotId: SnapshotId,
+    windowId: WindowId.optional().describe(
+      "Window reference from that snapshot. Omit to send the key to the app's focused window.",
+    ),
     key: z.string().min(1).max(128).describe("One key or chord, such as Enter, Tab, Escape, or Cmd+A."),
   }),
   z.strictObject({ action: z.literal("scroll"), ...ElementTarget, ...ObserveAfter, ...ScrollOffsets }),