chore: copy as prompt in ui should have codeframe (#34943)

Signed-off-by: Simon Knott <info@simonknott.de>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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

@@ -29,6 +29,7 @@ import { AIConversation } from './aiConversation';
 import { ToolbarButton } from '@web/components/toolbarButton';
 import { useIsLLMAvailable, useLLMChat } from './llm';
 import { useAsyncMemo } from '@web/uiUtils';
+import { useSources } from './sourceTab';
 
 const CommitInfoContext = React.createContext<MetadataWithCommitInfo | undefined>(undefined);
 
@@ -53,18 +54,47 @@ function usePageSnapshot(actions: modelUtil.ActionTraceEventInContext[]) {
   }, [actions], undefined);
 }
 
+function useCodeFrame(stack: StackFrame[] | undefined, sources: Map<string, modelUtil.SourceModel>, width: number) {
+  const selectedFrame = stack?.[0];
+  const { source } = useSources(stack, 0, sources);
+  return React.useMemo(() => {
+    if (!source.content)
+      return '';
+
+    const targetLine = selectedFrame?.line ?? 0;
+
+    const lines = source.content.split('\n');
+    const start = Math.max(0, targetLine - width);
+    const end = Math.min(lines.length, targetLine + width);
+    const lineNumberWidth = String(end).length;
+    const codeFrame = lines.slice(start, end).map((line, i) => {
+      const lineNumber = start + i + 1;
+      const paddedLineNumber = String(lineNumber).padStart(lineNumberWidth, ' ');
+      if (lineNumber !== targetLine)
+        return `  ${(paddedLineNumber)} | ${line}`;
+
+      let highlightLine = `> ${paddedLineNumber} | ${line}`;
+      if (selectedFrame?.column)
+        highlightLine += `\n${' '.repeat(4 + lineNumberWidth + selectedFrame.column)}^`;
+      return highlightLine;
+    }).join('\n');
+    return codeFrame;
+  }, [source, selectedFrame, width]);
+}
+
 const CopyPromptButton: React.FC<{
   error: string;
+  codeFrame: string;
   pageSnapshot?: string;
   diff?: string;
-}> = ({ error, pageSnapshot, diff }) => {
+}> = ({ error, codeFrame, pageSnapshot, diff }) => {
   const prompt = React.useMemo(
       () => fixTestPrompt(
-          error,
+          error + '\n\n' + codeFrame,
           diff,
           pageSnapshot
       ),
-      [error, diff, pageSnapshot]
+      [error, diff, codeFrame, pageSnapshot]
   );
 
   return (
@@ -97,7 +127,7 @@ export function useErrorsTabModel(model: modelUtil.MultiTraceModel | undefined):
   }, [model]);
 }
 
