dsh-codex-subscription 2.0.1 → 2.1.0-beta.1

package/lib/index.js

@@ -1,10 +1,12 @@
 import { clientRequestSchema } from "@deepseek-ai/dsh-client-connection";
+import { createHash, randomBytes, randomUUID } from "node:crypto";
+import { defineTool } from "@deepseek-ai/dsh-tools";
+import { lstat, mkdir, open, readFile, rename, rm, stat } from "node:fs/promises";
 import * as dshCredentials from "@deepseek-ai/dsh-credentials";
 import { dshHomePath, resolveDshHome } from "@deepseek-ai/dsh-home-paths";
 import { LlmError, createUserMessage } from "@deepseek-ai/dsh-llm";
 import { PiAiAdapter } from "@deepseek-ai/dsh-llm-pi-ai";
 import z from "@deepseek-ai/schemastery";
-import { createHash, randomBytes, randomUUID } from "node:crypto";
 import { execFile, spawn } from "node:child_process";
 import { request } from "node:https";
 import { AsyncLocalStorage } from "node:async_hooks";
@@ -14,9 +16,7 @@ import { HttpsProxyAgent } from "https-proxy-agent";
 import { openaiCodexProvider as createOpenAICodexProvider } from "@earendil-works/pi-ai/providers/openai-codex";
 import { createModels } from "@earendil-works/pi-ai";
 import { WebError } from "@deepseek-ai/dsh-web";
-import { defineTool } from "@deepseek-ai/dsh-tools";
 import { constants } from "node:fs";
-import { lstat, mkdir, open, readFile, rename, rm, stat } from "node:fs/promises";
 import { dirname, join, resolve } from "node:path";
 //#region src/image-models.js
 const DEFAULT_IMAGE_MODEL = "gpt-image-2";
@@ -59,7 +59,8 @@ const IMAGE_FEATURE_DEFAULTS = Object.freeze({
 	imageEditing: true,
 	imageViewer: true,
 	imageAnnotations: true,
-	imageSketch: true
+	imageSketch: false,
+	imageSketchAgent: false
 });
 function readImageFeatures(value = {}) {
 	return Object.fromEntries(Object.entries(IMAGE_FEATURE_DEFAULTS).map(([key, fallback]) => [key, typeof value?.[key] === "boolean" ? value[key] : fallback]));
@@ -346,7 +347,11 @@ const RPC_ENDPOINTS = Object.freeze([
 	"reset-credit/inspect",
 	"reset-credit/prepare",
 	"reset-credit/consume",
-	"image/original/chunk"
+	"image/original/chunk",
+	"sketch/connect",
+	"sketch/poll",
+	"sketch/result",
+	"sketch/disconnect"
 ]);
 //#endregion
 //#region src/subscription-transport.js
@@ -401,6 +406,220 @@ function registerSubscriptionTransport(connection, handler) {
 	};
 }
 //#endregion
