@cueai/omni-reader-mcp 1.0.2 → 1.1.1

package/dist/tools.js

@@ -1,32 +1,33 @@
-import { CallToolRequestSchema, } from "@modelcontextprotocol/sdk/types.js";
+import { createHash } from "node:crypto";
+import { CallToolRequestSchema, CancelTaskRequestSchema, ErrorCode, McpError, } from "@modelcontextprotocol/sdk/types.js";
 import { createClientRequestId as newClientRequestId, } from "./cube-client.js";
 import { OmniBridgeError } from "./errors.js";
 import { openAllowedFile, } from "./path-security.js";
 import { NOOP_PROGRESS } from "./progress.js";
-import { SAFETY_INSTRUCTIONS, discardResultSchema, parseDocumentSchema, readResultSchema, } from "./protocol.js";
-function bridgeError(code, message) {
+import { MACHINE_INSTRUCTIONS, cancelParseSchema, discardResultSchema, getParseStatusSchema, parseSchema, readResultSchema, } from "./protocol.js";
+import { discardResultOutputSchema, discardResultToolOutputSchema, parseResultSchema, parseToolOutputSchema, readResultOutputSchema, readResultToolOutputSchema, structuredResult, } from "./result-contract.js";
+import { classifySource } from "./source.js";
+import { TaskRuntime } from "./task-runtime.js";
+function bridgeError(code, message, facts = {}) {
     return new OmniBridgeError({
         code,
+        failureScope: "bridge",
         message,
-        fileUploaded: false,
-        billed: false,
-        contentReleased: false,
-        retryable: false,
+        operationCreated: facts.operationCreated ?? false,
+        fileUploaded: facts.fileUploaded ?? false,
+        parserStarted: facts.parserStarted ?? false,
+        billed: facts.billed ?? false,
+        contentReleased: facts.contentReleased ?? false,
+        retryable: facts.retryable ?? false,
     });
 }
-function errorResult(error) {
-    const stable = error instanceof OmniBridgeError
+function stableError(error) {
+    return error instanceof OmniBridgeError
         ? error
         : bridgeError("BRIDGE_INTERNAL_ERROR", "The local Omni Bridge could not complete the request.");
-    return {
-        isError: true,
-        content: [{ type: "text", text: JSON.stringify(stable.toJSON()) }],
-    };
 }
-function successResult(value) {
-    return {
-        content: [{ type: "text", text: JSON.stringify(value) }],
-    };
+function failed(error) {
+    return { status: "failed", error: stableError(error).toJSON() };
 }
 function progressSink(extra) {
     const progressToken = extra._meta?.progressToken;
@@ -41,41 +42,63 @@ function progressSink(extra) {
         },
     };
 }
-function parseSuccess(local) {
+function completedLocalResult(local) {
+    const dataHandling = {
+        processing_copy: "deleted",
+        temporary_data: "deleted",
+        delivery_result: "deleted_after_ack",
+        original_source: "unchanged",
+        remote_content_retained: false,
+    };
     if (local.kind === "inline") {
-        return successResult({
+        return {
+            status: "completed",
             operation_id: local.operationId,
-            text: local.text,
-        });
+            result: { kind: "inline", text: local.text },
+            data_handling: dataHandling,
+        };
     }
     if (local.nextCursor === undefined) {
-        throw bridgeError("LOCAL_RESULT_INTEGRITY_FAILED", "The local result artifact cursor is missing.");
-    }
-    return successResult({
-        result_id: local.resultId,
-        result_bytes: local.resultBytes,
-        artifact_path: local.artifactPath,
-        expires_at: local.expiresAt,
-        preview: local.preview,
-        next_cursor: local.nextCursor,
-    });
+        throw bridgeError("LOCAL_RESULT_INTEGRITY_FAILED", "The local result artifact cursor is missing.", {
+            operationCreated: true,
+            fileUploaded: true,
+            parserStarted: true,
+            contentReleased: true,
+        });
+    }
+    return {
+        status: "completed",
+        operation_id: local.operationId,
+        result: {
+            kind: "artifact",
+            result_id: local.resultId,
+            result_bytes: local.resultBytes,
+            expires_at: local.expiresAt,
+            preview: local.preview,
+            next_cursor: local.nextCursor,
+        },
+        data_handling: dataHandling,
+        local_result_cache: {
+            expires_at: local.expiresAt,
+            discard_action: "discard_result",
+        },
+    };
 }
-async function parseDocument(args, extra, dependencies) {
+async function parseLocal(args, clientRequestId, signal, progress, dependencies) {
     const openFile = dependencies.openFile ?? openAllowedFile;
-    const opened = await openFile(args.local_path, {
+    const opened = await openFile(args.source, {
         workspace: dependencies.workspace,
         extraRoots: dependencies.extraRoots,
     });
     let primaryError;
     try {
-        const clientRequestId = (dependencies.createClientRequestId ?? newClientRequestId)();
         const granted = await dependencies.cubeClient.createGrant({
             contentLength: opened.size,
             contentType: opened.contentType,
             fileExtension: opened.safeExtension,
-            noStore: args.no_store,
-            output: args.output,
-        }, clientRequestId, extra.signal);
+            noStore: true,
+            output: "markdown",
+        }, clientRequestId, signal);
         const retention = dependencies.artifactStore.createRetention();
         const operation = {
             parseGrant: granted.parseGrant,
@@ -85,12 +108,12 @@ async function parseDocument(args, extra, dependencies) {
             expiresAt: granted.expiresAt,
             openedFile: opened,
             retention,
-            progress: progressSink(extra),
-            signal: extra.signal,
+            progress,
+            signal,
         };
         await dependencies.iiisClient.uploadAndWait(operation);
         const local = retention.result();
-        const result = parseSuccess(local);
+        const result = completedLocalResult(local);
         await dependencies.iiisClient.ack(operation);
         return result;
     }
@@ -104,121 +127,303 @@ async function parseDocument(args, extra, dependencies) {
         }
         catch {
             if (primaryError === undefined) {
-                throw bridgeError("FILE_CLOSE_FAILED", "The local source file could not be closed safely.");
+                throw bridgeError("FILE_CLOSE_FAILED", "The local source file could not be closed safely.", {
+                    operationCreated: true,
+                    fileUploaded: true,
+                    parserStarted: true,
+                    contentReleased: true,
+                });
             }
         }
     }
 }
-function readSuccess(result) {
-    return successResult({
-        result_id: result.resultId,
-        result_bytes: result.resultBytes,
-        expires_at: result.expiresAt,
-        text: result.text,
-        ...(result.nextCursor === undefined ? {} : { next_cursor: result.nextCursor }),
-    });
+async function parseValue(args, signal, progress, dependencies) {
+    try {
+        const source = classifySource(args.source);
+        const clientRequestId = (dependencies.createClientRequestId ?? newClientRequestId)();
+        if (dependencies.operationManager !== undefined) {
+            const sourceHash = `sha256:${createHash("sha256")
+                .update(source.source, "utf8")
+                .digest("hex")}`;
+            let sourceFingerprint;
+            if (source.kind === "local") {
+                const openFile = dependencies.openFile ?? openAllowedFile;
+                try {
+                    const opened = await openFile(source.source, {
+                        workspace: dependencies.workspace,
+                        extraRoots: dependencies.extraRoots,
+                    });
+                    try {
+                        sourceFingerprint = opened.sourceFingerprint;
+                    }
+                    finally {
+                        await opened.close();
+                    }
+                }
+                catch (error) {
+                    if (!(error instanceof OmniBridgeError) || error.code !== "FILE_NOT_FOUND") {
+                        throw error;
+                    }
+                }
+            }
+            return await dependencies.operationManager.submitResult({
+                sourceKind: source.kind,
+                sourceFacts: {
+                    sourceHash,
+                    ...(sourceFingerprint === undefined ? {} : { sourceFingerprint }),
+                },
+                clientRequestId,
+                signal,
+                context: source.kind === "local"
+                    ? { source: source.source, progress }
+                    : { source: source.source },
+            });
+        }
+        if (source.kind === "url") {
+            if (dependencies.remoteClient === undefined) {
+                throw bridgeError("REMOTE_PARSE_UNAVAILABLE", "Remote URL parsing is not available in this Bridge build.", { retryable: true });
+            }
+            return await dependencies.remoteClient.parse(source.source, clientRequestId, signal);
+        }
+        return await parseLocal(args, clientRequestId, signal, progress, dependencies);
+    }
+    catch (error) {
+        return failed(error);
+    }
 }
-function invalidArgumentsResult() {
-    return errorResult(bridgeError("INVALID_TOOL_ARGUMENTS", "The Omni tool arguments are invalid."));
+async function callParse(args, extra, dependencies) {
+    return structuredResult(parseResultSchema, await parseValue(args, extra.signal, progressSink(extra), dependencies));
 }
-async function dispatchTool(name, rawArguments, extra, dependencies) {
-    if (name === "parse_document") {
-        const parsed = parseDocumentSchema.safeParse(rawArguments ?? {});
-        if (!parsed.success)
-            return invalidArgumentsResult();
-        try {
-            return await parseDocument(parsed.data, extra, dependencies);
+async function statusValue(operationId, waitMs, signal, dependencies) {
+    try {
+        if (dependencies.statusOperation !== undefined) {
+            return await dependencies.statusOperation(operationId, waitMs, signal);
         }
-        catch (error) {
-            return errorResult(error);
+        if (dependencies.operationManager !== undefined) {
+            try {
+                return await dependencies.operationManager.statusResult(operationId, waitMs, signal);
+            }
+            catch (error) {
+                if (!(error instanceof OmniBridgeError) || error.code !== "OPERATION_NOT_FOUND") {
+                    throw error;
+                }
+            }
         }
+        if (dependencies.remoteClient !== undefined) {
+            return await dependencies.remoteClient.status(operationId, waitMs, signal);
+        }
+        throw bridgeError("OPERATION_NOT_FOUND", "The requested parse operation is not available.", { operationCreated: true });
     }
-    if (name === "read_result") {
-        const parsed = readResultSchema.safeParse(rawArguments ?? {});
-        if (!parsed.success)
-            return invalidArgumentsResult();
-        try {
-            return readSuccess(await dependencies.artifactStore.read(parsed.data.result_id, parsed.data.cursor, parsed.data.max_bytes));
+    catch (error) {
+        return failed(error);
+    }
+}
+async function callStatus(operationId, waitMs, extra, dependencies) {
+    return structuredResult(parseResultSchema, await statusValue(operationId, waitMs, extra.signal, dependencies));
+}
+async function cancelValue(operationId, signal, dependencies) {
+    try {
+        if (dependencies.cancelOperation !== undefined) {
+            return await dependencies.cancelOperation(operationId, signal);
         }
-        catch (error) {
-            return errorResult(error);
+        if (dependencies.operationManager !== undefined) {
+            try {
+                return await dependencies.operationManager.cancelResult(operationId, signal);
+            }
+            catch (error) {
+                if (!(error instanceof OmniBridgeError) || error.code !== "OPERATION_NOT_FOUND") {
+                    throw error;
+                }
+            }
         }
+        if (dependencies.remoteClient !== undefined) {
+            return await dependencies.remoteClient.cancel(operationId, signal);
+        }
+        throw bridgeError("OPERATION_NOT_FOUND", "The requested parse operation is not available.", { operationCreated: true });
+    }
+    catch (error) {
+        return failed(error);
+    }
+}
+async function callCancel(operationId, extra, dependencies) {
+    return structuredResult(parseResultSchema, await cancelValue(operationId, extra.signal, dependencies));
+}
+async function callReadResult(resultId, cursor, maxBytes, dependencies) {
+    try {
+        const result = await dependencies.artifactStore.read(resultId, cursor, maxBytes);
+        return structuredResult(readResultOutputSchema, {
+            status: "completed",
+            result: {
+                result_id: result.resultId,
+                result_bytes: result.resultBytes,
+                expires_at: result.expiresAt,
+                text: result.text,
+                ...(result.nextCursor === undefined ? {} : { next_cursor: result.nextCursor }),
+            },
+        });
+    }
+    catch (error) {
+        return structuredResult(readResultOutputSchema, failed(error));
+    }
+}
+async function callDiscardResult(resultId, dependencies) {
+    try {
+        return structuredResult(discardResultOutputSchema, {
+            status: "completed",
+            result_id: resultId,
+            discarded: await dependencies.artifactStore.discard(resultId),
+        });
+    }
+    catch (error) {
+        return structuredResult(discardResultOutputSchema, failed(error));
+    }
+}
+function invalidArgumentsResult(name) {
+    const error = failed(bridgeError("INVALID_TOOL_ARGUMENTS", "The Omni tool arguments are invalid."));
+    if (name === "read_result")
+        return structuredResult(readResultOutputSchema, error);
+    if (name === "discard_result")
+        return structuredResult(discardResultOutputSchema, error);
+    return structuredResult(parseResultSchema, error);
+}
+async function dispatchTool(name, rawArguments, extra, dependencies) {
+    if (name === "parse") {
+        const parsed = parseSchema.safeParse(rawArguments ?? {});
+        return parsed.success
+            ? callParse(parsed.data, extra, dependencies)
+            : invalidArgumentsResult(name);
+    }
+    if (name === "get_parse_status") {
+        const parsed = getParseStatusSchema.safeParse(rawArguments ?? {});
+        return parsed.success
+            ? callStatus(parsed.data.operation_id, parsed.data.wait_ms, extra, dependencies)
+            : invalidArgumentsResult(name);
+    }
+    if (name === "cancel_parse") {
+        const parsed = cancelParseSchema.safeParse(rawArguments ?? {});
+        return parsed.success
+            ? callCancel(parsed.data.operation_id, extra, dependencies)
+            : invalidArgumentsResult(name);
+    }
+    if (name === "read_result") {
+        const parsed = readResultSchema.safeParse(rawArguments ?? {});
+        return parsed.success
+            ? callReadResult(parsed.data.result_id, parsed.data.cursor, parsed.data.max_bytes, dependencies)
+            : invalidArgumentsResult(name);
     }
     if (name === "discard_result") {
         const parsed = discardResultSchema.safeParse(rawArguments ?? {});
-        if (!parsed.success)
-            return invalidArgumentsResult();
-        try {
-            return successResult({
-                result_id: parsed.data.result_id,
-                discarded: await dependencies.artifactStore.discard(parsed.data.result_id),
-            });
-        }
-        catch (error) {
-            return errorResult(error);
-        }
+        return parsed.success
+            ? callDiscardResult(parsed.data.result_id, dependencies)
+            : invalidArgumentsResult(name);
     }
-    return errorResult(bridgeError("TOOL_NOT_FOUND", "The requested Omni tool is not available."));
+    return structuredResult(parseResultSchema, failed(bridgeError("TOOL_NOT_FOUND", "The requested Omni tool is not available.")));
 }
-export function registerOmniTools(server, dependencies) {
-    server.registerTool("parse_document", {
-        description: `${SAFETY_INSTRUCTIONS}\n\nParse one local file through Omni without sending source bytes through the model or Cube-MCP.`,
-        inputSchema: parseDocumentSchema,
+export function registerOmniTools(server, dependencies, taskStore) {
+    const taskRuntime = new TaskRuntime({
+        parse: (source, signal) => parseValue({ source }, signal, NOOP_PROGRESS, dependencies),
+        status: (operationId, waitMs, signal) => statusValue(operationId, waitMs, signal, dependencies),
+        cancel: (operationId, signal) => cancelValue(operationId, signal, dependencies),
+    });
+    const createParseTask = async (args, extra) => {
+        const requestedTtl = extra.taskRequestedTtl ?? 600_000;
+        const task = await extra.taskStore.createTask({
+            ttl: Math.min(600_000, Math.max(1_000, requestedTtl)),
+            pollInterval: 1_000,
+        });
+        void taskRuntime.start(task.taskId, args.source, extra.taskStore)
+            .catch(() => undefined);
+        return { task };
+    };
+    const parseInputShape = parseSchema.shape;
+    const parseTaskHandler = {
+        createTask: createParseTask,
+        getTask: async (_args, extra) => extra.taskStore.getTask(extra.taskId),
+        getTaskResult: async (_args, extra) => await extra.taskStore.getTaskResult(extra.taskId),
+    };
+    server.experimental.tasks.registerToolTask("parse", {
+        description: `${MACHINE_INSTRUCTIONS}\n\nParse one local file or HTTP(S) URL through Omni.`,
+        inputSchema: parseInputShape,
+        outputSchema: parseToolOutputSchema,
+        execution: { taskSupport: "optional" },
         annotations: {
             readOnlyHint: false,
             destructiveHint: false,
             idempotentHint: false,
             openWorldHint: true,
         },
-    }, async (args, extra) => {
-        try {
-            return await parseDocument(args, extra, dependencies);
-        }
-        catch (error) {
-            return errorResult(error);
-        }
-    });
+    }, parseTaskHandler);
+    server.registerTool("get_parse_status", {
+        description: "Read the current state of one existing Omni parse operation.",
+        inputSchema: getParseStatusSchema,
+        outputSchema: parseToolOutputSchema,
+        annotations: {
+            readOnlyHint: true,
+            destructiveHint: false,
+            idempotentHint: true,
+            openWorldHint: true,
+        },
+    }, (args, extra) => callStatus(args.operation_id, args.wait_ms, extra, dependencies));
+    server.registerTool("cancel_parse", {
+        description: "Request cancellation of one existing Omni parse operation.",
+        inputSchema: cancelParseSchema,
+        outputSchema: parseToolOutputSchema,
+        annotations: {
+            readOnlyHint: false,
+            destructiveHint: true,
+            idempotentHint: true,
+            openWorldHint: true,
+        },
+    }, (args, extra) => callCancel(args.operation_id, extra, dependencies));
     server.registerTool("read_result", {
         description: "Read one UTF-8 chunk from a Bridge-created local result artifact.",
         inputSchema: readResultSchema,
+        outputSchema: readResultToolOutputSchema,
         annotations: {
             readOnlyHint: true,
             destructiveHint: false,
             idempotentHint: true,
             openWorldHint: false,
         },
-    }, async (args) => {
-        try {
-            return readSuccess(await dependencies.artifactStore.read(args.result_id, args.cursor, args.max_bytes));
-        }
-        catch (error) {
-            return errorResult(error);
-        }
-    });
+    }, (args) => callReadResult(args.result_id, args.cursor, args.max_bytes, dependencies));
     server.registerTool("discard_result", {
         description: "Immediately delete one Bridge-created local result artifact.",
         inputSchema: discardResultSchema,
+        outputSchema: discardResultToolOutputSchema,
         annotations: {
             readOnlyHint: false,
             destructiveHint: true,
             idempotentHint: true,
             openWorldHint: false,
         },
-    }, async (args) => {
-        try {
-            return successResult({
-                result_id: args.result_id,
-                discarded: await dependencies.artifactStore.discard(args.result_id),
-            });
-        }
-        catch (error) {
-            return errorResult(error);
-        }
-    });
+    }, (args) => callDiscardResult(args.result_id, dependencies));
     server.server.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
         if (request.params.task !== undefined) {
-            return errorResult(bridgeError("TOOL_TASK_UNSUPPORTED", "Task-augmented tool calls are not supported by this Omni Bridge."));
+            if (request.params.name !== "parse") {
+                return structuredResult(parseResultSchema, failed(bridgeError("TOOL_TASK_UNSUPPORTED", "Task augmentation is supported only for Omni parse calls.")));
+            }
+            const parsed = parseSchema.safeParse(request.params.arguments ?? {});
+            if (!parsed.success || extra.taskStore === undefined) {
+                return invalidArgumentsResult(request.params.name);
+            }
+            return createParseTask(parsed.data, extra);
         }
         return dispatchTool(request.params.name, request.params.arguments, extra, dependencies);
     });
+    server.server.setRequestHandler(CancelTaskRequestSchema, async (request, extra) => {
+        const task = await taskStore.getTask(request.params.taskId, extra.sessionId);
+        if (task === null) {
+            throw new McpError(ErrorCode.InvalidParams, `Task not found: ${request.params.taskId}`);
+        }
+        if (["completed", "failed", "cancelled"].includes(task.status)) {
+            throw new McpError(ErrorCode.InvalidParams, `Cannot cancel task in terminal status: ${task.status}`);
+        }
+        await taskStore.updateTaskStatus(request.params.taskId, "cancelled", "Client cancelled task execution.", extra.sessionId);
+        await taskRuntime.cancel(request.params.taskId);
+        const canceled = await taskStore.getTask(request.params.taskId, extra.sessionId);
+        if (canceled === null) {
+            throw new McpError(ErrorCode.InvalidParams, `Task not found: ${request.params.taskId}`);
+        }
+        return { _meta: {}, ...canceled };
+    });
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@cueai/omni-reader-mcp",
-  "version": "1.0.2",
+  "version": "1.1.1",
   "description": "Local stdio MCP bridge for direct Omni document parsing",
   "type": "module",
   "license": "UNLICENSED",
@@ -19,12 +19,12 @@
   },
   "scripts": {
     "build": "tsc -p tsconfig.json",
-    "test": "vitest run",
+    "test": "npm run build && vitest run",
     "typecheck": "tsc -p tsconfig.json --noEmit && tsc -p tsconfig.test.json --noEmit",
     "prepack": "npm run test && npm run build"
   },
   "dependencies": {
-    "@modelcontextprotocol/sdk": "^1.24.0",
+    "@modelcontextprotocol/sdk": "1.30.0",
     "zod": "^3.25.0"
   },
   "devDependencies": {