@cueai/omni-reader-mcp 1.0.2 → 1.1.0

package/dist/result-contract.js

@@ -0,0 +1,235 @@
+import { z } from "zod";
+const operationIdSchema = z.string().min(1).max(128);
+const resultIdSchema = z.string().regex(/^result_[A-Za-z0-9_-]{16,64}$/u);
+const constraintsSchema = z
+    .object({
+    max_bytes: z.number().int().nonnegative().optional(),
+    supported_extensions: z.array(z.string().regex(/^[a-z0-9]{1,10}$/u)).optional(),
+})
+    .strict();
+export const stableErrorSchema = z
+    .object({
+    ok: z.literal(false),
+    code: z.string().min(1).max(128),
+    failure_scope: z
+        .enum([
+        "source",
+        "local_capability",
+        "authentication",
+        "billing",
+        "service",
+        "parser",
+        "operation",
+        "cleanup",
+        "bridge",
+    ])
+        .optional(),
+    source_kind: z.enum(["local", "url"]).optional(),
+    retryable: z.boolean(),
+    message: z.string().min(1).max(2048),
+    user_action: z.string().min(1).max(2048).optional(),
+    request_id: z.string().min(1).max(128).optional(),
+    operation_created: z.boolean(),
+    file_uploaded: z.boolean(),
+    parser_started: z.boolean(),
+    billed: z.boolean(),
+    content_released: z.boolean(),
+    retry_after: z.number().int().nonnegative().optional(),
+    constraints: constraintsSchema.optional(),
+})
+    .strict();
+export const dataHandlingSchema = z
+    .object({
+    processing_copy: z.enum(["in_use", "pending", "deleted"]),
+    temporary_data: z.enum(["in_use", "pending", "deleted"]),
+    delivery_result: z.union([
+        z.literal("not_created"),
+        z.literal("pending"),
+        z.literal("deleted_after_ack"),
+        z.object({ expires_at: z.string().datetime({ offset: true }) }).strict(),
+    ]),
+    original_source: z.literal("unchanged"),
+    remote_content_retained: z.literal(false).optional(),
+})
+    .strict();
+export const progressDetailSchema = z
+    .object({
+    unit: z.enum(["page", "sheet", "slide", "frame", "segment"]),
+    completed: z.number().int().nonnegative(),
+    total: z.number().int().positive(),
+})
+    .strict()
+    .refine((value) => value.completed <= value.total, {
+    message: "completed must not exceed total",
+});
+const inlineResultSchema = z
+    .object({
+    kind: z.literal("inline"),
+    text: z.string(),
+})
+    .strict();
+const artifactResultSchema = z
+    .object({
+    kind: z.literal("artifact"),
+    result_id: resultIdSchema,
+    result_bytes: z.number().int().nonnegative(),
+    expires_at: z.string().datetime({ offset: true }),
+    preview: z.string(),
+    next_cursor: z.string().min(1).max(2048),
+})
+    .strict();
+const billingSchema = z
+    .object({
+    credits_charged: z.number().nonnegative(),
+    credits_remaining: z.number().nonnegative(),
+})
+    .strict();
+const localResultCacheSchema = z
+    .object({
+    expires_at: z.string().datetime({ offset: true }),
+    discard_action: z.literal("discard_result"),
+})
+    .strict();
+const failedResultSchema = z
+    .object({
+    status: z.literal("failed"),
+    error: stableErrorSchema,
+})
+    .strict();
+export const parseResultSchema = z.discriminatedUnion("status", [
+    z
+        .object({
+        status: z.literal("processing"),
+        operation_id: operationIdSchema,
+        stage: z.string().min(1).max(128),
+        progress: z.number().min(0).max(100),
+        progress_detail: progressDetailSchema.optional(),
+        next_action: z.literal("check_status"),
+        poll_after_seconds: z.number().int().positive(),
+        can_cancel: z.boolean(),
+        data_handling: dataHandlingSchema,
+    })
+        .strict(),
+    z
+        .object({
+        status: z.literal("completed"),
+        operation_id: operationIdSchema,
+        result: z.union([inlineResultSchema, artifactResultSchema]),
+        billing: billingSchema.optional(),
+        data_handling: dataHandlingSchema,
+        local_result_cache: localResultCacheSchema.optional(),
+    })
+        .strict(),
+    z
+        .object({
+        status: z.literal("cleanup_pending"),
+        operation_id: operationIdSchema,
+        result: z.union([inlineResultSchema, artifactResultSchema]).optional(),
+        cleanup_deadline: z.string().datetime({ offset: true }),
+        data_handling: dataHandlingSchema,
+    })
+        .strict(),
+    failedResultSchema,
+    z
+        .object({
+        status: z.literal("canceled"),
+        operation_id: operationIdSchema,
+        cleanup_deadline: z.string().datetime({ offset: true }).optional(),
+        data_handling: dataHandlingSchema,
+    })
+        .strict(),
+    z
+        .object({
+        status: z.literal("expired"),
+        operation_id: operationIdSchema,
+        requires_user_confirmation: z.literal(true),
+    })
+        .strict(),
+]);
+const resultChunkSchema = z
+    .object({
+    result_id: resultIdSchema,
+    result_bytes: z.number().int().nonnegative(),
+    expires_at: z.string().datetime({ offset: true }),
+    text: z.string(),
+    next_cursor: z.string().min(1).max(2048).optional(),
+})
+    .strict();
+export const readResultOutputSchema = z.discriminatedUnion("status", [
+    z.object({ status: z.literal("completed"), result: resultChunkSchema }).strict(),
+    failedResultSchema,
+]);
+export const discardResultOutputSchema = z.discriminatedUnion("status", [
+    z
+        .object({
+        status: z.literal("completed"),
+        result_id: resultIdSchema,
+        discarded: z.boolean(),
+    })
+        .strict(),
+    failedResultSchema,
+]);
+export const parseToolOutputSchema = z
+    .object({
+    status: z.enum([
+        "processing",
+        "completed",
+        "cleanup_pending",
+        "failed",
+        "canceled",
+        "expired",
+    ]),
+    operation_id: operationIdSchema.optional(),
+    stage: z.string().min(1).max(128).optional(),
+    progress: z.number().min(0).max(100).optional(),
+    progress_detail: progressDetailSchema.optional(),
+    next_action: z.literal("check_status").optional(),
+    poll_after_seconds: z.number().int().positive().optional(),
+    can_cancel: z.boolean().optional(),
+    result: z.union([inlineResultSchema, artifactResultSchema]).optional(),
+    billing: billingSchema.optional(),
+    data_handling: dataHandlingSchema.optional(),
+    local_result_cache: localResultCacheSchema.optional(),
+    cleanup_deadline: z.string().datetime({ offset: true }).optional(),
+    error: stableErrorSchema.optional(),
+    requires_user_confirmation: z.literal(true).optional(),
+})
+    .strict();
+export const readResultToolOutputSchema = z
+    .object({
+    status: z.enum(["completed", "failed"]),
+    result: resultChunkSchema.optional(),
+    error: stableErrorSchema.optional(),
+})
+    .strict();
+export const discardResultToolOutputSchema = z
+    .object({
+    status: z.enum(["completed", "failed"]),
+    result_id: resultIdSchema.optional(),
+    discarded: z.boolean().optional(),
+    error: stableErrorSchema.optional(),
+})
+    .strict();
+function summarizeForTextFallback(value) {
+    if (value.status === "failed") {
+        const error = value.error;
+        return `Omni failed: ${typeof error?.code === "string" ? error.code : "UNKNOWN"}.`;
+    }
+    if (value.status === "processing")
+        return "Omni processing.";
+    if (value.status === "cleanup_pending")
+        return "Omni cleanup pending.";
+    if (value.status === "canceled")
+        return "Omni canceled.";
+    if (value.status === "expired")
+        return "Omni result expired.";
+    return "Omni completed.";
+}
+export function structuredResult(schema, value) {
+    const structuredContent = schema.parse(value);
+    return {
+        ...(structuredContent.status === "failed" ? { isError: true } : {}),
+        structuredContent,
+        content: [{ type: "text", text: summarizeForTextFallback(structuredContent) }],
+    };
+}

