improvement: better UX for AI context items (#6649)

AI Context Improvements

- Improved table context format: added sample values for columns and
inline formatting
- Styling of Resource's attached when using the agent (support
text/plain, markdown, and images)
- Reorganized completion sections with explicit ranking system (boosts
dont apply across sections, so need to rank the sections)
- Hide empty datasources from context menu

Various UI fixes
- Added hover states and improved visual hierarchy for resource and
resource-link blocks
- Implemented image preview popover for image resource links
- Added overflow handling to prevent horizontal scrolling in user
messages and tool notifications
- Fixed truncation for long accordion titles and adjusted z-index for
session tabs

Author: mscolnick

frontend/src/components/chat/acp/__tests__/context-utils.test.ts

@@ -133,7 +133,7 @@ describe("parseContextFromPrompt", () => {
       type: "resource",
       resource: {
         uri: "context.md",
-        mimeType: "text/markdown",
+        mimeType: "text/plain",
         text: "formatted context",
       },
     });

frontend/src/components/chat/acp/agent-panel.tsx

@@ -867,7 +867,7 @@ const AgentPanel: React.FC = () => {
             type: "resource",
             resource: {
               uri: "marimo_rules.md",
-              mimeType: "text/markdown",
+              mimeType: "text/plain",
               text: getAgentPrompt(filename),
             },
           },

frontend/src/components/chat/acp/blocks.tsx

@@ -30,7 +30,6 @@ import {
 } from "@/components/ui/popover";
 import { logNever } from "@/utils/assertNever";
 import { cn } from "@/utils/cn";
-import { type Base64String, base64ToDataURL } from "@/utils/json/base64";
 import { Strings } from "@/utils/strings";
 import { MarkdownRenderer } from "../markdown-renderer";
 import { SimpleAccordion } from "./common";
