Dispel AI Agent: Tools, Prompts, and CSS Streaming

import { tool } from "ai";
import { z } from "zod";

const generateCss = tool({
  description:
    "Generate CSS styles to the page. Return only valid CSS rules with no markdown, explanations, or comments.",
  inputSchema: z.object({
    css: z
      .string()
      .describe("Raw CSS rules targeting specific selectors on the page"),
  }),
  execute: async ({ css }: { css: string }) => css,
});

const reportErrorSchema = z.object({
  code: z.enum(LLM_ERROR_CODES),
  message: z
    .string()
    .describe("Specific reason the request cannot be fulfilled"),
});

const reportError = tool({
  description:
    "Report why the CSS request cannot be fulfilled with a specific error code and message.",
  inputSchema: reportErrorSchema,
  execute: async ({ code, message }: z.infer<typeof reportErrorSchema>) => ({
    code,
    message,
  }),
});

export function createGetPageHTMLTool(tabId: number) {
  return tool({
    description:
      "Fetch cleaned HTML of the current viewport. Call with no selector to see all visible top-level elements. Call with a selector to inspect a specific element. Returns the HTML string, or 'Element not found in current viewport' if the selector does not match a visible element.",
    inputSchema: z.object({
      selector: z
        .string()
        .optional()
        .describe(
          "Optional CSS selector to scope extraction to a specific element"
        ),
    }),
    execute: async ({ selector }: { selector?: string }) => {
      const result = await sendTabMessage(tabId, "GET_PAGE_HTML", { selector });
      if (result._tag === "Err") {
        logError("agentTools.getPageHTML", result.error);
        return `Failed to fetch page HTML (${result.error.code})`;
      }
      return result.value;
    },
  });
}

export const STATIC_GUIDELINES = `
===================
SELECTOR STRATEGY
===================
- Use precise, stable, and minimal selectors that exist on page reload
- Prefer:
  1. id selectors (#id)
  2. data-testid, data-test, data-qa attributes
  3. class selectors (.class)
  4. structural selectors (tag, nth-child, descendant combinators)
- Use selectors from the provided node list when available
- Avoid overly deep selector chains
- Avoid generic selectors like div, span unless necessary

IMPORTANT - USE PROVIDED SELECTORS:
- The node list provides stable selectors (id, data-testid, class combinations)
- Prefer these over generating new selectors
- Only fall back to structural selectors if no stable selector exists

===================
SCOPING & SAFETY
===================
- Scope styles to the smallest relevant element or container
- Do NOT affect unrelated parts of the page
- Do NOT modify global layout unless explicitly requested

FORBIDDEN:
- Universal selectors (*) unless absolutely required
- all: unset / all: initial / aggressive resets
- Styling html or body unless explicitly requested
- Excessive use of !important
- position: fixed or absolute on major layout elements unless required

===================
MODERN CSS PRACTICES
===================
- Use modern CSS (flexbox, grid, gap, clamp, etc.)
- Prefer responsive units (%, rem, vw, vh) over fixed px where appropriate
- Use shorthand properties when clean and readable
- Use CSS variables ONLY if they improve clarity and reuse

===================
EXTERNAL RESOURCES
===================
- Avoid external resources (remote fonts, @import URLs, CDN-hosted assets, external images)
  unless explicitly required by the user prompt
- Prefer pure CSS solutions and system fonts when possible

===================
ATTACHED FILES
===================
- Attached files may contain reference styles, design tokens, implementation constraints,
  or existing CSS provided by the user
- Follow the guidance in attached files unless it directly conflicts with explicit instructions
  in the user's prompt
- When conflicts exist, the user's prompt text takes precedence

===================
VISUAL QUALITY
===================
- Maintain clean spacing, alignment, and hierarchy
- Ensure readability and contrast
- Avoid redundant or conflicting rules

===================
ROBUSTNESS
===================
- Handle incomplete or noisy data gracefully
- Ignore duplicate or irrelevant elements
- If elements cannot be confidently identified, use the reportError tool

===================
SHADOW DOM CONSTRAINT
===================
- Do NOT assume styles can affect elements inside Shadow DOM
- Only target elements accessible from the main DOM tree
- If element is inside Shadow DOM, use the reportError tool
`;

export type StreamChatChunk =
  | { type: "css"; css: string }
  | { type: "error"; code: AppErrorCode; message: string };

interface StreamChatOptions {
  abortSignal?: AbortSignal;
  apiKey: string;
  baseUrl: string;
  compiledCss?: string;
  message: string;
  modelId: string;
  pageContext: IPageContext;
  pickedElements?: { outerHTML: string; label: string }[];
  provider: PROVIDER_TYPES;
  tabId: number;
}

const CONTEXT_LIMIT_RE =
  /(context[_\s-]?length|context_length_exceeded|context window|maximum (context|input|content)( length)?|token(s)? (limit|exceeded|too many|count exceeds)|too many tokens|input too long|prompt too long|request (exceeds|too large)|maximum size|exceeded.*(context|token))/i;

const errorMatchers: ErrorMatcher[] = [
  (e) =>
    APICallError.isInstance(e) && e.statusCode === 429
      ? {
          code: "RATE_LIMITED",
          message: "Rate limit reached. Please wait a moment and try again.",
        }
      : null,
  (e) =>
    APICallError.isInstance(e) && e.statusCode === 413
      ? { code: "CONTEXT_LIMIT_REACHED", message: CONTEXT_LIMIT_MESSAGE }
      : null,
  (e) =>
    APICallError.isInstance(e) &&
    (e.statusCode === 400 || e.statusCode === 422) &&
    CONTEXT_LIMIT_RE.test(`${e.message ?? ""} ${e.responseBody ?? ""}`)
      ? { code: "CONTEXT_LIMIT_REACHED", message: CONTEXT_LIMIT_MESSAGE }
      : null,
];

type ModelFactory = (
  apiKey: string,
  baseUrl?: string
) => (modelId: string) => LanguageModel;

const modelFactories: Record<PROVIDER_TYPES, ModelFactory> = {
  openai:      (apiKey, baseUrl) => createOpenAI({ apiKey, baseURL: baseUrl }),
  anthropic:   (apiKey, baseUrl) => createAnthropic({ apiKey, baseURL: baseUrl }),
  google:      (apiKey, baseUrl) =>
    createGoogleGenerativeAI({ apiKey, baseURL: baseUrl }),
  local:       (apiKey, baseUrl) =>
    createOpenAICompatible({ name: "local", apiKey, baseURL: baseUrl ?? "" }),
  openrouter:  (apiKey, baseUrl) => createOpenRouter({ apiKey, baseUrl }),
};