package/dist/server.js

@@ -1,10 +1,27 @@
+import { InMemoryTaskStore } from "@modelcontextprotocol/sdk/experimental/tasks/stores/in-memory.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { SAFETY_INSTRUCTIONS } from "./protocol.js";
+import { BRIDGE_RELEASE_VERSION } from "./constants.js";
+import { MACHINE_INSTRUCTIONS } from "./protocol.js";
 import { registerOmniTools } from "./tools.js";
 const SERVER_NAME = "omni-reader-mcp";
-const SERVER_VERSION = "1.0.2";
 export function createOmniMcpServer(dependencies) {
-    const server = new McpServer({ name: SERVER_NAME, version: SERVER_VERSION }, { instructions: SAFETY_INSTRUCTIONS });
-    registerOmniTools(server, dependencies);
+    const taskStore = new InMemoryTaskStore();
+    const server = new McpServer({ name: SERVER_NAME, version: BRIDGE_RELEASE_VERSION }, {
+        instructions: MACHINE_INSTRUCTIONS,
+        taskStore,
+        capabilities: {
+            tasks: {
+                list: {},
+                cancel: {},
+                requests: { tools: { call: {} } },
+            },
+        },
+    });
+    registerOmniTools(server, dependencies, taskStore);
+    const close = server.close.bind(server);
+    server.close = async () => {
+        taskStore.cleanup();
+        await close();
+    };
     return server;
 }

