Function loadTestsavant

browse/src/security-classifier.ts:204–247

(onProgress?: (msg: string) => void)

let loadPromise: Promise<void> | null = null;

export function loadTestsavant(onProgress?: (msg: string) => void): Promise<void> {
  if (process.env.GSTACK_SECURITY_OFF === '1') {
    testsavantState = 'failed';
    testsavantLoadError = 'GSTACK_SECURITY_OFF=1 — ML classifier kill switch engaged';
    return Promise.resolve();
  }
  if (testsavantState === 'loaded') return Promise.resolve();
  if (loadPromise) return loadPromise;
  testsavantState = 'loading';
  loadPromise = (async () => {
    try {
      await ensureTestsavantStaged(onProgress);
      // Dynamic import — keeps the module boundary clean so static analyzers
      // don't pull @huggingface/transformers into compiled contexts.
      onProgress?.('initializing classifier');
      const { pipeline, env } = await import('@huggingface/transformers');
      env.allowLocalModels = true;
      env.allowRemoteModels = false;
      env.localModelPath = MODELS_DIR;
      testsavantClassifier = await pipeline(
        'text-classification',
        'testsavant-small',
        { dtype: 'fp32' },
      );
      // TestSavantAI's tokenizer_config.json ships with model_max_length
      // set to a huge placeholder (1e18) which disables automatic truncation
      // in the TextClassificationPipeline. The underlying BERT-small has
      // max_position_embeddings: 512 — passing anything longer throws a
      // broadcast error. Override via _tokenizerConfig (the internal source
      // the computed model_max_length getter reads from) so the pipeline's
      // implicit truncation: true actually kicks in.
      const tok = testsavantClassifier?.tokenizer as any;
      if (tok?._tokenizerConfig) {
        tok._tokenizerConfig.model_max_length = 512;
      }
      testsavantState = 'loaded';
    } catch (err: any) {
      testsavantState = 'failed';
      testsavantLoadError = err?.message ?? String(err);
      console.error('[security-classifier] Failed to load TestSavantAI:', testsavantLoadError);
    }
  })();
  return loadPromise;
}

/**
 * Scan text content for prompt injection. Intended for page snapshots, tool