-function Error({ message, error, errorId, sdkLanguage, pageSnapshot, revealInSource }: { message: string, error: ErrorDescription, errorId: string, sdkLanguage: Language, pageSnapshot?: string, revealInSource: (error: ErrorDescription) => void  }) {
+function Error({ message, error, errorId, sdkLanguage, pageSnapshot, revealInSource, sources }: { message: string, error: ErrorDescription, errorId: string, sdkLanguage: Language, pageSnapshot?: string, revealInSource: (error: ErrorDescription) => void, sources: Map<string, modelUtil.SourceModel> }) {
   const [showLLM, setShowLLM] = React.useState(false);
   const llmAvailable = useIsLLMAvailable();
   const metadata = useCommitInfo();
@@ -111,6 +141,8 @@ function Error({ message, error, errorId, sdkLanguage, pageSnapshot, revealInSou
     longLocation = stackFrame.file + ':' + stackFrame.line;
   }
 
+  const codeFrame = useCodeFrame(error.stack, sources, 3);
+
   return <div style={{ display: 'flex', flexDirection: 'column', overflowX: 'clip' }}>
     <div className='hbox' style={{
       alignItems: 'center',
@@ -127,7 +159,7 @@ function Error({ message, error, errorId, sdkLanguage, pageSnapshot, revealInSou
       <span style={{ position: 'absolute', right: '5px' }}>
         {llmAvailable
           ? <FixWithAIButton conversationId={errorId} onChange={setShowLLM} value={showLLM} error={message} diff={metadata?.gitDiff} pageSnapshot={pageSnapshot} />
-          : <CopyPromptButton error={message} pageSnapshot={pageSnapshot} diff={metadata?.gitDiff} />}
+          : <CopyPromptButton error={message} codeFrame={codeFrame} pageSnapshot={pageSnapshot} diff={metadata?.gitDiff} />}
       </span>
     </div>
 
@@ -179,9 +211,10 @@ export const ErrorsTab: React.FunctionComponent<{
   errorsModel: ErrorsTabModel,
   actions: modelUtil.ActionTraceEventInContext[],
   wallTime: number,
+  sources: Map<string, modelUtil.SourceModel>,
   sdkLanguage: Language,
   revealInSource: (error: ErrorDescription) => void,
-}> = ({ errorsModel, sdkLanguage, revealInSource, actions, wallTime }) => {
+}> = ({ errorsModel, sdkLanguage, revealInSource, actions, wallTime, sources }) => {
   const pageSnapshot = usePageSnapshot(actions);
 
   if (!errorsModel.errors.size)
@@ -190,7 +223,7 @@ export const ErrorsTab: React.FunctionComponent<{
   return <div className='fill' style={{ overflow: 'auto' }}>
     {[...errorsModel.errors.entries()].map(([message, error]) => {
       const errorId = `error-${wallTime}-${message}`;
-      return <Error key={errorId} errorId={errorId} message={message} error={error} revealInSource={revealInSource} sdkLanguage={sdkLanguage} pageSnapshot={pageSnapshot} />;
+      return <Error key={errorId} errorId={errorId} message={message} error={error} sources={sources} revealInSource={revealInSource} sdkLanguage={sdkLanguage} pageSnapshot={pageSnapshot} />;
     })}
   </div>;
 };

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

@@ -27,25 +27,8 @@ import { CopyToClipboard } from './copyToClipboard';
 import { ToolbarButton } from '@web/components/toolbarButton';
 import { Toolbar } from '@web/components/toolbar';
 
-export const SourceTab: React.FunctionComponent<{
-  stack?: StackFrame[],
-  stackFrameLocation: 'bottom' | 'right',
-  sources: Map<string, SourceModel>,
-  rootDir?: string,
-  fallbackLocation?: SourceLocation,
-  onOpenExternally?: (location: SourceLocation) => void,
-}> = ({ stack, sources, rootDir, fallbackLocation, stackFrameLocation, onOpenExternally }) => {
-  const [lastStack, setLastStack] = React.useState<StackFrame[] | undefined>();
-  const [selectedFrame, setSelectedFrame] = React.useState<number>(0);
-
-  React.useEffect(() => {
-    if (lastStack !== stack) {
-      setLastStack(stack);
-      setSelectedFrame(0);
-    }
-  }, [stack, lastStack, setLastStack, setSelectedFrame]);
-
-  const { source, highlight, targetLine, fileName, location } = useAsyncMemo<{ source: SourceModel, targetLine?: number, fileName?: string, highlight: SourceHighlight[], location?: SourceLocation }>(async () => {
+export function useSources(stack: StackFrame[] | undefined, selectedFrame: number, sources: Map<string, SourceModel>, rootDir?: string, fallbackLocation?: SourceLocation) {
+  return useAsyncMemo<{ source: SourceModel, targetLine?: number, fileName?: string, highlight: SourceHighlight[], location?: SourceLocation }>(async () => {
     const actionLocation = stack?.[selectedFrame];
     const shouldUseFallback = !actionLocation?.file;
     if (shouldUseFallback && !fallbackLocation)
@@ -84,6 +67,27 @@ export const SourceTab: React.FunctionComponent<{
     }
     return { source, highlight, targetLine, fileName, location };
   }, [stack, selectedFrame, rootDir, fallbackLocation], { source: { errors: [], content: 'Loading\u2026' }, highlight: [] });
+}
+
+export const SourceTab: React.FunctionComponent<{
+  stack?: StackFrame[],
+  stackFrameLocation: 'bottom' | 'right',
+  sources: Map<string, SourceModel>,
+  rootDir?: string,
+  fallbackLocation?: SourceLocation,
+  onOpenExternally?: (location: SourceLocation) => void,
+}> = ({ stack, sources, rootDir, fallbackLocation, stackFrameLocation, onOpenExternally }) => {
+  const [lastStack, setLastStack] = React.useState<StackFrame[] | undefined>();
+  const [selectedFrame, setSelectedFrame] = React.useState<number>(0);
+
+  React.useEffect(() => {
+    if (lastStack !== stack) {
+      setLastStack(stack);
+      setSelectedFrame(0);
+    }
+  }, [stack, lastStack, setLastStack, setSelectedFrame]);
+
+  const { source, highlight, targetLine, fileName, location } = useSources(stack, selectedFrame, sources, rootDir, fallbackLocation);
 
   const openExternally = React.useCallback(() => {
     if (!location)

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

@@ -193,7 +193,7 @@ export const Workbench: React.FunctionComponent<{
     id: 'errors',
     title: 'Errors',
     errorCount: errorsModel.errors.size,
-    render: () => <ErrorsTab errorsModel={errorsModel} sdkLanguage={sdkLanguage} revealInSource={error => {
+    render: () => <ErrorsTab errorsModel={errorsModel} sources={sources} sdkLanguage={sdkLanguage} revealInSource={error => {
       if (error.action)
         setSelectedAction(error.action);
       else

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

@@ -503,11 +503,11 @@ test('skipped steps should have an indicator', async ({ runUITest }) => {
 test('should show copy prompt button in errors tab', async ({ runUITest }) => {
   const { page } = await runUITest({
     'a.spec.ts': `
-      import { test, expect } from '@playwright/test';
-      test('fails', async () => {
-        expect(1).toBe(2);
-      });
-    `,
+import { test, expect } from '@playwright/test';
+test('fails', async () => {
+  expect(1).toBe(2);
+});
+    `.trim(),
   });
 
   await page.getByText('fails').dblclick();
@@ -517,4 +517,11 @@ test('should show copy prompt button in errors tab', async ({ runUITest }) => {
   await page.locator('.tab-errors').getByRole('button', { name: 'Copy as Prompt' }).click();
   const prompt = await page.evaluate(() => navigator.clipboard.readText());
   expect(prompt, 'contains error').toContain('expect(received).toBe(expected)');
+  expect(prompt, 'contains codeframe').toContain(`
+  1 | import { test, expect } from '@playwright/test';
+  2 | test('fails', async () => {
+> 3 |   expect(1).toBe(2);
+                  ^
+  4 | });
+    `.trim());
 });