+//#region src/sketch-agent-bridge.js
+function createSketchAgentBridge({ enabled, now = Date.now, timeoutMs = 2e4 }) {
+	const sessions = /* @__PURE__ */ new Map();
+	const fail = (entry, message) => {
+		for (const task of entry.tasks.values()) task.reject(Error(message));
+		entry.tasks.clear();
+	};
+	const find = (payload) => {
+		const entry = sessions.get(payload.sessionId);
+		if (!entry || entry.token !== payload.token || now() - entry.seen > 1e4) throw Error("Sketch connection expired");
+		entry.seen = now();
+		return entry;
+	};
+	return {
+		async rpc(endpoint, payload) {
+			try {
+				if (!enabled()) throw Error("Sketch is disabled");
+				if (!payload || typeof payload.sessionId !== "string" || !payload.sessionId.length || payload.sessionId.length > 200) throw Error("Invalid session");
+				if (endpoint === "sketch/connect") {
+					for (const [id, entry] of sessions) if (now() - entry.seen >= 1e4) {
+						fail(entry, "Sketch connection expired");
+						sessions.delete(id);
+					}
+					const previous = sessions.get(payload.sessionId);
+					if (previous && now() - previous.seen < 1e4) throw Error("Another board is connected to this session");
+					if (previous) fail(previous, "Sketch connection replaced");
+					const entry = {
+						token: randomUUID(),
+						seen: now(),
+						tasks: /* @__PURE__ */ new Map()
+					};
+					sessions.set(payload.sessionId, entry);
+					return {
+						ok: true,
+						value: { token: entry.token }
+					};
+				}
+				const entry = find(payload);
+				if (endpoint === "sketch/poll") return {
+					ok: true,
+					value: [...entry.tasks].filter(([, t]) => !t.delivered).map(([id, t]) => {
+						t.delivered = true;
+						return {
+							id,
+							request: t.request,
+							expiresAt: t.expiresAt
+						};
+					})
+				};
+				if (endpoint === "sketch/disconnect") {
+					fail(entry, "Sketch board closed");
+					sessions.delete(payload.sessionId);
+					return {
+						ok: true,
+						value: null
+					};
+				}
+				if (endpoint === "sketch/result") {
+					const task = entry.tasks.get(payload.id);
+					if (task) {
+						entry.tasks.delete(payload.id);
+						payload.error ? task.reject(Error(String(payload.error).slice(0, 500))) : task.resolve(payload.value);
+					}
+					return {
+						ok: true,
+						value: null
+					};
+				}
+				throw Error("Unknown sketch route");
+			} catch (error) {
+				return {
+					ok: false,
+					error: {
+						code: "invalid-input",
+						message: error.message,
+						details: { issues: [] }
+					}
+				};
+			}
+		},
+		request(sessionId, request, signal) {
+			if (!enabled()) return Promise.reject(Error("Sketch is disabled"));
+			const entry = sessions.get(sessionId);
+			if (!entry || now() - entry.seen > 1e4) return Promise.reject(Error("Switch to this session in DSH with sketch editing enabled"));
+			if (entry.tasks.size) return Promise.reject(Error("Another sketch operation is pending"));
+			if (JSON.stringify(request).length > 2e6) return Promise.reject(Error("Sketch batch is too large"));
+			return new Promise((resolve, reject) => {
+				const id = randomUUID();
+				const finish = (callback, value) => {
+					clearTimeout(timer);
+					signal?.removeEventListener("abort", abort);
+					entry.tasks.delete(id);
+					callback(value);
+				};
+				const abort = () => finish(reject, Error("Sketch operation interrupted; inspect before retrying"));
+				const timer = setTimeout(() => finish(reject, Error("Sketch response timed out; inspect before retrying")), timeoutMs);
+				entry.tasks.set(id, {
+					request,
+					expiresAt: now() + timeoutMs,
+					delivered: false,
+					resolve: (value) => finish(resolve, value),
+					reject: (error) => finish(reject, error)
+				});
+				signal?.addEventListener("abort", abort, { once: true });
+				if (signal?.aborted) abort();
+			});
+		},
+		dispose() {
+			for (const entry of sessions.values()) fail(entry, "Sketch service stopped");
+			sessions.clear();
+		}
+	};
+}
+//#endregion
+//#region src/sketch-agent-tool.js
+function createSketchAgentTool(bridge, attachments) {
+	return defineTool({
+		name: "codex_sketch",
+		description: "Edit the sketch board in this session using native editable strokes and layers. Only use when asked to draw or edit a sketch. Start with inspect for the documentId, revision and command reference. Apply atomic batches, preview between stages, and save the finished draft. Never generates AI images, sends messages or attaches images automatically. Inspect automatically opens the board in the currently viewed session. Do not ask the user to open it first. If the session is not visible in DSH, ask them to switch to it. On timeout inspect before retrying; reuse the exact requestId only for the same request.",
+		parameters: {
+			action: {
+				type: "string",
+				required: true,
+				enum: [
+					"inspect",
+					"apply",
+					"preview",
+					"save"
+				]
+			},
+			documentId: {
+				type: "string",
+				description: "From inspect; required except for inspect."
+			},
+			revision: {
+				type: "integer",
+				description: "From latest response; required for apply/save."
+			},
+			requestId: {
+				type: "string",
+				description: "Unique id for apply/save; exact retries are deduplicated."
+			},
+			commands: {
+				type: "string",
+				description: "JSON array of native commands described by inspect. Required for apply."
+			},
+			name: {
+				type: "string",
+				description: "Draft name for save."
+			}
+		},
+		timeoutMs: 25e3,
+		isConcurrencySafe: () => false,
+		async execute(args, exec) {
+			const sessionId = exec.agent?.id;
+			if (typeof sessionId !== "string") throw Error("A session-owned sketch call is required");
+			const request = { ...args };
+			if (args.action === "apply") try {
+				request.commands = JSON.parse(args.commands);
+			} catch {
+				throw Error("commands must be a JSON array");
+			}
+			const value = await bridge.request(sessionId, request, exec.signal);
+			if (value.png) {
+				if (!/^data:image\/png;base64,/.test(value.png) || value.png.length > 8 * 1024 * 1024) throw Error("Invalid sketch preview");
+				const image = await attachments.saveImage({
+					data: new Uint8Array(Buffer.from(value.png.split(",")[1], "base64")),
+					mediaType: "image/png",
+					name: "sketch-preview.png"
+				});
+				const { png, ...snapshot } = value;
+				return {
+					...snapshot,
+					image
+				};
+			}
+			return value;
+		},
+		output: {
+			schema: {
+				type: "object",
+				additionalProperties: true
+			},
+			render: (_args, value) => [{
+				type: "text",
+				text: JSON.stringify({
+					...value,
+					image: void 0
+				})
+			}, ...value.image ? [{
+				type: "image",
+				attachment: value.image
+			}] : []]
+		}
+	});
+}
+//#endregion
+//#region src/sketch-codec-route.js
+function registerSketchCodec(connection) {
+	let source;
+	return connection.fetch.register({
+		path: "/api/codex-subscription/sketch-psd-worker",
+		methods: ["GET"],
+		requestBody: "buffered",
+		async fetch() {
+			source ??= await readFile(new URL("./sketch-psd-worker.js", import.meta.url));
+			return new Response(source, { headers: {
+				"content-type": "text/javascript; charset=utf-8",
+				"cache-control": "no-store"
+			} });
+		}
+	});
+}
+//#endregion
 //#region src/account-vault.js
 const VERSION = 1;
 const DEFAULT_LABEL = "Account 1";
