Add options and make parsing more lenient

Author: steipete

src/schemas.ts

@@ -68,13 +68,23 @@ export type GetScriptingTipsInput = z.infer<typeof GetScriptingTipsInputSchema>;
 // AX Query Input Schema
 export const AXQueryInputSchema = z.object({
     command: z.enum(['query', 'perform']).describe('The operation to perform. (Formerly cmd)'),
+
+    // Fields for lenient parsing if locator is flattened
+    app: z.string().optional().describe('Top-level app name (used if locator is a string and app is not specified within a locator object)'),
+    role: z.string().optional().describe('Top-level role (used if locator is a string/flattened)'),
+    match: z.record(z.string()).optional().describe('Top-level match (used if locator is a string/flattened)'),
+
+    locator: z.union([
+        z.object({
+            app: z.string().describe('Bundle ID or display name of the application to query'),
+            role: z.string().describe('Accessibility role to match, e.g., "AXButton", "AXStaticText"'),
+            match: z.record(z.string()).describe('Attributes to match for the element'),
+            navigation_path_hint: z.array(z.string()).optional().describe('Optional path to navigate within the application hierarchy, e.g., ["window[1]", "toolbar[1]"]. (Formerly pathHint)'),
+        }),
+        z.string().describe('Bundle ID or display name of the application to query (used if role/match are provided at top level and this string serves as the app name)')
+    ]).describe('Specifications to find the target element(s). Can be a full locator object or just an app name string (if role/match are top-level).'),
+
     return_all_matches: z.boolean().optional().describe('When true, returns all matching elements rather than just the first match. Default is false. (Formerly multi)'),
-    locator: z.object({
-        app: z.string().describe('Bundle ID or display name of the application to query'),
-        role: z.string().describe('Accessibility role to match, e.g., "AXButton", "AXStaticText"'),
-        match: z.record(z.string()).describe('Attributes to match for the element'),
-        navigation_path_hint: z.array(z.string()).optional().describe('Optional path to navigate within the application hierarchy, e.g., ["window[1]", "toolbar[1]"]. (Formerly pathHint)'),
-    }),
     attributes_to_query: z.array(z.string()).optional().describe('Attributes to query for matched elements. If not provided, common attributes will be included. (Formerly attributes)'),
     required_action_name: z.string().optional().describe('Filter elements to only those supporting this action, e.g., "AXPress". (Formerly requireAction)'),
     action_to_perform: z.string().optional().describe('Only used with command: "perform" - The action to perform on the matched element. (Formerly action)'),
@@ -83,6 +93,18 @@ export const AXQueryInputSchema = z.object({
     ),
     limit: z.number().int().positive().optional().default(500).describe(
       'Maximum number of lines to return in the output. Defaults to 500. Output will be truncated if it exceeds this limit.'
+    ),
+    max_elements: z.number().int().positive().optional().describe(
+      'For return_all_matches: true queries, specifies the maximum number of UI elements to fully process and return. If omitted, a default (e.g., 200) is used internally by the ax binary. Helps control performance for very large result sets.'
+    ),
+    debug_logging: z.boolean().optional().default(false).describe(
+      'If true, enables detailed debug logging from the ax binary, which will be returned as part of the response. Defaults to false.'
+    ),
+    output_format: z.enum(['smart', 'verbose', 'text_content']).optional().default('smart').describe(
+      "Controls the format and verbosity of the attribute output. \n" +
+      "'smart': (Default) Omits empty/placeholder values. Key-value pairs. \n" +
+      "'verbose': Includes all attributes, even empty/placeholders. Key-value pairs. Useful for debugging. \n" +
+      "'text_content': Returns only concatenated text values of common textual attributes (e.g., AXValue, AXTitle, AXDescription). No keys. Ideal for fast text extraction."
     )
 }).refine(
     (data) => {
@@ -93,7 +115,43 @@ export const AXQueryInputSchema = z.object({
         message: "When command is 'perform', an action_to_perform must be provided",
         path: ["action_to_perform"],
     }
-);
+).superRefine((data, ctx) => {
+    if (typeof data.locator === 'string') { // Case 1: locator is a string (app name)
+        if (data.role === undefined) {
+            ctx.addIssue({
+                code: z.ZodIssueCode.custom,
+                message: "If 'locator' is a string (app name), top-level 'role' must be provided.",
+                path: ['role'], // Path refers to the top-level role
+            });
+        }
+        // data.match will default to {} if undefined later in the handler
+        // data.app (top-level) is ignored if data.locator (string) is present, as the locator string *is* the app name.
+    } else { // Case 2: locator is an object
+        // Ensure top-level app, role, match are not present if locator is a full object, to avoid ambiguity.
+        // This is a stricter interpretation. Alternatively, we could prioritize the locator object's fields.
+        if (data.app !== undefined) {
+            ctx.addIssue({
+                code: z.ZodIssueCode.custom,
+                message: "Top-level 'app' should not be provided if 'locator' is a detailed object. Define 'app' inside the 'locator' object.",
+                path: ['app'],
+            });
+        }
+        if (data.role !== undefined) {
+            ctx.addIssue({
+                code: z.ZodIssueCode.custom,
+                message: "Top-level 'role' should not be provided if 'locator' is a detailed object. Define 'role' inside the 'locator' object.",
+                path: ['role'],
+            });
+        }
+        if (data.match !== undefined) {
+            ctx.addIssue({
+                code: z.ZodIssueCode.custom,
+                message: "Top-level 'match' should not be provided if 'locator' is a detailed object. Define 'match' inside the 'locator' object.",
+                path: ['match'],
+            });
+        }
+    }
+});
 
 export type AXQueryInput = z.infer<typeof AXQueryInputSchema>;
 

src/server.ts

@@ -75,18 +75,28 @@ const GetScriptingTipsInputShape = {
 
 const AXQueryInputShape = {
   command: z.enum(['query', 'perform']),
+  // Top-level fields for lenient parsing
+  app: z.string().optional(),
+  role: z.string().optional(),
+  match: z.record(z.string()).optional(),
+
+  locator: z.union([
+      z.object({
+          app: z.string(),
+          role: z.string(),
+          match: z.record(z.string()),
+          navigation_path_hint: z.array(z.string()).optional(),
+      }),
+      z.string()
+  ]),
   return_all_matches: z.boolean().optional(),
-  locator: z.object({
-      app: z.string(),
-      role: z.string(),
-      match: z.record(z.string()),
-      navigation_path_hint: z.array(z.string()).optional(),
-  }),
   attributes_to_query: z.array(z.string()).optional(),
   required_action_name: z.string().optional(),
   action_to_perform: z.string().optional(),
   report_execution_time: z.boolean().optional().default(false),
   limit: z.number().int().positive().optional().default(500),
+  debug_logging: z.boolean().optional().default(false),
+  output_format: z.enum(['smart', 'verbose', 'text_content']).optional().default('smart'),
 } as const;
 
 async function main() {
@@ -445,6 +455,15 @@ This tool exposes the complete macOS accessibility API capabilities, allowing de
 
 *   \`limit\` (integer, optional): Maximum number of lines to return in the output. Defaults to 500. Output will be truncated if it exceeds this limit.
 
+*   \`max_elements\` (integer, optional): For \`return_all_matches: true\` queries, this specifies the maximum number of UI elements the \`ax\` binary will fully process and return attributes for. If omitted, an internal default (e.g., 200) is used. This helps manage performance when querying UIs with a very large number of matching elements (like numerous text fields on a complex web page). This is different from \`limit\`, which truncates the final text output based on lines.
+
+*   \`debug_logging\` (boolean, optional): If true, enables detailed debug logging from the underlying \`ax\` binary. This diagnostic information will be included in the response, which can be helpful for troubleshooting complex queries or unexpected behavior. Defaults to false.
+
+*   \`output_format\` (enum: 'smart' | 'verbose' | 'text_content', optional, default: 'smart'): Controls the format and verbosity of the attribute output from the \`ax\` binary.
+    *   \`'smart'\`: (Default) Optimized for readability. Omits attributes with empty or placeholder values. Returns key-value pairs.
+    *   \`'verbose'\`: Maximum detail. Includes all attributes, even empty/placeholders. Key-value pairs. Best for debugging element properties.
+    *   \`'text_content'\`: Highly compact for text extraction. Returns only concatenated text values of common textual attributes (e.g., AXValue, AXTitle). No keys are returned. Ideal for quickly getting all text from elements; the \`attributes_to_query\` parameter is ignored in this mode.
+
 **Example Queries (Note: key names have changed to snake_case):**
 
 1.  **Find all text elements in the front Safari window:**
@@ -490,17 +509,48 @@ This tool exposes the complete macOS accessibility API capabilities, allowing de
 **Note:** Using this tool requires that the application running this server has the necessary Accessibility permissions in macOS System Settings > Privacy & Security > Accessibility.`,
     AXQueryInputShape,
     async (args: unknown) => {
-      let input: AXQueryInput; // Declare input here to make it accessible in catch
+      let inputFromZod: AXQueryInput; 
       try {
-        input = AXQueryInputSchema.parse(args);
-        logger.info('accessibility_query called with input:', input);
+        inputFromZod = AXQueryInputSchema.parse(args);
+        logger.info('accessibility_query called with raw Zod-parsed input:', inputFromZod);
+
+        // Normalize the input to the canonical structure AXQueryExecutor expects
+        let canonicalInput: AXQueryInput;
+
+        if (typeof inputFromZod.locator === 'string') {
+            logger.debug('Normalizing malformed input (locator is string). Top-level data:', { appLocatorString: inputFromZod.locator, role: inputFromZod.role, match: inputFromZod.match });
+            // Zod superRefine should have already ensured inputFromZod.role is defined.
+            // The top-level inputFromZod.app is ignored here because inputFromZod.locator (the string) is the app.
+            canonicalInput = {
+                // Spread all other fields from inputFromZod first
+                ...inputFromZod,
+                // Then explicitly define the locator object
+                locator: {
+                    app: inputFromZod.locator, // The string locator is the app name
+                    role: inputFromZod.role!,   // Role from top level (assert non-null due to Zod refine)
+                    match: inputFromZod.match || {}, // Match from top level, or default to empty
+                    navigation_path_hint: undefined // No path hint in this malformed case typically
+                },
+                // Nullify the top-level fields that are now part of the canonical locator
+                // to avoid confusion if they were passed, though AXQueryExecutor won't use them.
+                app: undefined,
+                role: undefined,
+                match: undefined
+            };
+        } else {
+            // Well-formed case: locator is an object. Zod superRefine ensures top-level app/role/match are undefined.
+            logger.debug('Input is well-formed (locator is object).');
+            canonicalInput = inputFromZod;
+        }
+
+        // logger.info('accessibility_query using canonical input for executor:', JSON.parse(JSON.stringify(canonicalInput))); // Commented out due to persistent linter issue
 
-        const result = await axQueryExecutor.execute(input);
+        const result = await axQueryExecutor.execute(canonicalInput);
 
         // For cleaner output, especially for multi-element queries, format the response
         let formattedOutput: string;
 
-        if (input.command === 'query' && input.return_all_matches === true) {
+        if (inputFromZod.command === 'query' && inputFromZod.return_all_matches === true) {
           // For multi-element queries, format the results more readably
           if ('elements' in result) {
             formattedOutput = JSON.stringify(result, null, 2);
@@ -515,15 +565,22 @@ This tool exposes the complete macOS accessibility API capabilities, allowing de
         // Apply line limit
         let finalOutputText = formattedOutput;
         const lines = finalOutputText.split('\n');
-        if (input.limit !== undefined && lines.length > input.limit) {
-          finalOutputText = lines.slice(0, input.limit).join('\n');
-          const truncationNotice = `\n\n--- Output truncated to ${input.limit} lines. Original length was ${lines.length} lines. ---`;
+        if (inputFromZod.limit !== undefined && lines.length > inputFromZod.limit) {
+          finalOutputText = lines.slice(0, inputFromZod.limit).join('\n');
+          const truncationNotice = `\n\n--- Output truncated to ${inputFromZod.limit} lines. Original length was ${lines.length} lines. ---`;
           finalOutputText += truncationNotice;
         }
 
         const responseContent: Array<{ type: 'text'; text: string }> = [{ type: 'text', text: finalOutputText }];
 
-        if (input.report_execution_time) {
+        // Add debug logs if they exist in the result
+        if (result.debug_logs && Array.isArray(result.debug_logs) && result.debug_logs.length > 0) {
+          const debugHeader = "\n\n--- AX Binary Debug Logs ---";
+          const logsString = result.debug_logs.join('\n');
+          responseContent.push({ type: 'text', text: `${debugHeader}\n${logsString}` });
+        }
+
+        if (inputFromZod.report_execution_time) {
           const ms = result.execution_time_seconds * 1000;
           let timeMessage = "Script executed in ";
           if (ms < 1) { // Less than 1 millisecond
@@ -545,7 +602,15 @@ This tool exposes the complete macOS accessibility API capabilities, allowing de
       } catch (error: unknown) {
         const err = error as Error;
         logger.error('Error in accessibility_query tool handler', { message: err.message });
-        throw new sdkTypes.McpError(sdkTypes.ErrorCode.InternalError, `Failed to execute accessibility query: ${err.message}`);
+        // If the error object from AXQueryExecutor contains debug_logs, include them
+        let errorMessage = `Failed to execute accessibility query: ${err.message}`;
+        const errorWithLogs = err as (Error & { debug_logs?: string[] }); // Cast here
+        if (errorWithLogs.debug_logs && Array.isArray(errorWithLogs.debug_logs) && errorWithLogs.debug_logs.length > 0) {
+          const debugHeader = "\n\n--- AX Binary Debug Logs (from error) ---";
+          const logsString = errorWithLogs.debug_logs.join('\n');
+          errorMessage += `\n${debugHeader}\n${logsString}`;
+        }
+        throw new sdkTypes.McpError(sdkTypes.ErrorCode.InternalError, errorMessage);
       }
     }
   );