Convert MCP Resource metadata to atomic resource #951

Author: Polleps

browser/data-browser/src/components/AI/chatConversionUtils.ts

@@ -14,7 +14,14 @@ import type {
   ToolResultPart,
 } from 'ai';
 import { newContextItem } from './AISidebarContext';
-import { type AIChatDisplayMessage, isMessageWithContext } from './types';
+import {
+  type AIAtomicResourceMessageContext,
+  type AIChatDisplayMessage,
+  type AIMCPResourceMessageContext,
+  type AIMessageContext,
+  isAtomicResource,
+  isMessageWithContext,
+} from './types';
 
 // Not exported from 'ai' for some reason, for now we need to define it ourselves.
 type ReasoningPart = {
@@ -40,17 +47,14 @@ export const displayMessageToResource = async (
   message: AIChatDisplayMessage,
   parent: Resource<Ai.AiChat>,
   store: Store,
-  context?: string[],
+  context?: AIMessageContext[],
 ): Promise<Resource<Ai.AiMessage>> => {
   if (isMessageWithContext(message)) {
-    // TODO: Add context to the resource
-    const contextSubjects = message.context.map(c => c.subject);
-
     return displayMessageToResource(
       message.message,
       parent,
       store,
-      contextSubjects,
+      message.context,
     );
   }
 
@@ -63,7 +67,11 @@ export const displayMessageToResource = async (
   });
 
   if (context && context.length > 0) {
-    messageResource.props.providedContext = context;
+    const subjects = await Promise.all(
+      context.map(c => contextToResource(c, messageResource, store)),
+    );
+
+    messageResource.props.providedContext = subjects;
   }
 
   if (typeof message.content === 'string') {
@@ -110,6 +118,33 @@ export const displayMessageToResource = async (
   return messageResource;
 };
 
+const contextToResource = async (
+  context: AIMessageContext,
+  message: Resource<Ai.AiMessage>,
+  store: Store,
+): Promise<string> => {
+  if (isAtomicResource(context)) {
+    return context.subject;
+  }
+
+  const contextResource = await store.newResource<Ai.AiMessage>({
+    isA: ai.classes.mcpResource,
+    parent: message.subject,
+    propVals: {
+      [core.properties.name]: context.name,
+      [ai.properties.mcpUri]: context.uri,
+      [ai.properties.mcpServerId]: context.serverId,
+      ...(context.mimetype
+        ? { [server.properties.mimetype]: context.mimetype }
+        : {}),
+    },
+  });
+
+  contextResource.save();
+
+  return contextResource.subject;
+};
+
 export const messageResourcesToDisplayMessages = async (
   subjects: string[],
   store: Store,
@@ -126,6 +161,7 @@ export const messageResourcesToDisplayMessages = async (
     }
 
     const role = tagToRole(resource.props.role);
+
     const contentResources = await Promise.all(
       resource.props.content.map(s => store.getResource(s)),
     );
@@ -151,15 +187,20 @@ export const messageResourcesToDisplayMessages = async (
       };
 
       if (resource.props.providedContext) {
+        const context = (
+          await Promise.allSettled(
+            resource.props.providedContext.map(c =>
+              resourceToAIMessageContext(c, store),
+            ),
+          )
+        )
+          .filter(c => c.status === 'fulfilled')
+          .map(c => c.value);
+
         message = {
           role: 'annotated-message',
           message,
-          context: resource.props.providedContext.map(c =>
-            newContextItem({
-              subject: c,
-              type: 'resource',
-            }),
-          ),
+          context,
         };
       }
     }
@@ -240,6 +281,32 @@ export const messageResourcesToDisplayMessages = async (
   return messages;
 };
 
+const resourceToAIMessageContext = async (
+  subject: string,
+  store: Store,
+): Promise<AIMessageContext> => {
+  const resource = await store.getResource(subject);
+
+  if (resource.error) {
+    throw resource.error;
+  }
+
+  if (resource.hasClasses(ai.classes.mcpResource)) {
+    return newContextItem<AIMCPResourceMessageContext>({
+      type: 'mcp-resource',
+      name: resource.props.name,
+      uri: resource.props.mcpUri,
+      serverId: resource.props.mcpServerId,
+      mimetype: resource.props.mimetype,
+    });
+  }
+
+  return newContextItem<AIAtomicResourceMessageContext>({
+    type: 'atomic-resource',
+    subject: resource.subject,
+  });
+};
+
 const tagToRole = (subject: string) => {
   const tag = TAG_TO_ROLE_MAPPING[subject as keyof typeof TAG_TO_ROLE_MAPPING];
 

browser/data-browser/src/components/AI/types.ts

@@ -37,6 +37,7 @@ export type AIMCPResourceMessageContext = {
   id: string;
   uri: string;
   name: string;
+  mimetype?: string;
   serverId: string;
 };
 
@@ -77,3 +78,15 @@ export function isMessageWithContext(
 ): message is MessageWithContext {
   return message.role === 'annotated-message';
 }
+
+export function isMCPResource(
+  context: AIMessageContext,
+): context is AIMCPResourceMessageContext {
+  return context.type === 'mcp-resource';
+}
+
+export function isAtomicResource(
+  context: AIMessageContext,
+): context is AIAtomicResourceMessageContext {
+  return context.type === 'atomic-resource';
+}

browser/data-browser/src/components/AI/useProcessMessages.ts

@@ -145,7 +145,7 @@ const processMCPResources = async (
       try {
         const resourceData = await readMCPResource(ctx.serverId, ctx.uri);
 
-        return `MCP resource "${ctx.name}" (${ctx.uri}):\n\`\`\`${resourceData.mimeType || 'text'}
+        return `\`\`\`${resourceData.mimeType || 'text'}
 ${typeof resourceData.contents === 'string' ? resourceData.contents : JSON.stringify(resourceData.contents, null, 2)}
 \`\`\``;
       } catch (error) {
@@ -195,7 +195,7 @@ const addContextToMessage = async (
 
   // Add MCP context if we have any MCP resources
   if (mcpContent) {
-    messageWithContext += `\n<extra-context>\n${mcpContent}\n</extra-context>`;
+    messageWithContext += `\n<context>\n${mcpContent}\n</context>`;
   }
 
   return messageWithContext;

browser/lib/src/ontologies/ai.ts

@@ -27,6 +27,8 @@ export const ai = {
     content:
       'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/content',
     data: 'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/data',
+    mcpServerId:
+      'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/mcp-server-id',
     mcpUri:
       'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/mcp-uri',
     messages:
@@ -35,10 +37,6 @@ export const ai = {
       'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/provided-context',
     reasoningSignature:
       'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/reasoning-signature',
-    resourceBlob:
-      'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/resource-blob',
-    resourceText:
-      'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/resource-text',
     role: 'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/role',
     toolArguments:
       'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/tool-arguments',
@@ -69,9 +67,8 @@ export const ai = {
     ['https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/class/mcp-resource']: [
       'https://atomicdata.dev/properties/name',
       'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/mcp-uri',
+      'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/mcp-server-id',
       'https://atomicdata.dev/properties/mimetype',
-      'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/resource-text',
-      'https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/property/resource-blob',
       'https://atomicdata.dev/properties/description',
     ],
     ['https://atomicdata.dev/01jtjxtsa9syxmfca2zx5gcnmj/class/reasoning-part']:
@@ -133,11 +130,10 @@ declare module '../index.js' {
       requires:
         | BaseProps
         | 'https://atomicdata.dev/properties/name'
-        | typeof ai.properties.mcpUri;
+        | typeof ai.properties.mcpUri
+        | typeof ai.properties.mcpServerId;
       recommends:
         | 'https://atomicdata.dev/properties/mimetype'
-        | typeof ai.properties.resourceText
-        | typeof ai.properties.resourceBlob
         | 'https://atomicdata.dev/properties/description';
     };
     [ai.classes.reasoningPart]: {
@@ -169,12 +165,11 @@ declare module '../index.js' {
   interface PropTypeMapping {
     [ai.properties.content]: string[];
     [ai.properties.data]: string;
+    [ai.properties.mcpServerId]: string;
     [ai.properties.mcpUri]: string;
     [ai.properties.messages]: string[];
     [ai.properties.providedContext]: string[];
     [ai.properties.reasoningSignature]: string;
-    [ai.properties.resourceBlob]: string;
-    [ai.properties.resourceText]: string;
     [ai.properties.role]: string;
     [ai.properties.toolArguments]: string;
     [ai.properties.toolId]: string;
@@ -186,12 +181,11 @@ declare module '../index.js' {
   interface PropSubjectToNameMapping {
     [ai.properties.content]: 'content';
     [ai.properties.data]: 'data';
+    [ai.properties.mcpServerId]: 'mcpServerId';
     [ai.properties.mcpUri]: 'mcpUri';
     [ai.properties.messages]: 'messages';
     [ai.properties.providedContext]: 'providedContext';
     [ai.properties.reasoningSignature]: 'reasoningSignature';
-    [ai.properties.resourceBlob]: 'resourceBlob';
-    [ai.properties.resourceText]: 'resourceText';
     [ai.properties.role]: 'role';
     [ai.properties.toolArguments]: 'toolArguments';
     [ai.properties.toolId]: 'toolId';