llmgw/internal/vscodeext/ext/extension.js

'use strict';

// llmgw for VS Code: the models a gateway serves, offered to VS Code as a
// language-model provider, so its chat and any extension that asks VS Code
// for a model can use them.
//
// The key lives in VS Code's secret storage, one per gateway, and is put
// there one of two ways: the gateway's account page makes one and hands it
// back through a vscode:// link, or the person pastes one.

const vscode = require('vscode');
const crypto = require('crypto');
const lib = require('./lib');

const VENDOR = 'llmgw';
const SIGN_IN_TTL_MS = 15 * 60 * 1000;

/** @param {vscode.ExtensionContext} context */
function activate(context) {
  const changed = new vscode.EventEmitter();
  const version = context.extension.packageJSON.version || '0.0.0';
  const userAgent = `llmgw-vscode/${version} vscode/${vscode.version}`;

  const gateway = () => lib.gatewayBase(vscode.workspace.getConfiguration('llmgw').get('gatewayUrl'));
  const secretName = (base) => `llmgw.key:${base}`;
  const keyFor = async (base) => (base ? context.secrets.get(secretName(base)) : undefined);

  // The sign-in this window started: what the link back must carry to be
  // taken without asking.
  let pending = null;

  async function askGateway() {
    const now = gateway();
    const typed = await vscode.window.showInputBox({
      title: 'llmgw gateway',
      prompt: 'The address of your llmgw gateway, such as https://llm.example.com',
      value: now,
      ignoreFocusOut: true,
      validateInput: (v) => (lib.gatewayBase(v) ? undefined : 'An http:// or https:// address'),
    });
    if (typed === undefined) return '';
    const base = lib.gatewayBase(typed);
    if (base !== now) {
      await vscode.workspace.getConfiguration('llmgw').update('gatewayUrl', base, vscode.ConfigurationTarget.Global);
    }
    return base;
  }

  async function signIn() {
    const base = gateway() || (await askGateway());
    if (!base) return;
    pending = { state: crypto.randomBytes(18).toString('base64url'), base, at: Date.now() };
    const page = vscode.Uri.parse(`${base}/profile?tool=vscode&state=${pending.state}#connect`);
    const opened = await vscode.env.openExternal(page);
    if (!opened) return;
    vscode.window.showInformationMessage(
      `Sign in on ${base} in your browser and choose "Make a key and the configuration". The page sends the key back to VS Code.`,
    );
  }

  async function enterKey() {
    const base = gateway() || (await askGateway());
    if (!base) return;
    const key = await vscode.window.showInputBox({
      title: `llmgw key for ${base}`,
      prompt: 'A key from your account page on the gateway, starting sk-llmgw-',
      password: true,
      ignoreFocusOut: true,
      validateInput: (v) => (v.trim() ? undefined : 'The key'),
    });
    if (!key) return;
    await storeKey(base, key.trim());
  }

  async function storeKey(base, key) {
    // Tried before it is kept: a wrong key or a gateway that is not there is
    // better said now than as an empty model list later.
    try {
      await listModels(base, key);
    } catch (err) {
      vscode.window.showErrorMessage(String(err && err.message ? err.message : err));
      return;
    }
    await context.secrets.store(secretName(base), key);
    changed.fire();
    vscode.window.showInformationMessage(`llmgw: connected to ${base}. Its models are in the chat's model picker.`);
  }

  async function signOut() {
    const base = gateway();
    if (!base) return;
    await context.secrets.delete(secretName(base));
    changed.fire();
    vscode.window.showInformationMessage(
      `llmgw: signed out of ${base}. The key still works until it is revoked on the gateway's account page.`,
    );
  }

  // The account page's "Open in VS Code" link: vscode://llmgw.llmgw/connect
  // with the gateway, the key it just made, and the state this window sent.
  // A link from anywhere else is asked about first: a page that could hand
  // VS Code a key of its own choosing would also choose where code is sent.
  async function handleUri(uri) {
    if (uri.path !== '/connect') return;
    const q = new URLSearchParams(uri.query);
    let base = lib.gatewayBase(q.get('gateway'));
    const key = (q.get('key') || '').trim();
    if (!base || !key) {
      vscode.window.showErrorMessage('llmgw: the link has no gateway or no key in it.');
      return;
    }
    // The same gateway by plain http, behind a proxy that ends TLS without
    // saying so: kept at https, so the key is never sent in the clear.
    const current = gateway();
    if (current && lib.downgrades(current, base)) {
      base = current;
    }
    const expected =
      pending && pending.base === base && q.get('state') === pending.state && Date.now() - pending.at < SIGN_IN_TTL_MS;
    pending = null;
    if (!expected) {
      const yes = 'Connect';
      const choice = await vscode.window.showWarningMessage(
        `Connect VS Code to the llmgw gateway at ${base}?`,
        {
          modal: true,
          detail: 'What you ask its models, and the code they are given, will go to that address. Connect only if you just made this key on its account page.',
        },
        yes,
      );
      if (choice !== yes) return;
    }
    if (base !== gateway()) {
      await vscode.workspace.getConfiguration('llmgw').update('gatewayUrl', base, vscode.ConfigurationTarget.Global);
    }
    await storeKey(base, key);
  }

  async function call(base, key, path, init) {
    const res = await fetch(base + path, {
      ...init,
      headers: {
        Authorization: `Bearer ${key}`,
        'Content-Type': 'application/json',
        'User-Agent': userAgent,
        ...(init && init.headers),
      },
    });
    if (!res.ok) throw new Error(lib.errorMessage(res.status, await res.text()));
    return res;
  }

  async function listModels(base, key) {
    let res;
    try {
      res = await call(base, key, '/v1/models', { method: 'GET', signal: AbortSignal.timeout(15000) });
    } catch (err) {
      if (err && err.message && err.message.startsWith('llmgw:')) throw err;
      throw new Error(`llmgw: could not reach ${base} (${err && err.cause ? err.cause.message || err.cause : err}).`);
    }
    const body = await res.json();
    return Array.isArray(body.data) ? body.data : [];
  }

  const provider = {
    onDidChangeLanguageModelChatInformation: changed.event,

    async provideLanguageModelChatInformation(options) {
      const base = gateway();
      const key = await keyFor(base);
      if (!base || !key) {
        // Offered, not waited for: VS Code is waiting on this list, and a
        // notification can sit unanswered for as long as the person likes.
        // Signing in fires the change event, and VS Code asks again.
        if (!options.silent) {
          vscode.window
            .showInformationMessage('llmgw: sign in to a gateway to use its models.', 'Sign In', 'Enter API Key')
            .then((pick) => {
              if (pick === 'Sign In') return signIn();
              if (pick === 'Enter API Key') return enterKey();
            });
        }
        return [];
      }
      try {
        const show = vscode.workspace.getConfiguration('llmgw').get('showModelsWithoutTools', true);
        return (await listModels(base, key)).map((e) => lib.modelInfo(e, show)).filter(Boolean);
      } catch (err) {
        if (!options.silent) vscode.window.showErrorMessage(String(err.message || err));
        return [];
      }
    },

    async provideLanguageModelChatResponse(model, messages, options, progress, token) {
      const base = gateway();
      const key = await keyFor(base);
      if (!base || !key) throw new Error('llmgw: not signed in. Run "llmgw: Sign In".');
      const abort = new AbortController();
      const sub = token.onCancellationRequested(() => abort.abort());
      try {
        const res = await call(base, key, '/v1/chat/completions', {
          method: 'POST',
          body: JSON.stringify(lib.requestBody(model, messages, options)),
          signal: abort.signal,
        });
        const answer = new lib.Answer();
        for await (const data of lib.sseData(res.body)) {
          const text = answer.add(data);
          if (text) progress.report(new vscode.LanguageModelTextPart(text));
        }
        for (const c of answer.toolCalls()) {
          progress.report(new vscode.LanguageModelToolCallPart(c.id, c.name, c.input));
        }
      } catch (err) {
        if (token.isCancellationRequested) return;
        throw err;
      } finally {
        sub.dispose();
      }
    },

    async provideTokenCount(_model, text) {
      if (typeof text === 'string') return lib.tokenGuess(text);
      const parts = (text && text.content) || [];
      return parts.reduce((n, p) => n + lib.tokenGuess(p && typeof p.value === 'string' ? p.value : JSON.stringify(p)), 0);
    },
  };

  context.subscriptions.push(
    changed,
    vscode.lm.registerLanguageModelChatProvider(VENDOR, provider),
    vscode.window.registerUriHandler({ handleUri }),
    vscode.commands.registerCommand('llmgw.signIn', signIn),
    vscode.commands.registerCommand('llmgw.enterKey', enterKey),
    vscode.commands.registerCommand('llmgw.signOut', signOut),
    vscode.commands.registerCommand('llmgw.refresh', () => changed.fire()),
    vscode.workspace.onDidChangeConfiguration((e) => {
      if (e.affectsConfiguration('llmgw')) changed.fire();
    }),
  );

  // For another extension, or a test, that has a gateway and a key in hand:
  // the same as the account page's link, without the question, since the
  // caller already runs with every right VS Code gives an extension.
  return {
    async connect(gatewayUrl, key) {
      const base = lib.gatewayBase(gatewayUrl);
      if (!base || !key) throw new Error('llmgw: connect needs a gateway address and a key');
      if (base !== gateway()) {
        await vscode.workspace.getConfiguration('llmgw').update('gatewayUrl', base, vscode.ConfigurationTarget.Global);
      }
      await storeKey(base, String(key).trim());
    },
  };
}

function deactivate() {}

module.exports = { activate, deactivate };