draw-things-chat

// src/index.ts
import { type PluginContext } from "@lmstudio/sdk";
import { toolsProvider } from "./toolsProvider.js";
import {
  configSchematics,
  globalConfigSchematics,
  engineConnectionDefaults,
  preprocess,
  getLogsDir,
  getActiveChatContext,
  resolveActiveLMStudioChatId,
  findLMStudioHome,
} from "./core-bundle.mjs";
import { generate } from "./orchestrator.js";
import fs from "fs";
import path from "path";
import os from "os";
import { checkVisionPrimerStatus, loadVisionPrimerModel, type VisionPrimerQuickCheck } from "./core-bundle.mjs";
import { warmupBackendAtStartup } from "./core/tools.js";

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",
    contextLength: 4096,
    gpuMode: "off" as const,
    ttlSeconds: 7200,
    identifier: "vision-capability-priming",
  };

  // 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(
      primerConfig.modelKey,
      primerConfig.identifier
    );
    // Mark as complete for orchestrator
    (globalThis as any).__dtc_visionPrimerResult = { ok: true, alreadyLoaded: true };
  } else if (quickCheck.needsLoad && quickCheck.lmsCli) {
    // 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(
            primerConfig.modelKey,
            primerConfig.identifier
          );
        } 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, "generate-image-plugin.log"),
      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, "generate-image-plugin.log"),
          `${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, "generate-image-plugin.log"),
          `${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, "generate-image-plugin.log"),
        `${ts()} - Startup warmup(main): error: ${String(
          (e as any)?.message || e
        )}\n`
      );
    } catch {}
  }
}

/**
 * Inject lastUsedModel into the newest conversation.json file, but ONLY if:
 * 1. We (ceveyne/draw-things-chat) are in the plugins array
 * 2. There is no lastUsedModel entry yet
 *
 * @param modelKey - The model key (e.g., "qwen/qwen3-vl-4b") for indexedModelIdentifier
 * @param identifier - The runtime identifier (e.g., "vision-capability-priming") for API access
 */
