feat(ui): add AI tab

package.json

@@ -102,6 +102,7 @@
     "node-stream-zip": "^1.15.0",
     "react": "^18.1.0",
     "react-dom": "^18.1.0",
+    "react-markdown": "^9.0.3",
     "ssim.js": "^3.5.0",
     "typescript": "^5.7.3",
     "vite": "^5.4.14",

packages/playwright-core/src/server/index.ts

@@ -29,6 +29,7 @@ export { createPlaywright } from './playwright';
 export type { DispatcherScope } from './dispatchers/dispatcher';
 export type { Playwright } from './playwright';
 export { installRootRedirect, openTraceInBrowser, openTraceViewerApp, runTraceViewerApp, startTraceViewerServer } from './trace/viewer/traceViewer';
+export { getLLMFromEnv } from './llm';
 export { serverSideCallMetadata } from './instrumentation';
 export { SocksProxy } from '../common/socksProxy';
 export * from './fileUtils';

packages/playwright-core/src/server/llm.ts

@@ -0,0 +1,118 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export type LLMMessage = {
+  role: 'user' | 'assistant' | 'developer';
+  content: string;
+};
+
+export interface LLM {
+  chatCompletion(messages: LLMMessage[]): AsyncGenerator<string>;
+}
+
+async function *parseSSE(body: Response['body']): AsyncGenerator<string> {
+  const reader = body!.getReader();
+  const decoder = new TextDecoder();
+  let buffer = '';
+  while (true) {
+    const { value, done } = await reader.read();
+    if (done)
+      break;
+    buffer += decoder.decode(value, { stream: true });
+    const events = buffer.split('\n\n');
+    buffer = events.pop()!;
+    for (const event of events) {
+      const contentStart = event.indexOf('data: ');
+      if (contentStart === -1)
+        continue;
+      yield event.substring(contentStart + 'data: '.length);
+    }
+  }
+}
+
+class OpenAI implements LLM {
+  private baseURL = process.env.OPENAI_BASE_URL ?? 'https://api.openai.com';
+
+  constructor(private apiKey: string) {}
+
+  async *chatCompletion(messages: LLMMessage[])  {
+    const url = new URL('./v1/chat/completions', this.baseURL);
+    const response = await fetch(url, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'Authorization': `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify({
+        model: 'gpt-4o',
+        messages,
+        stream: true,
+      }),
+    });
+
+    if (response.status !== 200 || !response.body)
+      throw new Error('Failed to chat with OpenAI, unexpected status: ' + response.status + await response.text());
+
+    for await (const eventString of parseSSE(response.body)) {
+      if (eventString === '[DONE]')
+        break;
+      const event = JSON.parse(eventString);
+      if (event.object === 'chat.completion.chunk') {
+        if (event.choices[0].finish_reason)
+          break;
+        yield event.choices[0].delta.content;
+      }
+    }
+  }
+}
+
+class Anthropic implements LLM {
+  private baseURL = process.env.ANTHROPIC_BASE_URL ?? 'https://api.anthropic.com';
+  constructor(private apiKey: string) {}
+  async *chatCompletion(messages: LLMMessage[]): AsyncGenerator<string> {
+    const response = await fetch(new URL('./v1/messages', this.baseURL), {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'x-api-key': this.apiKey,
+        'anthropic-version': '2023-06-01'
+      },
+      body: JSON.stringify({
+        model: 'claude-3-5-sonnet-20241022',
+        messages: messages.filter(({ role }) => role !== 'developer'),
+        system: messages.find(({ role }) => role === 'developer')?.content,
+        max_tokens: 1024,
+        stream: true,
+      })
+    });
+
+    if (response.status !== 200 || !response.body)
+      throw new Error('Failed to chat with Anthropic, unexpected status: ' + response.status + await response.text());
+
+    for await (const eventString of parseSSE(response.body)) {
+      const event = JSON.parse(eventString);
+      if (event.type  === 'content_block_delta')
+        yield event.delta.text;
+    }
+  }
+}
+
+export function getLLMFromEnv(): LLM | undefined {
+  if (process.env.OPENAI_API_KEY)
+    return new OpenAI(process.env.OPENAI_API_KEY);
+  if (process.env.ANTHROPIC_API_KEY)
+    return new Anthropic(process.env.ANTHROPIC_API_KEY);
+}

packages/playwright-core/src/server/trace/viewer/traceViewer.ts

@@ -17,7 +17,7 @@
 import * as fs from 'fs';
 import * as path from 'path';
 
-import { gracefullyProcessExitDoNotHang } from '../../../server';
+import { getLLMFromEnv, gracefullyProcessExitDoNotHang } from '../../../server';
 import { isUnderTest } from '../../../utils';
 import { HttpServer } from '../../../utils/httpServer';
 import { open } from '../../../utilsBundle';
@@ -29,6 +29,7 @@ import { ProgressController } from '../../progress';
 
 import type { Transport } from '../../../utils/httpServer';
 import type { BrowserType } from '../../browserType';
+import type { LLM, LLMMessage } from '../../llm';
 import type { Page } from '../../page';
 
 export type TraceViewerServerOptions = {
@@ -36,6 +37,7 @@ export type TraceViewerServerOptions = {
   port?: number;
   isServer?: boolean;
   transport?: Transport;
+  llm?: LLM;
 };
 
 export type TraceViewerRedirectOptions = {
@@ -46,6 +48,7 @@ export type TraceViewerRedirectOptions = {
   reporter?: string[];
   webApp?: string;
   isServer?: boolean;
+  llm?: LLM;
 };
 
 export type TraceViewerAppOptions = {
@@ -76,6 +79,24 @@ export async function startTraceViewerServer(options?: TraceViewerServerOptions)
     }
     if (relativePath.endsWith('/stall.js'))
       return true;
+    if (relativePath.startsWith('/llm/chat-completion') && options?.llm) {
+      const chunks: Buffer[] = [];
+      request.on('data', chunk => chunks.push(chunk));
+      request.on('end', async () => {
+        const messages = JSON.parse(Buffer.concat(chunks).toString()) as LLMMessage[];
+        response.setHeader('Content-Type', 'text/plain');
+        try {
+          const stream = options.llm!.chatCompletion(messages);
+          for await (const chunk of stream)
+            response.write(chunk);
+          response.end();
+        } catch (error) {
+          response.statusCode = 500;
+          response.end(error.toString());
+        }
+      });
+      return true;
+    }
     if (relativePath.startsWith('/file')) {
       try {
         const filePath = url.searchParams.get('path')!;
@@ -131,6 +152,8 @@ export async function installRootRedirect(server: HttpServer, traceUrls: string[
     params.append('project', project);
   for (const reporter of options.reporter || [])
     params.append('reporter', reporter);
+  if (options.llm)
+    params.append('llm', 'true');
 
   let baseUrl = '.';
   if (process.env.PW_HMR) {
@@ -149,6 +172,7 @@ export async function installRootRedirect(server: HttpServer, traceUrls: string[
 
 export async function runTraceViewerApp(traceUrls: string[], browserName: string, options: TraceViewerServerOptions & { headless?: boolean }, exitOnClose?: boolean) {
   validateTraceUrls(traceUrls);
+  options.llm = getLLMFromEnv();
   const server = await startTraceViewerServer(options);
   await installRootRedirect(server, traceUrls, options);
   const page = await openTraceViewerApp(server.urlPrefix('precise'), browserName, options);
@@ -159,6 +183,7 @@ export async function runTraceViewerApp(traceUrls: string[], browserName: string
 
 export async function runTraceInBrowser(traceUrls: string[], options: TraceViewerServerOptions) {
   validateTraceUrls(traceUrls);
+  options.llm = getLLMFromEnv();
   const server = await startTraceViewerServer(options);
   await installRootRedirect(server, traceUrls, options);
   await openTraceInBrowser(server.urlPrefix('human-readable'));

packages/playwright/src/runner/testServer.ts

@@ -17,7 +17,7 @@
 import * as fs from 'fs';
 import * as path from 'path';
 
-import { gracefullyProcessExitDoNotHang, installRootRedirect, openTraceInBrowser, openTraceViewerApp, registry, startTraceViewerServer } from 'playwright-core/lib/server';
+import { getLLMFromEnv, gracefullyProcessExitDoNotHang, installRootRedirect, openTraceInBrowser, openTraceViewerApp, registry, startTraceViewerServer } from 'playwright-core/lib/server';
 import { ManualPromise, isUnderTest } from 'playwright-core/lib/utils';
 import { open } from 'playwright-core/lib/utilsBundle';
 
@@ -59,7 +59,7 @@ class TestServer {
 
   async start(options: { host?: string, port?: number }): Promise<HttpServer> {
     this._dispatcher = new TestServerDispatcher(this._configLocation, this._configCLIOverrides);
-    return await startTraceViewerServer({ ...options, transport: this._dispatcher.transport });
+    return await startTraceViewerServer({ ...options, transport: this._dispatcher.transport, llm: getLLMFromEnv() });
   }
 
   async stop() {
@@ -437,6 +437,7 @@ export class TestServerDispatcher implements TestServerInterface {
 
 export async function runUIMode(configFile: string | undefined, configCLIOverrides: ConfigCLIOverrides, options: TraceViewerServerOptions & TraceViewerRedirectOptions): Promise<reporterTypes.FullResult['status'] | 'restarted'> {
   const configLocation = resolveConfigLocation(configFile);
+  options.llm = getLLMFromEnv();
   return await innerRunTestServer(configLocation, configCLIOverrides, options, async (server: HttpServer, cancelPromise: ManualPromise<void>) => {
     await installRootRedirect(server, [], { ...options, webApp: 'uiMode.html' });
     if (options.host !== undefined || options.port !== undefined) {

packages/trace-viewer/src/sw/main.ts

@@ -36,16 +36,25 @@ const scopePath = new URL(self.registration.scope).pathname;
 
 const loadedTraces = new Map<string, { traceModel: TraceModel, snapshotServer: SnapshotServer }>();
 
-const clientIdToTraceUrls = new Map<string, { limit: number | undefined, traceUrls: Set<string>, traceViewerServer: TraceViewerServer }>();
+const clientIdToTraceViewerServer = new Map<string, TraceViewerServer>();
+const clientIdToTraceUrls = new Map<string, { limit: number | undefined, traceUrls: Set<string> }>();
+
+function getTraceViewerServer(client: any) {
+  const clientId: string = client?.id ?? '';
+  if (!clientIdToTraceViewerServer.has(clientId)) {
+    const clientURL = new URL(client?.url ?? self.registration.scope);
+    const traceViewerServerBaseUrl = new URL(clientURL.searchParams.get('server') ?? '../', clientURL);
+    clientIdToTraceViewerServer.set(clientId, new TraceViewerServer(traceViewerServerBaseUrl));
+  }
+  return clientIdToTraceViewerServer.get(clientId)!;
+}
 
 async function loadTrace(traceUrl: string, traceFileName: string | null, client: any | undefined, limit: number | undefined, progress: (done: number, total: number) => undefined): Promise<TraceModel> {
   await gc();
   const clientId = client?.id ?? '';
   let data = clientIdToTraceUrls.get(clientId);
   if (!data) {
-    const clientURL = new URL(client?.url ?? self.registration.scope);
+    data = { limit, traceUrls: new Set() };
-    const traceViewerServerBaseUrl = new URL(clientURL.searchParams.get('server') ?? '../', clientURL);
-    data = { limit, traceUrls: new Set(), traceViewerServer: new TraceViewerServer(traceViewerServerBaseUrl) };
     clientIdToTraceUrls.set(clientId, data);
   }
   data.traceUrls.add(traceUrl);
@@ -54,7 +63,8 @@ async function loadTrace(traceUrl: string, traceFileName: string | null, client:
   try {
     // Allow 10% to hop from sw to page.
     const [fetchProgress, unzipProgress] = splitProgress(progress, [0.5, 0.4, 0.1]);
-    const backend = traceUrl.endsWith('json') ? new FetchTraceModelBackend(traceUrl, data.traceViewerServer) : new ZipTraceModelBackend(traceUrl, data.traceViewerServer, fetchProgress);
+    const traceViewerServer = getTraceViewerServer(client);
+    const backend = traceUrl.endsWith('json') ? new FetchTraceModelBackend(traceUrl, traceViewerServer) : new ZipTraceModelBackend(traceUrl, traceViewerServer, fetchProgress);
     await traceModel.load(backend, unzipProgress);
   } catch (error: any) {
     // eslint-disable-next-line no-console
@@ -156,15 +166,18 @@ async function doFetch(event: FetchEvent): Promise<Response> {
 
     if (relativePath.startsWith('/file/')) {
       const path = url.searchParams.get('path')!;
-      const traceViewerServer = clientIdToTraceUrls.get(event.clientId ?? '')?.traceViewerServer;
+      const traceViewerServer = getTraceViewerServer(client);
-      if (!traceViewerServer)
-        throw new Error('client is not initialized');
       const response = await traceViewerServer.readFile(path);
       if (!response)
         return new Response(null, { status: 404 });
       return response;
     }
 
+    if (relativePath.startsWith('/llm/chat-completion')) {
+      const traceViewerServer = getTraceViewerServer(client);
+      return traceViewerServer.chatCompletion(await request.text());
+    }
+
     // Fallback for static assets.
     return fetch(event.request);
   }

packages/trace-viewer/src/sw/traceModelBackends.ts

@@ -160,4 +160,15 @@ export class TraceViewerServer {
       return;
     return response;
   }
+
+  async chatCompletion(body: string): Promise<Response> {
+    const url = new URL('trace/llm/chat-completion', this.baseUrl);
+    return await fetch(url, {
+      body,
+      method: "POST",
+      headers: {
+        'content-type': 'application/json'
+      }
+    });
+  }
 }

packages/trace-viewer/src/ui/aiTab.css

@@ -0,0 +1,95 @@
+/* Main container */
+.chat-container {
+  display: flex;
+  flex-direction: column;
+  justify-content: space-between;
+  height: 100vh;
+  color: #e0e0e0;
+}
+
+.messages-container {
+  overflow-y: auto;
+  padding: 16px;
+  gap: 16px;
+}
+
+.message {
+  gap: 12px;
+  max-width: 85%;
+}
+
+.user-message {
+  flex-direction: row-reverse;
+  margin-left: auto;
+  width: fit-content
+}
+
+.message-icon {
+  width: 28px;
+  height: 28px;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  background-color: var(--vscode-titleBar-inactiveBackground);
+  flex-shrink: 0;
+}
+
+.message-content {
+  background-color: var(--vscode-titleBar-inactiveBackground);
+  color: var(--vscode-titleBar-activeForeground);
+  padding: 1px 8px;
+  margin-bottom: 16px;
+}
+
+.message-content pre {
+  text-wrap: auto;
+}
+
+.user-message .message-content {
+  background-color: var(--vscode-titleBar-activeBackground);
+}
+
+/* Input form styles */
+.input-form {
+  display: flex;
+  height: 80px;
+  gap: 8px;
+  padding: 16px;
+  background-color: var(--vscode-sideBar-background);
+  border-top: 1px solid var(--vscode-sideBarSectionHeader-border);
+}
+
+.message-input {
+  flex: 1;
+  padding: 8px 12px;
+  border: 1px solid var(--vscode-settings-textInputBorder);
+  background-color: var(--vscode-settings-textInputBackground);
+  font-size: 14px;
+  outline: none;
+  transition: border-color 0.2s;
+}
+
+.message-input:focus {
+  border-color: #0078d4;
+}
+
+.send-button {
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  padding: 8px;
+  background-color: var(--vscode-button-background);
+  border: none;
+  color: white;
+  cursor: pointer;
+  transition: background-color 0.2s;
+}
+
+.send-button:hover {
+  background-color: var(--vscode-button-hoverBackground);
+}
+
+.send-button:disabled {
+  background-color: var(--vscode-disabledForeground);
+  cursor: not-allowed;
+}

packages/trace-viewer/src/ui/aiTab.tsx

@@ -0,0 +1,124 @@
+import { useCallback, useEffect, useState } from 'react';
+import Markdown from 'react-markdown'
+import './aiTab.css';
+import type { LLMMessage } from 'playwright-core/lib/server/llm';
+import { clsx } from '@web/uiUtils';
+
+export interface AIState {
+  prompt?: string;
+  variables: Record<string, string>;
+}
+
+export function AITab({ state }: { state?: AIState }) {
+  const [input, setInput] = useState('');
+  const [messages, setMessages] = useState<LLMMessage[]>([
+    {
+      role: "developer",
+      content: 'You are a helpful assistant, skilled in programming and software testing with Playwright. Help me write good code. Be bold, creative and assertive when you suggest solutions.'
+    }
+  ]);
+  const [abort, setAbort] = useState<AbortController>();
+
+  const onSubmit = useCallback(async (event: React.FormEvent<HTMLFormElement>) => {
+    event.preventDefault();
+    setInput('');
+    const content = new FormData(event.target as any).get('content') as string;
+    const messages = await new Promise<LLMMessage[]>(resolve => {
+      setMessages(messages => {
+          const newMessages = [...messages, { role: 'user', content } as LLMMessage];
+          resolve(newMessages);
+          return newMessages;
+      })
+    });
+    const controller = new AbortController();
+    setAbort(controller);
+
+    const hydratedMessages = messages.map(message => {
+      let content = message.content;
+      for (const [variable, value] of Object.entries(state?.variables || {})) {
+        content = content.replaceAll(variable, value);
+      }
+      return { ...message, content };
+    })
+
+    const response = await fetch('./llm/chat-completion', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json'
+      },
+      body: JSON.stringify(hydratedMessages),
+      signal: controller.signal,
+    });
+    try {
+      const decoder = new TextDecoder();
+      let reply = '';
+      function update() {
+        setMessages(messages => {
+          return messages.slice(0, -1).concat([{ role: 'assistant', content: reply }]);
+        });
+      }
+      await response.body?.pipeTo(new WritableStream({
+        write(chunk) {
+          reply += decoder.decode(chunk, { stream: true });
+          update();
+        },
+        close() {
+          reply += decoder.decode();
+          update();
+        }
+      }));
+    } finally {
+      setAbort(undefined);
+    }
+  }, []);
+
+  useEffect(() => {
+    if (state?.prompt)
+      setInput(state?.prompt);
+  }, [state])
+
+  return (
+    <div className="chat-container">
+      <div className="messages-container">
+        {messages.filter(({ role }) => role !== 'developer').map((message, index) => (
+          <div
+            key={'' + index}
+            className={clsx('message', message.role === 'user' && 'user-message')}
+          >
+            {message.role === 'assistant' && (
+              <div className="message-icon">
+                <img src="playwright-logo.svg" />
+              </div>
+            )}
+            <div className="message-content">
+              <Markdown>{message.content}</Markdown>
+            </div>
+          </div>
+        ))}
+      </div>
+
+      <form onSubmit={onSubmit} className="input-form">
+        <input
+          type="text"
+          name='content'
+          value={input}
+          onChange={(e) => setInput(e.target.value)}
+          placeholder="Ask a question..."
+          className="message-input"
+        />
+        {abort ? (
+          <button type="button" className="send-button" onClick={(evt) => {
+            evt.preventDefault()
+            abort.abort()
+          }}>
+            Cancel
+          </button>
+        ) : (
+          <button type="submit" className="send-button" disabled={!input.trim()}>
+            Send
+          </button>  
+        )}
+      </form>
+    </div>
+  );
+}

packages/trace-viewer/src/ui/errorsTab.tsx

@@ -25,6 +25,8 @@ import { CopyToClipboardTextButton } from './copyToClipboard';
 import { attachmentURL } from './attachmentsTab';
 import { fixTestPrompt } from '@web/components/prompts';
 import type { GitCommitInfo } from '@testIsomorphic/types';
+import type { AIState } from './aiTab';
+import { ToolbarButton } from '@web/components/toolbarButton';
 
 const GitCommitInfoContext = React.createContext<GitCommitInfo | undefined>(undefined);
 
@@ -36,7 +38,45 @@ export function useGitCommitInfo() {
   return React.useContext(GitCommitInfoContext);
 }
 
-const PromptButton: React.FC<{
+const OpenInAITabButton: React.FC<{
+  error: string;
+  actions: modelUtil.ActionTraceEventInContext[];
+  fixWithAI(state: AIState): void;
+}> = ({ error, actions, fixWithAI }) => {
+  const [pageSnapshot, setPageSnapshot] = React.useState<string>();
+
+  React.useEffect(( )=> {
+    for (const action of actions) {
+      for (const attachment of action.attachments ?? []) {
+        if (attachment.name === 'pageSnapshot') {
+          fetch(attachmentURL({ ...attachment, traceUrl: action.context.traceUrl })).then(async response => {
+            setPageSnapshot(await response.text());
+          });
+          return;
+        }
+      }
+    }
+  }, [actions]);
+
+  const gitCommitInfo = useGitCommitInfo();
+  const prompt = React.useMemo(
+      () => fixTestPrompt(
+          error,
+          gitCommitInfo?.['pull.diff'] ?? gitCommitInfo?.['revision.diff'],
+          pageSnapshot
+      ),
+      [error, gitCommitInfo, pageSnapshot]
+  );
+
+  return <ToolbarButton onClick={() => {
+    fixWithAI({
+      prompt: "What's going wrong in this error? @error",
+      variables: { '@error': prompt },
+    });
+  }} title="Fix with AI" className='copy-to-clipboard-text-button'>Fix with AI</ToolbarButton>;
+};
+
+const CopyPromptButton: React.FC<{
   error: string;
   actions: modelUtil.ActionTraceEventInContext[];
 }> = ({ error, actions }) => {
@@ -100,7 +140,8 @@ export const ErrorsTab: React.FunctionComponent<{
   actions: modelUtil.ActionTraceEventInContext[],
   sdkLanguage: Language,
   revealInSource: (error: ErrorDescription) => void,
-}> = ({ errorsModel, sdkLanguage, revealInSource, actions }) => {
+  fixWithAI?(state: AIState): void;
+}> = ({ errorsModel, sdkLanguage, revealInSource, actions, fixWithAI }) => {
   if (!errorsModel.errors.size)
     return <PlaceholderPanel text='No errors' />;
 
@@ -127,9 +168,10 @@ export const ErrorsTab: React.FunctionComponent<{
             @ <span title={longLocation} onClick={() => revealInSource(error)}>{location}</span>
           </div>}
           <span style={{ position: 'absolute', right: '5px' }}>
-            <PromptButton error={message} actions={actions} />
+            {fixWithAI ? <OpenInAITabButton error={message} actions={actions} fixWithAI={fixWithAI} /> : <CopyPromptButton error={message} actions={actions} />}
           </span>
         </div>
+        
         <ErrorMessage error={message} />
       </div>;
     })}

packages/trace-viewer/src/ui/uiModeTraceView.tsx

@@ -31,7 +31,8 @@ export const TraceView: React.FC<{
   onOpenExternally?: (location: SourceLocation) => void,
   revealSource?: boolean,
   pathSeparator: string,
-}> = ({ item, rootDir, onOpenExternally, revealSource, pathSeparator }) => {
+  llmAvailable?: boolean,
+}> = ({ item, rootDir, onOpenExternally, revealSource, pathSeparator, llmAvailable }) => {
   const [model, setModel] = React.useState<{ model: MultiTraceModel, isLive: boolean } | undefined>();
   const [counter, setCounter] = React.useState(0);
   const pollTimer = React.useRef<NodeJS.Timeout | null>(null);
@@ -97,6 +98,7 @@ export const TraceView: React.FC<{
     annotations={item.testCase?.annotations || []}
     onOpenExternally={onOpenExternally}
     revealSource={revealSource}
+    llmAvailable={llmAvailable}
   />;
 };
 

packages/trace-viewer/src/ui/uiModeView.tsx

@@ -61,6 +61,7 @@ const queryParams = {
   updateSnapshots: (searchParams.get('updateSnapshots') as 'all' | 'none' | 'missing' | undefined) || undefined,
   reporters: searchParams.has('reporter') ? searchParams.getAll('reporter') : undefined,
   pathSeparator: searchParams.get('pathSeparator') || '/',
+  llm: searchParams.has('llm'),
 };
 if (queryParams.updateSnapshots && !['all', 'none', 'missing'].includes(queryParams.updateSnapshots))
   queryParams.updateSnapshots = undefined;
@@ -438,6 +439,7 @@ export const UIModeView: React.FC<{}> = ({
               rootDir={testModel?.config?.rootDir}
               revealSource={revealSource}
               onOpenExternally={location => testServerConnection?.openNoReply({ location: { file: location.file, line: location.line, column: location.column } })}
+              llmAvailable={queryParams.llm}
             />
           </GitCommitInfoProvider>
         </div>

packages/trace-viewer/src/ui/workbench.tsx

@@ -43,6 +43,7 @@ import { testStatusIcon, testStatusText } from './testUtils';
 import type { UITestStatus } from './testUtils';
 import type { AfterActionTraceEventAttachment } from '@trace/trace';
 import type { HighlightedElement } from './snapshotTab';
+import { AIState, AITab } from './aiTab';
 
 export const Workbench: React.FunctionComponent<{
   model?: modelUtil.MultiTraceModel,
@@ -56,7 +57,8 @@ export const Workbench: React.FunctionComponent<{
   inert?: boolean,
   onOpenExternally?: (location: modelUtil.SourceLocation) => void,
   revealSource?: boolean,
-}> = ({ model, showSourcesFirst, rootDir, fallbackLocation, isLive, hideTimeline, status, annotations, inert, onOpenExternally, revealSource }) => {
+  llmAvailable?: boolean,
+}> = ({ model, showSourcesFirst, rootDir, fallbackLocation, isLive, hideTimeline, status, annotations, inert, onOpenExternally, revealSource, llmAvailable }) => {
   const [selectedCallId, setSelectedCallId] = React.useState<string | undefined>(undefined);
   const [revealedError, setRevealedError] = React.useState<ErrorDescription | undefined>(undefined);
   const [revealedAttachment, setRevealedAttachment] = React.useState<[attachment: AfterActionTraceEventAttachment, renderCounter: number] | undefined>(undefined);
@@ -67,8 +69,10 @@ export const Workbench: React.FunctionComponent<{
   const [selectedPropertiesTab, setSelectedPropertiesTab] = useSetting<string>('propertiesTab', showSourcesFirst ? 'source' : 'call');
   const [isInspecting, setIsInspectingState] = React.useState(false);
   const [highlightedElement, setHighlightedElement] = React.useState<HighlightedElement>({ lastEdited: 'none' });
+  const [aiState, setAIState] = React.useState<AIState>();
   const [selectedTime, setSelectedTime] = React.useState<Boundaries | undefined>();
   const [sidebarLocation, setSidebarLocation] = useSetting<'bottom' | 'right'>('propertiesSidebarLocation', 'bottom');
+  const [showAITab, setShowAITab] = useSetting<'on' | 'off'>('aiTab', 'off')
 
   const setSelectedAction = React.useCallback((action: modelUtil.ActionTraceEventInContext | undefined) => {
     setSelectedCallId(action?.callId);
@@ -199,7 +203,11 @@ export const Workbench: React.FunctionComponent<{
       else
         setRevealedError(error);
       selectPropertiesTab('source');
-    }} actions={model?.actions ?? []} />
+    }} actions={model?.actions ?? []} fixWithAI={showLLMChat ? prompt => {
+      setShowAITab('on');
+      setAIState(prompt);
+      selectPropertiesTab('ai');
+    } : undefined} />
   };
 
   // Fallback location w/o action stands for file / test.
@@ -267,6 +275,14 @@ export const Workbench: React.FunctionComponent<{
     tabs.push(annotationsTab);
   }
 
+  const showLLMChat = llmAvailable && showAITab === 'on';
+  if (showLLMChat)
+    tabs.push({
+      id: 'ai',
+      title: 'AI',
+      render: () => <AITab state={aiState} />,
+    });
+
   if (showSourcesFirst) {
     const sourceTabIndex = tabs.indexOf(sourceTab);
     tabs.splice(sourceTabIndex, 1);

tests/playwright-test/ui-mode-llm.spec.ts

@@ -0,0 +1,134 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { test, expect, retries } from './ui-mode-fixtures';
+
+test.describe.configure({ mode: 'parallel', retries });
+
+test('openai', async ({ runUITest, server }) => {
+  const OPENAI_API_KEY = 'fake-key';
+
+  server.setRoute('/v1/chat/completions', (req, res) => {
+    expect(req.headers['authorization']).toBe('Bearer ' + OPENAI_API_KEY);
+    const event = {
+      object: 'chat.completion.chunk',
+      choices: [{ delta: { content: 'This is a mock response' } }]
+    };
+    res.end(`\n\ndata: ${JSON.stringify(event)}\n\n`);
+  });
+
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('trace test', async ({ page }) => {
+        await page.setContent('<button>Submit</button>');
+        expect(1).toBe(2);
+      });
+    `,
+  }, {
+    OPENAI_API_KEY,
+    OPENAI_BASE_URL: server.PREFIX,
+  });
+
+  await page.getByTitle('Run all').click();
+  await page.getByText('Errors', { exact: true }).click();
+  await page.getByRole('button', { name: 'Fix with AI' }).click();
+  await page.getByRole('button', { name: 'Send' }).click();
+
+  await expect(page.getByRole('tabpanel', { name: 'AI' })).toMatchAriaSnapshot(`
+    - tabpanel "AI":
+      - paragraph: /Here is the error:/
+      - code: /Expected. 2 Received. 1/
+      - paragraph: This is a mock response
+  `);
+});
+
+test('anthropic', async ({ runUITest, server }) => {
+  const ANTHROPIC_API_KEY = 'fake-key';
+
+  server.setRoute('/v1/messages', (req, res) => {
+    expect(req.headers['x-api-key']).toBe(ANTHROPIC_API_KEY);
+    const event = {
+      type: 'content_block_delta',
+      delta: { text: 'This is a mock response' }
+    };
+    res.end(`\n\ndata: ${JSON.stringify(event)}\n\n`);
+  });
+
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('trace test', async ({ page }) => {
+        await page.setContent('<button>Submit</button>');
+        expect(1).toBe(2);
+      });
+    `,
+  }, {
+    ANTHROPIC_API_KEY,
+    ANTHROPIC_BASE_URL: server.PREFIX,
+  });
+
+  await page.getByTitle('Run all').click();
+  await page.getByText('Errors', { exact: true }).click();
+  await page.getByRole('button', { name: 'Fix with AI' }).click();
+  await page.getByRole('button', { name: 'Send' }).click();
+
+  await expect(page.getByRole('tabpanel', { name: 'AI' })).toMatchAriaSnapshot(`
+    - tabpanel "AI":
+      - paragraph: /Here is the error:/
+      - code: /Expected. 2 Received. 1/
+      - paragraph: This is a mock response
+  `);
+});
+
+test('invisible without key', async ({ runUITest }) => {
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('trace test', async ({ page }) => {
+        await page.setContent('<button>Submit</button>');
+        expect(1).toBe(2);
+      });
+    `,
+  });
+
+  await page.getByTitle('Run all').click();
+  await page.getByText('Errors', { exact: true }).click();
+  await page.getByRole('button', { name: 'Fix with AI' }).click();
+  await expect(page.getByRole('tabpanel', { name: 'AI' })).not.toBeVisible();
+});
+
+test('tab stays visible once activated', async ({ runUITest }) => {
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('trace test', async ({ page }) => {
+        await page.setContent('<button>Submit</button>');
+        expect(1).toBe(2);
+      });
+    `,
+  }, {
+    ANTHROPIC_API_KEY: 'fake-key',
+  });
+
+  await page.getByTitle('Run all').click();
+  await page.getByText('Errors', { exact: true }).click();
+  await page.getByRole('button', { name: 'Fix with AI' }).click();
+  await expect(page.getByRole('tabpanel', { name: 'AI' })).toBeVisible();
+
+  await page.reload();
+  await expect(page.getByRole('tabpanel', { name: 'AI' })).toBeVisible();
+});