package/dist/source.d.ts

@@ -0,0 +1,8 @@
+export type ClassifiedSource = {
+    kind: "url";
+    source: string;
+} | {
+    kind: "local";
+    source: string;
+};
+export declare function classifySource(source: string): ClassifiedSource;

package/dist/source.js

@@ -0,0 +1,37 @@
+import { OmniBridgeError } from "./errors.js";
+function unsupportedSource() {
+    return new OmniBridgeError({
+        code: "UNSUPPORTED_SOURCE",
+        failureScope: "source",
+        message: "The source is not supported.",
+        userAction: "Provide an HTTP(S) URL or an ordinary local path.",
+        operationCreated: false,
+        fileUploaded: false,
+        parserStarted: false,
+        billed: false,
+        contentReleased: false,
+        retryable: false,
+    });
+}
+export function classifySource(source) {
+    if (source.includes("\0"))
+        throw unsupportedSource();
+    const scheme = /^([A-Za-z][A-Za-z0-9+.-]*):/u.exec(source)?.[1]?.toLowerCase();
+    if (scheme === undefined)
+        return { kind: "local", source };
+    if (/^[A-Za-z]$/u.test(scheme) && /^[A-Za-z]:[\\/]/u.test(source)) {
+        return { kind: "local", source };
+    }
+    if (scheme !== "http" && scheme !== "https")
+        throw unsupportedSource();
+    let url;
+    try {
+        url = new URL(source);
+    }
+    catch {
+        throw unsupportedSource();
+    }
+    if (url.username !== "" || url.password !== "")
+        throw unsupportedSource();
+    return { kind: "url", source };
+}

package/dist/task-runtime.d.ts

@@ -0,0 +1,13 @@
+import type { RequestTaskStore } from "@modelcontextprotocol/sdk/shared/protocol.js";
+import { type ParseResult } from "./result-contract.js";
+export interface TaskRuntimeOperations {
+    parse(source: string, signal: AbortSignal): Promise<ParseResult>;
+    status(operationId: string, waitMs: number, signal: AbortSignal): Promise<ParseResult>;
+    cancel(operationId: string, signal: AbortSignal): Promise<ParseResult>;
+}
+export declare class TaskRuntime {
+    #private;
+    constructor(operations: TaskRuntimeOperations);
+    start(taskId: string, source: string, store: RequestTaskStore): Promise<void>;
+    cancel(taskId: string): Promise<void>;
+}

package/dist/task-runtime.js

