Support all content types in prompts from Zed (#6756)

packages/cli/src/zed-integration/schema.ts

@@ -84,6 +84,8 @@ export type AgentCapabilities = z.infer<typeof agentCapabilitiesSchema>;
 
 export type AuthMethod = z.infer<typeof authMethodSchema>;
 
+export type PromptCapabilities = z.infer<typeof promptCapabilitiesSchema>;
+
 export type ClientResponse = z.infer<typeof clientResponseSchema>;
 
 export type ClientNotification = z.infer<typeof clientNotificationSchema>;
@@ -270,8 +272,15 @@ export const mcpServerSchema = z.object({
   name: z.string(),
 });
 
+export const promptCapabilitiesSchema = z.object({
+  audio: z.boolean().optional(),
+  embeddedContext: z.boolean().optional(),
+  image: z.boolean().optional(),
+});
+
 export const agentCapabilitiesSchema = z.object({
-  loadSession: z.boolean(),
+  loadSession: z.boolean().optional(),
+  promptCapabilities: promptCapabilitiesSchema.optional(),
 });
 
 export const authMethodSchema = z.object({

packages/cli/src/zed-integration/zedIntegration.ts

@@ -99,6 +99,11 @@ class GeminiAgent {
       authMethods,
       agentCapabilities: {
         loadSession: false,
+        promptCapabilities: {
+          image: true,
+          audio: true,
+          embeddedContext: true,
+        },
       },
     };
   }
@@ -488,45 +493,59 @@ class Session {
     message: acp.ContentBlock[],
     abortSignal: AbortSignal,
   ): Promise<Part[]> {
+    const FILE_URI_SCHEME = 'file://';
+
+    const embeddedContext: acp.EmbeddedResourceResource[] = [];
+
     const parts = message.map((part) => {
       switch (part.type) {
         case 'text':
           return { text: part.text };
-        case 'resource_link':
+        case 'image':
+        case 'audio':
+          return {
+            inlineData: {
+              mimeType: part.mimeType,
+              data: part.data,
+            },
+          };
+        case 'resource_link': {
+          if (part.uri.startsWith(FILE_URI_SCHEME)) {
             return {
               fileData: {
                 mimeData: part.mimeType,
                 name: part.name,
-              fileUri: part.uri,
+                fileUri: part.uri.slice(FILE_URI_SCHEME.length),
               },
             };
+          } else {
+            return { text: `@${part.uri}` };
+          }
+        }
         case 'resource': {
-          return {
+          embeddedContext.push(part.resource);
-            fileData: {
+          return { text: `@${part.resource.uri}` };
-              mimeData: part.resource.mimeType,
-              name: part.resource.uri,
-              fileUri: part.resource.uri,
-            },
-          };
         }
         default: {
-          throw new Error(`Unexpected chunk type: '${part.type}'`);
+          const unreachable: never = part;
+          throw new Error(`Unexpected chunk type: '${unreachable}'`);
         }
       }
     });
 
     const atPathCommandParts = parts.filter((part) => 'fileData' in part);
 
-    if (atPathCommandParts.length === 0) {
+    if (atPathCommandParts.length === 0 && embeddedContext.length === 0) {
       return parts;
     }
 
+    const atPathToResolvedSpecMap = new Map<string, string>();
+
     // Get centralized file discovery service
     const fileDiscovery = this.config.getFileService();
     const respectGitIgnore = this.config.getFileFilteringRespectGitIgnore();
 
     const pathSpecsToRead: string[] = [];
-    const atPathToResolvedSpecMap = new Map<string, string>();
     const contentLabelsForDisplay: string[] = [];
     const ignoredPaths: string[] = [];
 
@@ -634,6 +653,7 @@ class Session {
         contentLabelsForDisplay.push(pathName);
       }
     }
+
     // Construct the initial part of the query for the LLM
     let initialQueryText = '';
     for (let i = 0; i < parts.length; i++) {
@@ -687,12 +707,16 @@ class Session {
         `Ignored ${ignoredPaths.length} ${ignoreType} files: ${ignoredPaths.join(', ')}`,
       );
     }
+
+    const processedQueryParts: Part[] = [{ text: initialQueryText }];
+
+    if (pathSpecsToRead.length === 0 && embeddedContext.length === 0) {
       // Fallback for lone "@" or completely invalid @-commands resulting in empty initialQueryText
-    if (pathSpecsToRead.length === 0) {
       console.warn('No valid file paths found in @ commands to read.');
       return [{ text: initialQueryText }];
     }
-    const processedQueryParts: Part[] = [{ text: initialQueryText }];
+
+    if (pathSpecsToRead.length > 0) {
       const toolArgs = {
         paths: pathSpecsToRead,
         respectGitIgnore, // Use configuration setting
@@ -755,7 +779,6 @@ class Session {
             'read_many_files tool returned no content or empty content.',
           );
         }
-      return processedQueryParts;
       } catch (error: unknown) {
         await this.sendUpdate({
           sessionUpdate: 'tool_call_update',
@@ -776,6 +799,33 @@ class Session {
       }
     }
 
+    if (embeddedContext.length > 0) {
+      processedQueryParts.push({
+        text: '\n--- Content from referenced context ---',
+      });
+
+      for (const contextPart of embeddedContext) {
+        processedQueryParts.push({
+          text: `\nContent from @${contextPart.uri}:\n`,
+        });
+        if ('text' in contextPart) {
+          processedQueryParts.push({
+            text: contextPart.text,
+          });
+        } else {
+          processedQueryParts.push({
+            inlineData: {
+              mimeType: contextPart.mimeType ?? 'application/octet-stream',
+              data: contextPart.blob,
+            },
+          });
+        }
+      }
+    }
+
+    return processedQueryParts;
+  }
+
   debug(msg: string) {
     if (this.config.getDebugMode()) {
       console.warn(msg);
@@ -784,6 +834,10 @@ class Session {
 }
 
 function toToolCallContent(toolResult: ToolResult): acp.ToolCallContent | null {
+  if (toolResult.error?.message) {
+    throw new Error(toolResult.error.message);
+  }
+
   if (toolResult.returnDisplay) {
     if (typeof toolResult.returnDisplay === 'string') {
       return {