@llmindset/hf-mcp 0.2.39 → 0.2.40

package/src/jobs/jobs-tool.ts

@@ -14,9 +14,10 @@ import {
 	scheduledSuspendCommand,
 	scheduledResumeCommand,
 } from './commands/scheduled.js';
-import { formatCommandHelp } from './schema-help.js';
+import { formatCommandHelp, extractFieldDetails, type AnyZodType } from './schema-help.js';
 import type { ToolResult } from '../types/tool-result.js';
 import { CPU_FLAVORS, GPU_FLAVORS, SPECIALIZED_FLAVORS } from './types.js';
+import { DEFAULT_LOG_WAIT_SECONDS } from './sse-handler.js';
 import type {
 	RunArgs,
 	UvArgs,
@@ -48,10 +49,90 @@ import {
 	scheduledJobArgsSchema,
 } from './types.js';
 
+const OPERATION_NAMES = [
+	'run',
+	'uv',
+	'ps',
+	'logs',
+	'inspect',
+	'cancel',
+	'scheduled run',
+	'scheduled uv',
+	'scheduled ps',
+	'scheduled inspect',
+	'scheduled delete',
+	'scheduled suspend',
+	'scheduled resume',
+] as const;
+
+type OperationName = (typeof OPERATION_NAMES)[number];
+
+const OPERATION_EXAMPLES: Partial<Record<OperationName, string>> = {
+	run: `{
+  "operation": "run",
+  "args": {
+    "image": "python:3.12",
+    "command": ["python", "-c", "print('Hello from HF Jobs!')"],
+    "flavor": "cpu-basic"
+  }
+}`,
+	uv: `{
+  "operation": "uv",
+  "args": {
+    "script": "import random\\nprint(42 + random.randint(1, 5))"
+  }
+}`,
+	ps: `{"operation": "ps"}`,
+	logs: `{
+  "operation": "logs",
+  "args": {"job_id": "your-job-id"}
+}`,
+	inspect: `{
+  "operation": "inspect",
+  "args": {"job_id": "your-job-id"}
+}`,
+	cancel: `{
+  "operation": "cancel",
+  "args": {"job_id": "your-job-id"}
+}`,
+	'scheduled run': `{
+  "operation": "scheduled run",
+  "args": {
+    "schedule": "@hourly",
+    "image": "python:3.12",
+    "command": ["python", "backup.py"]
+  }
+}`,
+	'scheduled uv': `{
+  "operation": "scheduled uv",
+  "args": {
+    "schedule": "0 9 * * 1-5",
+    "script": "import datetime\\nprint('daily check', datetime.datetime.utcnow())"
+  }
+}`,
+	'scheduled ps': `{"operation": "scheduled ps"}`,
+	'scheduled inspect': `{
+  "operation": "scheduled inspect",
+  "args": {"scheduled_job_id": "your-scheduled-job-id"}
+}`,
+	'scheduled delete': `{
+  "operation": "scheduled delete",
+  "args": {"scheduled_job_id": "your-scheduled-job-id"}
+}`,
+	'scheduled suspend': `{
+  "operation": "scheduled suspend",
+  "args": {"scheduled_job_id": "your-scheduled-job-id"}
+}`,
+	'scheduled resume': `{
+  "operation": "scheduled resume",
+  "args": {"scheduled_job_id": "your-scheduled-job-id"}
+}`,
+};
+
 /**
- * Map of command names to their validation schemas
+ * Map of operation names to their validation schemas
  */
-const COMMAND_SCHEMAS = {
+const OPERATION_SCHEMAS: Record<OperationName, z.ZodSchema> = {
 	run: runArgsSchema,
 	uv: uvArgsSchema,
 	ps: psArgsSchema,
@@ -65,9 +146,10 @@ const COMMAND_SCHEMAS = {
 	'scheduled delete': scheduledJobArgsSchema,
 	'scheduled suspend': scheduledJobArgsSchema,
 	'scheduled resume': scheduledJobArgsSchema,
-} as const;
+};
 
 const HELP_FLAG = 'help';
+const operationRequiresArgsCache = new Map<OperationName, boolean>();
 
 const CPU_FLAVOR_LIST = CPU_FLAVORS.join(', ');
 const GPU_FLAVOR_LIST = GPU_FLAVORS.join(', ');
@@ -98,19 +180,83 @@ function removeHelpFlag(args: Record<string, unknown> | undefined): Record<strin
 	return rest;
 }
 
+function isOperationName(value: string): value is OperationName {
+	return (OPERATION_NAMES as readonly string[]).includes(value);
+}
+
+function formatExampleSnippet(operation: OperationName): string | undefined {
+	const example = OPERATION_EXAMPLES[operation];
+	if (!example) {
+		return undefined;
+	}
+
+	return `Call this tool with:\n\`\`\`json\n${example}\n\`\`\``;
+}
+
+function renderExampleSection(title: string, operation: OperationName): string {
+	const snippet = formatExampleSnippet(operation);
+	if (!snippet) {
+		return '';
+	}
+
+	return `### ${title}
+${snippet}
+`;
+}
+
+function operationRequiresArgs(operation: OperationName): boolean {
+	const cached = operationRequiresArgsCache.get(operation);
+	if (cached !== undefined) {
+		return cached;
+	}
+
+	const schema = OPERATION_SCHEMAS[operation];
+	if (!schema) {
+		operationRequiresArgsCache.set(operation, false);
+		return false;
+	}
+
+	const fields = extractFieldDetails(schema as AnyZodType);
+	const requiresArgs = fields.some((field) => !field.isOptional);
+	operationRequiresArgsCache.set(operation, requiresArgs);
+	return requiresArgs;
+}
+
+function formatCommandHelpWithExample(operation: OperationName, schema: z.ZodSchema): string {
+	let help = formatCommandHelp(operation, schema);
+	const exampleSnippet = formatExampleSnippet(operation);
+
+	if (exampleSnippet) {
+		help += `\n\n### Example\n${exampleSnippet}`;
+	}
+
+	return help;
+}
+
+function extractTopLevelArgs(params: { [key: string]: unknown }): Record<string, unknown> {
+	const legacyArgs: Record<string, unknown> = {};
+	for (const [key, value] of Object.entries(params)) {
+		if (key === 'operation' || key === 'args') {
+			continue;
+		}
+		legacyArgs[key] = value;
+	}
+	return legacyArgs;
+}
+
 /**
- * Validate command arguments against a Zod schema
+ * Validate operation arguments against a Zod schema
  * Returns a ToolResult with detailed error message if validation fails
  */
-function validateArgs(
-	schema: z.ZodSchema,
+function validateArgs<T extends z.ZodTypeAny>(
+	schema: T,
 	args: unknown,
-	commandName: string
-): { success: true } | { success: false; errorResult: ToolResult } {
+	operationName: string
+): { success: true; data: z.infer<T> } | { success: false; errorResult: ToolResult } {
 	const result = schema.safeParse(args);
 
 	if (result.success) {
-		return { success: true };
+		return { success: true, data: result.data as z.infer<T> };
 	}
 
 	// Format Zod errors into a helpful message
@@ -127,7 +273,7 @@ function validateArgs(
 		}
 	}
 
-	let errorMessage = `Error: Invalid parameters for '${commandName}'\n\n`;
+	let errorMessage = `Error: Invalid parameters for '${operationName}'\n\n`;
 
 	if (missingFields.length > 0) {
 		errorMessage += `Missing required parameters:\n${missingFields.join('\n')}\n\n`;
@@ -137,7 +283,7 @@ function validateArgs(
 		errorMessage += `Invalid parameters:\n${invalidFields.join('\n')}\n\n`;
 	}
 
-	errorMessage += `Call hf_jobs("${commandName}", {"help": true}) to see valid arguments.`;
+	errorMessage += `Call this tool with {"operation": "${operationName}", "args": {"help": true}} to see valid arguments.`;
 
 	return {
 		success: false,
@@ -178,97 +324,7 @@ Manage compute jobs on Hugging Face infrastructure.
 
 ## Examples
 
-### Run a simple job
-\`\`\`
-hf_jobs("run", {
-  "image": "python:3.12",
-  "command": ["python", "-c", "print('Hello from HF Jobs!')"],
-  "flavor": "cpu-basic"
-})
-\`\`\`
-
-### Use a Hugging Face Space as the image
-\`\`\`
-hf_jobs("run", {
-  "image": "hf.co/spaces/username/spacename",
-  "command": ["python", "app.py"],
-  "flavor": "cpu-basic"
-})
-\`\`\`
-
-### Run multiline Python scripts
-\`\`\`
-hf_jobs("run", {
-  "image": "python:3.12",
-  "command": ["python", "-c", "import sys\\nprint('Line 1')\\nprint('Line 2')"],
-  "flavor": "cpu-basic"
-})
-\`\`\`
-
-### Run a Python Script from a URL with UV
-\`\`\`
-hf_jobs("uv", {
-  "script": "https://huggingface.co/datasets/uv-scripts/classification/blob/main/classify-dataset.py",
-  "with_deps": ["pandas"],
-  "script_args": ["--input", "data.csv"],
-  "flavor": "cpu-basic"
-})
-\`\`\`
-
-### Run an inline Python script with UV
-\`\`\`
-hf_jobs("uv", {
-  "script": "import math\\nprint('area:', math.pi * 4**2)"
-})
-\`\`\`
-
-
-### Run bash/shell commands
-\`\`\`
-hf_jobs("run", {
-  "image": "ubuntu:22.04",
-  "command": ["/bin/sh", "-lc", "apt-get update && apt-get install -y curl"],
-  "flavor": "cpu-basic"
-})
-\`\`\`
-
-### List running jobs
-\`\`\`
-hf_jobs("ps")
-\`\`\`
-
-### Get job logs
-\`\`\`
-hf_jobs("logs", {"job_id": "your-job-id"})
-\`\`\`
-
-### Run with GPU
-\`\`\`
-hf_jobs("run", {
-  "image": "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel",
-  "command": ["python", "train.py"],
-  "flavor": "a10g-small"
-})
-\`\`\`
-
-### Schedule a job
-\`\`\`
-hf_jobs("scheduled run", {
-  "schedule": "@hourly",
-  "image": "python:3.12",
-  "command": ["python", "backup.py"]
-})
-\`\`\`
-
-### Schedule a UV script
-\`\`\`
-hf_jobs("scheduled uv", {
-  "schedule": "0 9 * * 1-5",
-  "script": "https://huggingface.co/datasets/uv-scripts/classification/blob/main/classify-dataset.py",
-  "with_deps": ["pandas"],
-  "script_args": ["--input", "data.csv"]
-})
-\`\`\`
+${renderExampleSection('Run a simple job', 'run')}${renderExampleSection('Run a Python script with UV', 'uv')}
 
 ## Hardware Flavors
 
@@ -291,17 +347,18 @@ ${HARDWARE_FLAVORS_SECTION}
 - UV inline scripts are automatically base64-decoded inside the container; just send the raw script text
 
 ### Show command-specific help
-\`\`\`
-hf_jobs("<command>", {"help": true})
+Call this tool with:
+\`\`\`json
+{"operation": "<operation>", "args": {"help": true}}
 \`\`\`
 
 ## Tips
 
 - The uv-scripts organisation contains examples for common tasks. dataset_search {'author':'uv-scripts'}
-- Jobs default to detached mode, returning after 10 seconds..
+- Jobs default to non-detached mode (tail logs for up to ${DEFAULT_LOG_WAIT_SECONDS}s or until completion). Set \`detach: true\` to return immediately.
 - Prefer array commands to avoid shell parsing surprises
 - To access private Hub assets, include \`secrets: { "HF_TOKEN": "$HF_TOKEN" }\` (or \`${'${HF_TOKEN}'}\`) to inject your auth token.
-- Logs are time-limited (10s max) - check job page for full logs
+- When not detached, logs are time-limited (${DEFAULT_LOG_WAIT_SECONDS}s max or until job completes) - check job page for full logs
 `;
 
 /**
@@ -310,20 +367,15 @@ hf_jobs("<command>", {"help": true})
 export const HF_JOBS_TOOL_CONFIG = {
 	name: 'hf_jobs',
 	description:
-		'Manage HuggingFace compute jobs. Run commands in Docker containers, ' +
-		'execute Python scripts with UV, schedule and monitor jobs, status and logs. ' +
-		'Call hf_jobs with no command for full usage instructions and examples. ' +
-		'Supports CPU and GPU hardware.',
+		'Manage Hugging Face CPU/GPU compute jobs. Run commands in Docker containers, ' +
+		'execute Python scripts with UV. List, schedule and monitor jobs/logs. ' +
+		'Call this tool with no operation for full usage instructions and examples. ',
 	schema: z.object({
-		command: z
-			.string()
+		operation: z
+			.enum(OPERATION_NAMES)
 			.optional()
-			.describe(
-				'Command to execute: "run", "uv", "ps", "logs", "inspect", "cancel", ' +
-					'"scheduled run", "scheduled uv", "scheduled ps", "scheduled inspect", ' +
-					'"scheduled delete", "scheduled suspend", "scheduled resume"'
-			),
-		args: z.record(z.any()).optional().describe('Command-specific arguments as a JSON object'),
+			.describe(`Operation to execute. Valid values: ${OPERATION_NAMES.map((cmd) => `"${cmd}"`).join(', ')}`),
+		args: z.record(z.any()).optional().describe('Operation-specific arguments as a JSON object'),
 	}),
 	annotations: {
 		title: 'Hugging Face Jobs', // omit destructive hint.
@@ -347,9 +399,9 @@ export class HfJobsTool {
 	}
 
 	/**
-	 * Execute a jobs command
+	 * Execute a jobs operation
 	 */
-	async execute(params: { command?: string; args?: Record<string, unknown> }): Promise<ToolResult> {
+	async execute(params: { operation?: string; args?: Record<string, unknown> }): Promise<ToolResult> {
 		// If not authenticated, show upgrade message
 		if (!this.isAuthenticated) {
 			return {
@@ -360,8 +412,10 @@ export class HfJobsTool {
 			};
 		}
 
-		// If no command provided, return usage instructions
-		if (!params.command) {
+		const requestedOperation = params.operation;
+
+		// If no operation provided, return usage instructions
+		if (!requestedOperation) {
 			return {
 				formatted: USAGE_INSTRUCTIONS,
 				totalResults: 1,
@@ -369,101 +423,129 @@ export class HfJobsTool {
 			};
 		}
 
-		const command = params.command.toLowerCase();
-		const rawArgs = params.args || {};
-		const schema = COMMAND_SCHEMAS[command as keyof typeof COMMAND_SCHEMAS];
+		const normalizedOperation = requestedOperation.toLowerCase();
+		if (!isOperationName(normalizedOperation)) {
+			return {
+				formatted: `Unknown operation: "${requestedOperation}"
+Available operations:
+- run, uv, ps, logs, inspect, cancel
+- scheduled run, scheduled uv, scheduled ps, scheduled inspect, scheduled delete, scheduled suspend, scheduled resume
+
+Call this tool with no operation for full usage instructions.`,
+				totalResults: 0,
+				resultsShared: 0,
+			};
+		}
+
+		const operation = normalizedOperation;
+		const legacyArgs = extractTopLevelArgs(params as Record<string, unknown>);
+		const rawArgs = params.args ? params.args : Object.keys(legacyArgs).length > 0 ? legacyArgs : {};
+		const schema = OPERATION_SCHEMAS[operation];
+		const noArgsProvided = !params.args || Object.keys(params.args).length === 0;
+
+		if (schema && noArgsProvided && operationRequiresArgs(operation)) {
+			const helpText = formatCommandHelpWithExample(operation, schema);
+			return {
+				formatted: `No arguments provided for "${operation}".\n\n${helpText}`,
+				totalResults: 1,
+				resultsShared: 1,
+			};
+		}
 		const helpRequested = isHelpRequested(rawArgs);
 
 		if (helpRequested) {
 			if (!schema) {
 				return {
-					formatted: `No help available for '${params.command}'.`,
+					formatted: `No help available for '${requestedOperation}'.`,
 					totalResults: 0,
 					resultsShared: 0,
 				};
 			}
 
+			const helpText = formatCommandHelpWithExample(operation, schema);
 			return {
-				formatted: formatCommandHelp(command, schema),
+				formatted: helpText,
 				totalResults: 1,
 				resultsShared: 1,
 			};
 		}
 
-		const args = removeHelpFlag(rawArgs);
+		const cleanedArgs = removeHelpFlag(rawArgs);
+		let parsedArgs: Record<string, unknown> = cleanedArgs;
 
-		// Validate command arguments if schema exists
+		// Validate operation arguments if schema exists
 		if (schema) {
-			const validation = validateArgs(schema, args, command);
+			const validation = validateArgs(schema, cleanedArgs, operation);
 			if (!validation.success) {
 				return validation.errorResult;
 			}
+			parsedArgs = validation.data as Record<string, unknown>;
 		}
 
 		try {
 			let result: string;
 
-			switch (command) {
+			switch (operation) {
 				case 'run':
-					result = await runCommand(args as RunArgs, this.client, this.hfToken);
+					result = await runCommand(parsedArgs as RunArgs, this.client, this.hfToken);
 					break;
 
 				case 'uv':
-					result = await uvCommand(args as UvArgs, this.client, this.hfToken);
+					result = await uvCommand(parsedArgs as UvArgs, this.client, this.hfToken);
 					break;
 
 				case 'ps':
-					result = await psCommand(args as PsArgs, this.client);
+					result = await psCommand(parsedArgs as PsArgs, this.client);
 					break;
 
 				case 'logs':
-					result = await logsCommand(args as LogsArgs, this.client, this.hfToken);
+					result = await logsCommand(parsedArgs as LogsArgs, this.client, this.hfToken);
 					break;
 
 				case 'inspect':
-					result = await inspectCommand(args as InspectArgs, this.client);
+					result = await inspectCommand(parsedArgs as InspectArgs, this.client);
 					break;
 
 				case 'cancel':
-					result = await cancelCommand(args as CancelArgs, this.client);
+					result = await cancelCommand(parsedArgs as CancelArgs, this.client);
 					break;
 
 				case 'scheduled run':
-					result = await scheduledRunCommand(args as ScheduledRunArgs, this.client, this.hfToken);
+					result = await scheduledRunCommand(parsedArgs as ScheduledRunArgs, this.client, this.hfToken);
 					break;
 
 				case 'scheduled uv':
-					result = await scheduledUvCommand(args as ScheduledUvArgs, this.client, this.hfToken);
+					result = await scheduledUvCommand(parsedArgs as ScheduledUvArgs, this.client, this.hfToken);
 					break;
 
 				case 'scheduled ps':
-					result = await scheduledPsCommand(args as ScheduledPsArgs, this.client);
+					result = await scheduledPsCommand(parsedArgs as ScheduledPsArgs, this.client);
 					break;
 
 				case 'scheduled inspect':
-					result = await scheduledInspectCommand(args as ScheduledJobArgs, this.client);
+					result = await scheduledInspectCommand(parsedArgs as ScheduledJobArgs, this.client);
 					break;
 
 				case 'scheduled delete':
-					result = await scheduledDeleteCommand(args as ScheduledJobArgs, this.client);
+					result = await scheduledDeleteCommand(parsedArgs as ScheduledJobArgs, this.client);
 					break;
 
 				case 'scheduled suspend':
-					result = await scheduledSuspendCommand(args as ScheduledJobArgs, this.client);
+					result = await scheduledSuspendCommand(parsedArgs as ScheduledJobArgs, this.client);
 					break;
 
 				case 'scheduled resume':
-					result = await scheduledResumeCommand(args as ScheduledJobArgs, this.client);
+					result = await scheduledResumeCommand(parsedArgs as ScheduledJobArgs, this.client);
 					break;
 
 				default:
 					return {
-						formatted: `Unknown command: "${params.command}"
-Available commands:
+						formatted: `Unknown operation: "${requestedOperation ?? 'unknown'}"
+Available operations:
 - run, uv, ps, logs, inspect, cancel
 - scheduled run, scheduled uv, scheduled ps, scheduled inspect, scheduled delete, scheduled suspend, scheduled resume
 
-Call hf_jobs() with no arguments for full usage instructions.`,
+Call this tool with no operation for full usage instructions.`,
 						totalResults: 0,
 						resultsShared: 0,
 					};
@@ -493,7 +575,7 @@ Call hf_jobs() with no arguments for full usage instructions.`,
 			}
 
 			return {
-				formatted: `Error executing ${params.command}: ${errorMessage}`,
+				formatted: `Error executing ${requestedOperation ?? 'operation'}: ${errorMessage}`,
 				totalResults: 0,
 				resultsShared: 0,
 				isError: true,

package/src/jobs/schema-help.ts

@@ -5,7 +5,7 @@ import { z } from 'zod';
  * Provides introspection and documentation generation for command arguments
  */
 
-type AnyZodType = z.ZodType<unknown, z.ZodTypeDef, unknown>;
+export type AnyZodType = z.ZodType<unknown, z.ZodTypeDef, unknown>;
 
 export interface FieldDetails {
 	key: string;

package/src/jobs/sse-handler.ts

@@ -1,5 +1,20 @@
 import type { LogEvent } from './types.js';
 
+/**
+ * Default duration to wait for logs when not detached (in milliseconds)
+ */
+export const DEFAULT_LOG_WAIT_MS = 10000; // 10 seconds
+
+/**
+ * Default duration to wait for logs when not detached (in seconds, for display)
+ */
+export const DEFAULT_LOG_WAIT_SECONDS = DEFAULT_LOG_WAIT_MS / 1000;
+
+/**
+ * Default maximum number of log lines to return
+ */
+export const DEFAULT_MAX_LOG_LINES = 20;
+
 /**
  * Options for fetching logs via SSE
  */
@@ -33,7 +48,7 @@ export interface SseLogResult {
  * @returns Log result with collected lines and status
  */
 export async function fetchJobLogs(url: string, options: SseLogOptions = {}): Promise<SseLogResult> {
-	const { maxDuration = 10000, maxLines = 20, token } = options;
+	const { maxDuration = DEFAULT_LOG_WAIT_MS, maxLines = DEFAULT_MAX_LOG_LINES, token } = options;
 
 	const logLines: string[] = [];
 	let finished = false;

package/src/jobs/types.ts

@@ -133,7 +133,11 @@ const commonArgsSchema = z.object({
 
 // Run command args
 export const runArgsSchema = commonArgsSchema.extend({
-	image: z.string().describe('Docker image or HF Space URL (e.g., "python:3.12" or "hf.co/spaces/user/space")'),
+	image: z
+		.string()
+		.describe('Docker image or HF Space URL (e.g., "python:3.12" or "hf.co/spaces/user/space")')
+		.optional()
+		.default('python:3.12'), // NOTE -- this is a deviation from the hf jobs command (which has no default)
 	command: z
 		.union([z.string(), z.array(z.string())])
 		.describe(
@@ -155,8 +159,8 @@ export const runArgsSchema = commonArgsSchema.extend({
 	detach: z
 		.boolean()
 		.optional()
-		.default(true)
-		.describe('Run in background and return after 10 seconds (default: true)'),
+		.default(false)
+		.describe('If true, return immediately with job ID. If false (default), tail logs for up to 10 seconds.'),
 });
 
 // UV command args
@@ -175,7 +179,7 @@ export const uvArgsSchema = commonArgsSchema.extend({
 		.optional()
 		.describe('Secrets as key-value pairs. Use HF_TOKEN=$HF_TOKEN to include your token'),
 	timeout: z.string().optional().default('30m').describe('Max duration'),
-	detach: z.boolean().optional().default(true).describe('Run in background and return after 10 seconds'),
+	detach: z.boolean().optional().default(false).describe('If true, return immediately with job ID. If false (default), tail logs for up to 10 seconds.'),
 });
 
 // PS command args