@@ -289,7 +288,7 @@ export const PlansBlock = (props: { data: PlanNotificationEvent[] }) => {
 
 export const UserMessagesBlock = (props: { data: UserNotificationEvent[] }) => {
   return (
-    <div className="flex flex-col gap-2 text-muted-foreground border p-2 bg-background rounded break-words">
+    <div className="flex flex-col gap-2 text-muted-foreground border p-2 bg-background rounded break-words overflow-x-hidden">
       <ContentBlocks data={props.data.map((item) => item.content)} />
     </div>
   );
@@ -364,58 +363,28 @@ export const ResourceBlock = (props: { data: ContentBlockOf<"resource"> }) => {
     return (
       <Popover>
         <PopoverTrigger>
-          <span className="flex items-center gap-1">
+          <span className="flex items-center gap-1 hover:bg-muted rounded-md px-1">
             {props.data.resource.mimeType && (
               <MimeIcon mimeType={props.data.resource.mimeType} />
             )}
             {props.data.resource.uri}
           </span>
         </PopoverTrigger>
-        <PopoverContent className="max-h-96 overflow-y-auto scrollbar-thin">
-          <MarkdownRenderer content={props.data.resource.text} />
+        <PopoverContent className="max-h-96 overflow-y-auto scrollbar-thin whitespace-pre-wrap w-full max-w-[500px]">
+          <span className="text-muted-foreground text-xs mb-1 italic">
+            Formatted for agents, not humans.
+          </span>
+          {props.data.resource.mimeType === "text/plain" ? (
+            <pre className="text-xs whitespace-pre-wrap p-2 bg-muted rounded-md break-words">
+              {props.data.resource.text}
+            </pre>
+          ) : (
+            <MarkdownRenderer content={props.data.resource.text} />
+          )}
         </PopoverContent>
       </Popover>
     );
   }
-
-  if ("blob" in props.data.resource) {
-    if (props.data.resource.mimeType?.startsWith("image/")) {
-      return (
-        <ImageBlock
-          data={{
-            type: "image",
-            mimeType: props.data.resource.mimeType,
-            data: props.data.resource.blob,
-          }}
-        />
-      );
-    }
-    if (props.data.resource.mimeType?.startsWith("audio/")) {
-      return (
-        <AudioBlock
-          data={{
-            type: "audio",
-            mimeType: props.data.resource.mimeType,
-            data: props.data.resource.blob,
-          }}
-        />
-      );
-    }
-    const dataURL = base64ToDataURL(
-      props.data.resource.blob as Base64String,
-      props.data.resource.mimeType ?? "",
-    );
-    return (
-      <a href={dataURL} className="flex items-center gap-1" download={true}>
-        {props.data.resource.mimeType && (
-          <MimeIcon mimeType={props.data.resource.mimeType} />
-        )}
-        {props.data.resource.uri}
-      </a>
-    );
-  }
-  logNever(props.data.resource);
-  return null;
 };
 
 export const ResourceLinkBlock = (props: {
@@ -427,38 +396,62 @@ export const ResourceLinkBlock = (props: {
         href={props.data.uri}
         target="_blank"
         rel="noopener noreferrer"
-        className="text-link hover:underline"
+        className="text-link hover:underline px-1"
       >
         {props.data.name}
       </a>
     );
   }
 
+  // Show image in popover for image mime types
+  if (props.data.mimeType?.startsWith("image/")) {
+    return (
+      <div>
+        <Popover>
+          <PopoverTrigger>
+            <span className="flex items-center gap-1 hover:bg-muted rounded-md px-1 cursor-pointer">
+              <MimeIcon mimeType={props.data.mimeType} />
+              {props.data.name || props.data.title || props.data.uri}
+            </span>
+          </PopoverTrigger>
+          <PopoverContent className="w-auto max-w-[500px] p-2">
+            <img
+              src={props.data.uri}
+              alt={props.data.name || props.data.title || "Image"}
+              className="max-w-full max-h-96 object-contain"
+            />
+          </PopoverContent>
+        </Popover>
+      </div>
+    );
+  }
+
   return (
-    <span className="flex items-center gap-1">
+    <span className="flex items-center gap-1 px-1">
       {props.data.mimeType && <MimeIcon mimeType={props.data.mimeType} />}
       {props.data.name || props.data.title || props.data.uri}
     </span>
   );
 };
 
 export const MimeIcon = (props: { mimeType: string }) => {
+  const classNames = "h-2 w-2 flex-shrink-0";
   if (props.mimeType.startsWith("image/")) {
-    return <FileImageIcon className="h-2 w-2" />;
+    return <FileImageIcon className={classNames} />;
   }
   if (props.mimeType.startsWith("audio/")) {
-    return <FileAudio2Icon className="h-2 w-2" />;
+    return <FileAudio2Icon className={classNames} />;
   }
   if (props.mimeType.startsWith("video/")) {
-    return <FileVideoCameraIcon className="h-2 w-2" />;
+    return <FileVideoCameraIcon className={classNames} />;
   }
   if (props.mimeType.startsWith("text/")) {
-    return <FileTextIcon className="h-2 w-2" />;
+    return <FileTextIcon className={classNames} />;
   }
   if (props.mimeType.startsWith("application/")) {
-    return <FileJsonIcon className="h-2 w-2" />;
+    return <FileJsonIcon className={classNames} />;
   }
-  return <FileIcon className="h-2 w-2" />;
+  return <FileIcon className={classNames} />;
 };
 
 export const SessionNotificationsBlock = <
@@ -533,7 +526,7 @@ export const ToolNotificationsBlock = (props: {
   const toolCalls = mergeToolCalls(props.data);
 
   return (
-    <div className="flex flex-col text-muted-foreground">
+    <div className="flex flex-col text-muted-foreground overflow-x-hidden">
       {toolCalls.map((item) => (
         <SimpleAccordion
           key={item.toolCallId}

frontend/src/components/chat/acp/common.tsx

@@ -68,7 +68,7 @@ export const SimpleAccordion: React.FC<SimpleAccordionProps> = ({
         >
           <span className="flex items-center gap-1">
             {getStatusIcon()}
-            <code className="font-mono text-xs">{title}</code>
+            <code className="font-mono text-xs truncate">{title}</code>
           </span>
         </AccordionTrigger>
         <AccordionContent className="p-2">

frontend/src/components/chat/acp/context-utils.ts

@@ -93,7 +93,7 @@ export async function parseContextFromPrompt(
         type: "resource",
         resource: {
           uri: "context.md",
-          mimeType: "text/markdown",
+          mimeType: "text/plain",
           text: contextString,
         },
       });

frontend/src/components/chat/acp/session-tabs.tsx

@@ -30,7 +30,7 @@ const SessionTab: React.FC<SessionTabProps> = memo(
       <div
         className={cn(
           "flex items-center gap-1 px-2 py-1 text-xs border-r border-border bg-muted/30 hover:bg-muted/50 cursor-pointer min-w-0",
-          isActive && "bg-background border-b-0 relative z-10",
+          isActive && "bg-background border-b-0 relative z-1",
         )}
         onClick={() => onSelect(session.tabId)}
       >

frontend/src/components/editor/ai/__tests__/completion-utils.test.ts

@@ -55,15 +55,13 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="dataset1" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>
-
-      <data name="dataset2" source="unknown">
-      Columns:
-        - col3: boolean
-        - col4: date</data>",
+          "plainText": "<data name="dataset1" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>
+
+      <data name="dataset2" source="unknown">Columns:
+        col3 (boolean)
+        col4 (date)</data>",
           "schema": [],
           "variables": [],
         },
@@ -108,10 +106,9 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="existingDataset" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>",
+          "plainText": "<data name="existingDataset" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>",
           "schema": [],
           "variables": [],
         },
@@ -144,14 +141,12 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="dataset.with.dots" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>
-
-      <data name="regular_dataset" source="unknown">
-      Columns:
-        - col3: boolean</data>",
+          "plainText": "<data name="dataset.with.dots" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>
+
+      <data name="regular_dataset" source="unknown">Columns:
+        col3 (boolean)</data>",
           "schema": [],
           "variables": [],
         },
@@ -198,9 +193,8 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="table1" source="unknown">
-      Columns:
-        - col1: number</data>",
+          "plainText": "<data name="table1" source="unknown">Columns:
+        col1 (number)</data>",
           "schema": [],
           "variables": [],
         },
@@ -276,10 +270,9 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="dataset1" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>
+          "plainText": "<data name="dataset1" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>
 
       <variable name="var1" dataType="string">"string value"</variable>",
           "schema": [],
@@ -346,9 +339,8 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="conflict" source="unknown">
-      Columns:
-        - col1: number</data>",
+          "plainText": "<data name="conflict" source="unknown">Columns:
+        col1 (number)</data>",
           "schema": [],
           "variables": [],
         },

frontend/src/core/ai/context/providers/__tests__/__snapshots__/tables.test.ts.snap

@@ -1,34 +1,28 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
 exports[`TableContextProvider > formatContext > should format context for basic table > basic-table-context 1`] = `
-"<data name="products" source="memory">
-Shape: 100 rows, 3 columns
+"<data name="products" source="memory">Shape: 100 rows, 3 columns
 Columns:
-  - id: integer
-  - name: string
-  - active: boolean</data>"
+  id (integer) - samples: [sample_id_1, sample_id_2]
+  name (string) - samples: [sample_name_1, sample_name_2]
+  active (boolean) - samples: [sample_active_1, sample_active_2]</data>"
 `;
 
 exports[`TableContextProvider > formatContext > should format context for remote database table > remote-table-context 1`] = `
-"<data name="remote_table" source="postgresql://localhost:5432/mydb">
-Shape: 100 rows, 3 columns
+"<data name="remote_table" source="postgresql://localhost:5432/mydb">Shape: 100 rows, 3 columns
 Columns:
-  - uuid: string
-  - created_at: string
-  - metadata: string</data>"
+  uuid (string) - samples: [sample_uuid_1, sample_uuid_2]
+  created_at (string) - samples: [sample_created_at_1, sample_created_at_2]
+  metadata (string) - samples: [sample_metadata_1, sample_metadata_2]</data>"
 `;
 
-exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `
-"<data name="no_columns" source="memory">
-Shape: 100 rows, 3 columns</data>"
-`;
+exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `"<data name="no_columns" source="memory">Shape: 100 rows, 3 columns</data>"`;
 
 exports[`TableContextProvider > formatContext > should format context for table without shape info > no-shape-table-context 1`] = `
-"<data name="no_shape" source="memory">
-Columns:
-  - id: integer
-  - name: string
-  - active: boolean</data>"
+"<data name="no_shape" source="memory">Columns:
+  id (integer) - samples: [sample_id_1, sample_id_2]
+  name (string) - samples: [sample_name_1, sample_name_2]
+  active (boolean) - samples: [sample_active_1, sample_active_2]</data>"
 `;
 
 exports[`TableContextProvider > getItems > should handle dataframe tables with variable names > dataframe-table 1`] = `

frontend/src/core/ai/context/providers/__tests__/cell-output.test.ts

@@ -187,7 +187,6 @@ describe("CellOutputContextProvider", () => {
       expect(completion.displayLabel).toBe(item.data.cellName);
       expect(completion.detail).toContain("output");
       expect(completion.type).toBe("cell-output");
-      expect(completion.section).toBe("Cell Output");
       expect(typeof completion.info).toBe("function");
     });
   });