@@ -0,0 +1,94 @@
+import { STATUS_LONG_POLL_MAX_MS } from "./constants.js";
+import { OmniBridgeError } from "./errors.js";
+import { parseResultSchema, structuredResult, } from "./result-contract.js";
+function taskFailure(operationCreated) {
+    const error = new OmniBridgeError({
+        code: "TASK_RUNTIME_FAILED",
+        failureScope: "bridge",
+        message: "The Omni task could not complete.",
+        operationCreated,
+        fileUploaded: false,
+        parserStarted: false,
+        billed: false,
+        contentReleased: false,
+        retryable: true,
+    });
+    return { status: "failed", error: error.toJSON() };
+}
+function taskStatusMessage(result) {
+    return JSON.stringify({
+        operation_id: result.operation_id,
+        stage: result.stage,
+        progress: result.progress,
+        ...(result.progress_detail === undefined
+            ? {}
+            : { progress_detail: result.progress_detail }),
+    });
+}
+async function isCancelled(store, taskId) {
+    try {
+        return (await store.getTask(taskId)).status === "cancelled";
+    }
+    catch {
+        return true;
+    }
+}
+export class TaskRuntime {
+    #operations;
+    #running = new Map();
+    constructor(operations) {
+        this.#operations = operations;
+    }
+    async start(taskId, source, store) {
+        if (this.#running.has(taskId)) {
+            throw new Error(`Task ${taskId} is already running.`);
+        }
+        const running = {
+            controller: new AbortController(),
+        };
+        this.#running.set(taskId, running);
+        try {
+            let result = await this.#operations.parse(source, running.controller.signal);
+            while (result.status === "processing") {
+                running.operationId = result.operation_id;
+                if (await isCancelled(store, taskId)) {
+                    await this.#cancelRunning(running);
+                    return;
+                }
+                await store.updateTaskStatus(taskId, "working", taskStatusMessage(result));
+                result = await this.#operations.status(result.operation_id, STATUS_LONG_POLL_MAX_MS, running.controller.signal);
+            }
+            if (await isCancelled(store, taskId))
+                return;
+            if (result.status === "canceled") {
+                await store.updateTaskStatus(taskId, "cancelled", "Omni parse canceled.");
+                return;
+            }
+            await store.storeTaskResult(taskId, result.status === "failed" ? "failed" : "completed", structuredResult(parseResultSchema, result));
+        }
+        catch {
+            if (!await isCancelled(store, taskId)) {
+                const failed = taskFailure(running.operationId !== undefined);
+                await store.storeTaskResult(taskId, "failed", structuredResult(parseResultSchema, failed));
+            }
+        }
+        finally {
+            this.#running.delete(taskId);
+        }
+    }
+    async cancel(taskId) {
+        const running = this.#running.get(taskId);
+        if (running === undefined)
+            return;
+        await this.#cancelRunning(running);
+    }
+    async #cancelRunning(running) {
+        if (running.cancelRequested === true)
+            return;
+        running.cancelRequested = true;
+        running.controller.abort();
+        if (running.operationId === undefined)
+            return;
+        await this.#operations.cancel(running.operationId, new AbortController().signal);
+    }
+}

package/dist/tools.d.ts

@@ -1,8 +1,11 @@
 import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { TaskStore } from "@modelcontextprotocol/sdk/experimental/tasks/interfaces.js";
 import type { ArtifactReadResult, LocalResult } from "./artifact-store.js";
 import { type CubeGrantClient } from "./cube-client.js";
 import type { IiisClient, ResultRetentionSink } from "./iiis-client.js";
 import { type OpenAllowedFileOptions, type OpenedAllowedFile } from "./path-security.js";
+import type { RemoteOmniClient } from "./remote-client.js";
+import { type ParseResult } from "./result-contract.js";
 interface ToolRetention extends ResultRetentionSink {
     result(): LocalResult;
 }
@@ -11,6 +14,17 @@ interface ToolArtifactStore {
     read(resultId: string, cursor?: string, maxBytes?: number): Promise<ArtifactReadResult>;
     discard(resultId: string): Promise<boolean>;
 }
+export interface ParseOperationController {
+    submitResult(input: {
+        sourceKind: "local" | "url";
+        sourceFacts: Readonly<Record<string, unknown>>;
+        clientRequestId: string;
+        signal: AbortSignal;
+        context: unknown;
+    }): Promise<ParseResult>;
+    statusResult(operationId: string, waitMs: number | undefined, signal: AbortSignal): Promise<ParseResult>;
+    cancelResult(operationId: string, signal: AbortSignal): Promise<ParseResult>;
+}
 export interface OmniToolDependencies {
     readonly workspace: string;
     readonly extraRoots?: readonly string[];
@@ -18,7 +32,11 @@ export interface OmniToolDependencies {
     readonly iiisClient: Pick<IiisClient, "uploadAndWait" | "ack">;
     readonly artifactStore: ToolArtifactStore;
     readonly createClientRequestId?: () => string;
+    readonly remoteClient?: RemoteOmniClient;
+    readonly operationManager?: ParseOperationController;
     readonly openFile?: (localPath: string, options: OpenAllowedFileOptions) => Promise<OpenedAllowedFile>;
+    readonly statusOperation?: (operationId: string, waitMs: number | undefined, signal: AbortSignal) => Promise<ParseResult>;
+    readonly cancelOperation?: (operationId: string, signal: AbortSignal) => Promise<ParseResult>;
 }
-export declare function registerOmniTools(server: McpServer, dependencies: OmniToolDependencies): void;
+export declare function registerOmniTools(server: McpServer, dependencies: OmniToolDependencies, taskStore: TaskStore): void;
 export {};