draw-things-chat

// src/index.ts
import { setGlobalDispatcher, Agent } from "undici";
import { type PluginContext } from "@lmstudio/sdk";

// Undici's default bodyTimeout is 300 s. Long prompt-processing phases exceed
// that before the first streaming token arrives, causing UND_ERR_BODY_TIMEOUT.
// Setting bodyTimeout: 0 removes the limit for all outgoing fetch() calls.
setGlobalDispatcher(new Agent({ bodyTimeout: 0, headersTimeout: 0 }));
import { toolsProvider } from "./toolsProvider.js";
import {
  configSchematics,
  globalConfigSchematics,
  engineConnectionDefaults,
  preprocess,
  getLogsDir,
  getPluginLogFilename,
  getPluginMeta,
} from "./core-bundle.mjs";
import { generate } from "./orchestrator.js";
import fs from "fs";
import path from "path";
import {
  checkVisionPrimerStatus,
  loadVisionPrimerModel,
  injectLastUsedModelIntoNewestChat,
  type VisionPrimerQuickCheck,
} from "./core-bundle.mjs";
import { warmupBackendAtStartup } from "./core/tools.js";

function getPrimerApiConfig(context: PluginContext): { baseUrl?: string; apiKey?: string } {
  try {
    const getter: any =
      (context as any).getGlobalPluginConfig || (context as any).getGlobalConfig;
    const cfg = getter ? getter.call(context, globalConfigSchematics) : null;
    return {
      baseUrl: typeof cfg?.get?.("baseUrl") === "string" ? cfg.get("baseUrl") : undefined,
      apiKey: typeof cfg?.get?.("apiKey") === "string" ? cfg.get("apiKey") : undefined,
    };
  } catch {
    return {};
  }
}