@@ -1615,7 +1834,7 @@ function openaiCodexSubscriptionProvider({ resolveSpeedMode = () => void 0, reso
 }
 //#endregion
 //#region src/version.js
-const PACKAGE_VERSION = "2.0.1";
+const PACKAGE_VERSION = "2.1.0-beta.1";
 const USER_AGENT = `dsh-codex-subscription/${PACKAGE_VERSION}`;
 //#endregion
 //#region src/model-catalog.js
@@ -4011,7 +4230,27 @@ function apply(ctx) {
 		usageReader,
 		fetch: (input, init) => network.fetch("quota-reset", input, init)
 	});
-	const handler = createSubscriptionRpcHandler({
+	const sketchBridge = createSketchAgentBridge({ enabled: () => settings.get().imageSketchAgent && settings.get().imageSketch && settings.get().imageEditing });
+	ctx.effect(() => {
+		let dispose;
+		const sync = () => {
+			const value = settings.get();
+			if (value.imageSketchAgent && value.imageSketch && value.imageEditing) dispose ??= ctx.tools.register(createSketchAgentTool(sketchBridge, ctx.attachments));
+			else {
+				dispose?.();
+				dispose = void 0;
+				sketchBridge.dispose();
+			}
+		};
+		sync();
+		const unwatch = settings.watch(sync);
+		return () => {
+			unwatch();
+			dispose?.();
+			sketchBridge.dispose();
+		};
+	}, "codex-subscription: native sketch tool");
+	const subscriptionHandler = createSubscriptionRpcHandler({
 		authHandler: createCodexRpcHandler(coordinator, { openExternal: openCodexAuthUrl }),
 		usageReader,
 		resetCreditService,
@@ -4027,10 +4266,24 @@ function apply(ctx) {
 		originalImages,
 		resolveInheritedOriginal: (sessionId, assetId) => inheritedOriginalImageRef(ctx.get?.("sessions")?.get?.(sessionId), assetId)
 	});
+	const handler = (endpoint, payload, signal) => endpoint.startsWith("sketch/") ? sketchBridge.rpc(endpoint, payload) : subscriptionHandler(endpoint, payload, signal);
 	ctx.effect(() => {
 		modelCatalog.refresh().catch((error) => ctx.logger?.debug?.("could not refresh Codex model catalog: %s", error.message));
 	}, "codex-subscription: official model catalog");
-	ctx.inject(["connection"], (connectionContext) => connectionContext.effect(() => registerSubscriptionTransport(connectionContext.connection, handler), "codex-subscription: DSH-trusted account RPC"));
+	ctx.inject(["connection"], (connectionContext) => connectionContext.effect(() => {
+		const transport = registerSubscriptionTransport(connectionContext.connection, handler);
+		let codec;
+		try {
+			codec = registerSketchCodec(connectionContext.connection);
+		} catch (error) {
+			transport();
+			throw error;
+		}
+		return () => {
+			codec();
+			transport();
+		};
+	}, "codex-subscription: DSH-trusted account RPC"));
 }
 //#endregion
 export { CODEX_IMAGE_GENERATION_URL, CODEX_IMAGE_TOOL_NAME, CODEX_RESET_CONSUME_URL, CODEX_RESET_CREDITS_URL, CODEX_USAGE_URL, CodexLoginCoordinator, DshOAuthCredentialStore, apply, assertCodexAuthUrl, commandForCodexAuthUrl, createCodexAuthService, createCodexImageTool, createCodexResetCreditService, createCodexRpcHandler, createCodexUsageReader, createSearchProviderSwitcher, createSubscriptionDiagnostics, createSubscriptionRpcHandler, decodeCodexPng, inject, name, normalizeContextMode, normalizeCustomContextWindow, openCodexAuthUrl, parseCodexUsage };