async function injectLastUsedModelIntoNewestChat(
  modelKey: string,
  identifier?: string
): Promise<void> {
  const startTs = Date.now();

  // DEBUG: Check what context is available at plugin start
  console.debug("[VisionPrimer] === Plugin Start Debug ===");

  const existingContext = getActiveChatContext({ maxAgeMs: 10 * 60 * 1000 });
  if (existingContext) {
    const ageMs = Date.now() - existingContext.setAtMs;
    console.debug("[VisionPrimer] ✓ Active context available:", {
      chatId: existingContext.chatId,
      workingDir: existingContext.workingDir,
      ageMs: ageMs,
      ageSec: Math.round(ageMs / 1000),
      requestId: existingContext.requestId,
    });
  } else {
    console.debug(
      "[VisionPrimer] ✗ No active context available at plugin start"
    );
  }

  // DEBUG: Try resolving via heuristic
  try {
    const resolved = await resolveActiveLMStudioChatId({
      requireRecentMtimeSec: 600,
      retries: 2,
      delayMs: 100,
    });
    if (resolved.ok) {
      console.debug("[VisionPrimer] ✓ Heuristic resolver found:", {
        chatId: resolved.chatId,
        confidence: resolved.confidence,
        reason: resolved.reason,
        filePath: resolved.filePath,
        mtimeMs: resolved.mtimeMs,
        mtimeAgeSec: Math.round((Date.now() - resolved.mtimeMs) / 1000),
      });
    } else {
      console.debug(
        "[VisionPrimer] ✗ Heuristic resolver failed:",
        resolved.reason
      );
    }
  } catch (e: any) {
    console.warn("[VisionPrimer] ✗ Heuristic resolver error:", e.message || e);
  }

  console.debug("[VisionPrimer] === Starting Injection Logic ===");

  try {
    const conversationsDir = path.join(findLMStudioHome(), "conversations");
    console.debug("[VisionPrimer] Conversations dir:", conversationsDir);

    const dirExists = await fs.promises
      .stat(conversationsDir)
      .then(() => true)
      .catch(() => false);
    if (!dirExists) {
      console.debug(
        "[VisionPrimer] Conversations directory not found, skipping"
      );
      return;
    }

    const files = await fs.promises.readdir(conversationsDir);
    const conversationFiles = files.filter((f) =>
      f.endsWith(".conversation.json")
    );

    console.debug(
      `[VisionPrimer] Found ${conversationFiles.length} conversation files`
    );

    if (conversationFiles.length === 0) {
      console.debug("[VisionPrimer] No conversation files found");
      return;
    }

    // Find newest conversation by creation time
    let newestFile: string | null = null;
    let newestBirthtime = 0;

    for (const file of conversationFiles) {
      const filePath = path.join(conversationsDir, file);
      try {
        const stats = await fs.promises.stat(filePath);
        if (stats.birthtimeMs > newestBirthtime) {
          newestBirthtime = stats.birthtimeMs;
          newestFile = file;
        }
      } catch {
        // Skip files we can't stat
      }
    }

    if (!newestFile) {
      console.debug("[VisionPrimer] Could not determine newest conversation");
      return;
    }

    const fileAgeSec = Math.round((Date.now() - newestBirthtime) / 1000);
    console.debug(
      `[VisionPrimer] Newest file: ${newestFile} (created ${fileAgeSec}s ago)`
    );

    const filePath = path.join(conversationsDir, newestFile);
    const raw = await fs.promises.readFile(filePath, "utf8");
    const conversation = JSON.parse(raw);

    // Check if we're in the plugins array
    const ourPluginId = "ceveyne/draw-things-chat";
    const isActive =
      Array.isArray(conversation.plugins) &&
      conversation.plugins.includes(ourPluginId);

    console.debug(`[VisionPrimer] Plugin active in ${newestFile}:`, isActive);
    if (conversation.plugins) {
      console.debug(`[VisionPrimer] Plugins in file:`, conversation.plugins);
    }

    if (!isActive) {
      console.debug(
        `[VisionPrimer] We're not active in newest chat (${newestFile}), skipping`
      );
      return;
    }

    // Check if lastUsedModel already exists
    if (conversation.lastUsedModel) {
      console.debug(
        `[VisionPrimer] lastUsedModel already exists in ${newestFile}:`,
        conversation.lastUsedModel.identifier || conversation.lastUsedModel
      );
      return;
    }

    // Inject lastUsedModel
    // Use identifier if provided (to reference already-loaded model), else fallback to modelKey
    const runtimeIdentifier = identifier ?? modelKey;
    conversation.lastUsedModel = {
      indexedModelIdentifier: modelKey,
      identifier: runtimeIdentifier,
      instanceLoadTimeConfig: { fields: [] },
      instanceOperationTimeConfig: { fields: [] },
    };

    await fs.promises.writeFile(
      filePath,
      JSON.stringify(conversation, null, 2),
      "utf8"
    );

    // First touch to ensure LM Studio's file watcher picks up the change
    const now = new Date();
    await fs.promises.utimes(filePath, now, now);

    const elapsedMs = Date.now() - startTs;
    console.debug(
      `[VisionPrimer] ✓ Injected lastUsedModel (model=${modelKey}, id=${runtimeIdentifier}) into ${newestFile} (${elapsedMs}ms)`
    );
    console.debug(`[VisionPrimer] ✓ First touch to ${newestFile}`);

    // Second touch after 2 seconds to ensure UI update
    setTimeout(async () => {
      try {
        const now2 = new Date();
        await fs.promises.utimes(filePath, now2, now2);
        console.debug(`[VisionPrimer] ✓ Second touch to ${newestFile}`);
      } catch (e: any) {
        console.warn("[VisionPrimer] Failed second touch:", e.message || e);
      }
    }, 2000);
  } catch (e: any) {
    console.warn(
      "[VisionPrimer] Failed to process newest chat:",
      e.message || e
    );
  }
}