chrome-devtools-frontend/front_end/panels/ai_assistance/components/MarkdownRendererWithCodeBlock.ts

Chrome DevTools UI

// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

import * as Common from '../../../core/common/common.js';
import * as Platform from '../../../core/platform/platform.js';
import type * as SDK from '../../../core/sdk/sdk.js';
import * as AiAssistanceModel from '../../../models/ai_assistance/ai_assistance.js';
import * as Logs from '../../../models/logs/logs.js';
import type * as Marked from '../../../third_party/marked/marked.js';
import * as MarkdownView from '../../../ui/components/markdown_view/markdown_view.js';
import * as Lit from '../../../ui/lit/lit.js';

const {html} = Lit;

export interface MarkdownRendererWithCodeBlockOptions {
  /**
   * Retrieves the established origin locked for the active conversation.
   * Required to authorize #file-<id> links and prevent cross-origin file leakage.
   * When omitted or unauthorized, file links safely render as plain inert text.
   */
  getEstablishedOrigin?: () => SDK.SecurityOrigin.SecurityOrigin | undefined;
}

/**
 * Markdown renderer for AI assistance conversations.
 *
 * The model returns multiline code blocks in an erroneous way with the language being in new line.
 * This renderer takes that into account and correctly updates the parsed multiline token with the language
 * correctly identified and stripped from the content.
 * Example:
 * ```
 * css <-- This should have been on the first line.
 * * {
 * color: red;
 * }
 * ```
 *
 * Also handles linkifying DevTools internal resource links (#req- and origin-locked #file-).
 */
export class MarkdownRendererWithCodeBlock extends MarkdownView.MarkdownView.MarkdownInsightRenderer {
  constructor(private readonly options: MarkdownRendererWithCodeBlockOptions = {}) {
    super();
  }
  #revealableLink(revealable: unknown, label: string): Lit.LitTemplate {
    return html`<devtools-link @click=${(e: Event) => {
      e.preventDefault();
      e.stopPropagation();
      void Common.Revealer.reveal(revealable);
    }}>${Platform.StringUtilities.trimEndWithMaxLength(label, 100)}</devtools-link>`;
  }

  #renderLink(
      href: string,
      fallbackText: string,
      ): Lit.LitTemplate|null {
    if (href.startsWith('#req-')) {
      const request = Logs.NetworkLog.NetworkLog.instance().requests().find(
          req => req.requestId() === href.substring(5),
      );

      if (request) {
        return this.#revealableLink(request, request.url());
      }
      return html`${fallbackText}`;
    }
    if (href.startsWith('#file-')) {
      const fileId = Number(href.substring(6));
      const file = AiAssistanceModel.ContextSelectionAgent.ContextSelectionAgent.getSourceById(
          fileId,
          this.options.getEstablishedOrigin?.(),
      );

      if (file) {
        return this.#revealableLink(file, file.name());
      }
      // Fall back to plain text to prevent rendering unauthorized or unverified links.
      return html`${fallbackText}`;
    }
    return null;
  }

  override templateForToken(token: Marked.Marked.MarkedToken): Lit.LitTemplate|null {
    if (token.type === 'link') {
      const link = this.#renderLink(token.href, token.text);
      if (link) {
        return link;
      }
    }

    if (token.type === 'code') {
      const lines = (token.text).split('\n');
      if (lines[0]?.trim() === 'css') {
        token.lang = 'css';
        token.text = lines.slice(1).join('\n');
      }
    }

    if (token.type === 'codespan') {
      // LLM likes outputting the link inside a codespan block.
      // Remove the codespan and render the link directly.
      const matches = token.text.match(/^\[(.*)\]\((.+)\)$/);
      if (matches?.[2]) {
        const link = this.#renderLink(
            matches[2],
            matches[1],
        );
        if (link) {
          return link;
        }
      }
    }

    return super.templateForToken(token);
  }
}