export async function main(context: PluginContext) {
  // Register schematics first so getGlobalConfig() can access defaults
  context
    .withConfigSchematics(configSchematics)
    .withGlobalConfigSchematics(globalConfigSchematics)
    .withPromptPreprocessor(preprocess)
    .withGenerator(generate)
    .withToolsProvider(toolsProvider);

  // Vision Capability Primer (HYBRID: quick checks awaited, load fire-and-forget):
  // 1. Quick checks (CLI, installed, loaded) are AWAITED (fast, ~1-2s)
  // 2. If already loaded → inject lastUsedModel immediately
  // 3. If needs load → fire-and-forget the slow load operation
  // This ensures the injection happens before LM Studio checks capabilities,
  // while not blocking startup on slow model loading.
  const primerConfig = {
    modelKey: "qwen/qwen3-vl-4b",
    ...getPrimerApiConfig(context),
    contextLength: 4096,
    gpuMode: "off" as const,
    ttlSeconds: 7200,
    identifier: "vision-capability-priming",
  };
  const pluginId = getPluginMeta().pluginIdentifier;

  // AWAIT quick checks (fast)
  const quickCheck = await checkVisionPrimerStatus(primerConfig);

  if (quickCheck.alreadyLoaded) {
    // Model already loaded - inject immediately (BEFORE main() returns)
    console.debug("[VisionPrimer] Model already loaded, injecting lastUsedModel immediately");
    await injectLastUsedModelIntoNewestChat({
      modelKey: primerConfig.modelKey,
      identifier: primerConfig.identifier,
      contextLength: primerConfig.contextLength,
      pluginId,
    });
    // Mark as complete for orchestrator
    (globalThis as any).__dtc_visionPrimerResult = { ok: true, alreadyLoaded: true };
  } else if (quickCheck.needsLoad) {
    // Model installed but not loaded - fire-and-forget the load
    console.debug("[VisionPrimer] Model needs loading, starting fire-and-forget load");
    const loadPromise = loadVisionPrimerModel(quickCheck.lmsCli, primerConfig);

    // Store promise for orchestrator
    (globalThis as any).__dtc_visionPrimerPromise = loadPromise;

    // Handle completion asynchronously
    loadPromise
      .then(async (loadResult) => {
        (globalThis as any).__dtc_visionPrimerResult = loadResult;
        if (loadResult.ok) {
          console.debug(
            `[VisionPrimer] Model loaded: ${loadResult.size} in ${loadResult.loadTimeSec}s`
          );
          await injectLastUsedModelIntoNewestChat({
            modelKey: primerConfig.modelKey,
            identifier: primerConfig.identifier,
            contextLength: primerConfig.contextLength,
            pluginId,
          });
        } else {
          console.warn("[VisionPrimer] Load failed:", loadResult.error);
        }
      })
      .catch((err) => {
        console.warn("[VisionPrimer] Unexpected load error:", err?.message || err);
      });
  } else if (quickCheck.userFacingError) {
    // Model not installed or other user-facing error
    console.warn("[VisionPrimer] User-facing error:", quickCheck.error);
    (globalThis as any).__dtc_visionPrimerResult = {
      ok: false,
      notInstalled: quickCheck.notInstalled,
      userFacingError: quickCheck.userFacingError,
      error: quickCheck.error,
    };
  } else if (quickCheck.infrastructureError) {
    // Infrastructure error (silent)
    console.warn("[VisionPrimer] Infrastructure error (silent):", quickCheck.error);
  }

  // Startup file log for diagnostics
  try {
    const cwd = process.cwd();
    const logsDir = getLogsDir();
    const ts = () => {
      try {
        return new Date().toLocaleString(undefined, {
          year: "numeric",
          month: "2-digit",
          day: "2-digit",
          hour: "2-digit",
          minute: "2-digit",
          second: "2-digit",
          hour12: false,
          timeZoneName: "short",
        });
      } catch {
        return new Date().toString();
      }
    };
    if (!fs.existsSync(logsDir)) fs.mkdirSync(logsDir, { recursive: true });
    const lines = [
      `${ts()} - Plugin start`,
      `  cwd=${cwd}`,
      `  node=${process.version} platform=${process.platform} arch=${process.arch}`,
      `  transportDefault=${engineConnectionDefaults.transport}`,
      `  http.baseUrl=${engineConnectionDefaults.http?.baseUrl}`,
      `  grpc.target=${engineConnectionDefaults.grpc?.target}`,
    ];
    fs.appendFileSync(
      path.join(logsDir, getPluginLogFilename()),
      lines.join("\n") + "\n"
    );
  } catch {}

  // Load cached connection config from last run (if available)
  // This allows warmup to use the user's configured host/port instead of defaults
  try {
    const logsDir = getLogsDir();
    const cachePath = path.join(logsDir, "last-connection-config.json");
    if (fs.existsSync(cachePath)) {
      const cached = JSON.parse(fs.readFileSync(cachePath, "utf-8"));
      if (cached.DRAW_THINGS_HOST && !process.env.DRAW_THINGS_HOST) {
        process.env.DRAW_THINGS_HOST = cached.DRAW_THINGS_HOST;
      }
      if (cached.DRAW_THINGS_HTTP_PORT && !process.env.DRAW_THINGS_HTTP_PORT) {
        process.env.DRAW_THINGS_HTTP_PORT = cached.DRAW_THINGS_HTTP_PORT;
      }
      if (cached.DRAW_THINGS_GRPC_PORT && !process.env.DRAW_THINGS_GRPC_PORT) {
        process.env.DRAW_THINGS_GRPC_PORT = cached.DRAW_THINGS_GRPC_PORT;
      }
      console.log(`[Startup] Loaded cached connection config: ${cached.DRAW_THINGS_HOST}:${cached.DRAW_THINGS_HTTP_PORT}`);
    }
  } catch {}

  // Eager backend warmup (no lazy-on-first-tool-call):
  // As soon as the plugin starts, probe gRPC/HTTP reachability, select transport,
  // and (when gRPC is used) log the model/LoRA preflight.
  try {
    if (process.env.LMS_BACKEND_WARMED_UP !== "1") {
      const logsDir = getLogsDir();
      try {
        if (!fs.existsSync(logsDir)) fs.mkdirSync(logsDir, { recursive: true });
        const ts = () => {
          try {
            return new Date().toLocaleString(undefined, {
              year: "numeric",
              month: "2-digit",
              day: "2-digit",
              hour: "2-digit",
              minute: "2-digit",
              second: "2-digit",
              hour12: false,
              timeZoneName: "short",
            });
          } catch {
            return new Date().toString();
          }
        };
        fs.appendFileSync(
          path.join(logsDir, getPluginLogFilename()),
          `${ts()} - Startup warmup(main): begin\n`
        );
      } catch {}

      await warmupBackendAtStartup();
      process.env.LMS_BACKEND_WARMED_UP = "1";

      try {
        const ts = () => {
          try {
            return new Date().toLocaleString(undefined, {
              year: "numeric",
              month: "2-digit",
              day: "2-digit",
              hour: "2-digit",
              minute: "2-digit",
              second: "2-digit",
              hour12: false,
              timeZoneName: "short",
            });
          } catch {
            return new Date().toString();
          }
        };
        fs.appendFileSync(
          path.join(logsDir, getPluginLogFilename()),
          `${ts()} - Startup warmup(main): done\n`
        );
      } catch {}
    }
  } catch (e) {
    try {
      const logsDir = getLogsDir();
      if (!fs.existsSync(logsDir)) fs.mkdirSync(logsDir, { recursive: true });
      const ts = () => {
        try {
          return new Date().toLocaleString(undefined, {
            year: "numeric",
            month: "2-digit",
            day: "2-digit",
            hour: "2-digit",
            minute: "2-digit",
            second: "2-digit",
            hour12: false,
            timeZoneName: "short",
          });
        } catch {
          return new Date().toString();
        }
      };
      fs.appendFileSync(
        path.join(logsDir, getPluginLogFilename()),
        `${ts()} - Startup warmup(main): error: ${String(
          (e as any)?.message || e
        )}\n`
      );
    } catch {}
  }
}

// injectLastUsedModelIntoNewestChat is now provided by draw-things-chat-core (imported from core-bundle.mjs)