@dxos/plugin-assistant 0.8.4-main.ae835ea → 0.8.4-main.bc674ce

package/src/capabilities/app-graph-builder/app-graph-builder.ts

@@ -0,0 +1,180 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+
+import { Capability, Common } from '@dxos/app-framework';
+import { Prompt } from '@dxos/blueprints';
+import { Sequence } from '@dxos/conductor';
+import { DXN, type Database, Obj, type Ref } from '@dxos/echo';
+import { AtomObj } from '@dxos/echo-atom';
+import { invariant } from '@dxos/invariant';
+import { Operation, type OperationInvoker } from '@dxos/operation';
+import { ClientCapabilities } from '@dxos/plugin-client';
+import { ATTENDABLE_PATH_SEPARATOR, PLANK_COMPANION_TYPE } from '@dxos/plugin-deck/types';
+import { GraphBuilder, NodeMatcher } from '@dxos/plugin-graph';
+import { getActiveSpace } from '@dxos/plugin-space';
+import { SpaceOperation } from '@dxos/plugin-space/types';
+import { Query } from '@dxos/react-client/echo';
+
+import { ASSISTANT_DIALOG, meta } from '../../meta';
+import { Assistant, AssistantCapabilities, AssistantOperation } from '../../types';
+
+export default Capability.makeModule(
+  Effect.fnUntraced(function* () {
+    const capabilities = yield* Capability.Service;
+
+    const extensions = yield* Effect.all([
+      GraphBuilder.createTypeExtension({
+        id: `${meta.id}/root`,
+        type: Assistant.Chat,
+        actions: (chat) => {
+          const id = Obj.getDXN(chat).toString();
+          return Effect.succeed([
+            {
+              id: `${AssistantOperation.UpdateChatName.meta.key}/${id}`,
+              data: () => Operation.invoke(AssistantOperation.UpdateChatName, { chat }),
+              properties: {
+                label: ['chat update name label', { ns: meta.id }],
+                icon: 'ph--magic-wand--regular',
+                disposition: 'list-item',
+              },
+            },
+          ]);
+        },
+      }),
+
+      GraphBuilder.createExtension({
+        id: `${meta.id}/assistant`,
+        match: NodeMatcher.whenRoot,
+        actions: () =>
+          Effect.succeed([
+            {
+              id: `${Common.LayoutOperation.UpdateDialog.meta.key}/assistant/open`,
+              data: Effect.fnUntraced(function* () {
+                const capabilities = yield* Capability.Service;
+                const client = yield* Capability.get(ClientCapabilities.Client);
+                const operationInvoker = yield* Capability.get(Common.Capability.OperationInvoker);
+                const space = getActiveSpace(capabilities) ?? client.spaces.default;
+                const chat = yield* Effect.tryPromise(() => getOrCreateChat(operationInvoker.invokePromise, space.db));
+                if (!chat) {
+                  return;
+                }
+
+                yield* Operation.invoke(Common.LayoutOperation.UpdateDialog, {
+                  subject: ASSISTANT_DIALOG,
+                  state: true,
+                  blockAlign: 'end',
+                  props: {
+                    chat,
+                  },
+                });
+              }),
+              properties: {
+                label: ['open assistant label', { ns: meta.id }],
+                icon: 'ph--sparkle--regular',
+                disposition: 'pin-end',
+                position: 'hoist',
+                keyBinding: {
+                  macos: 'shift+meta+k',
+                  windows: 'shift+ctrl+k',
+                },
+              },
+            },
+          ]),
+      }),
+
+      GraphBuilder.createExtension({
+        id: `${meta.id}/companion-chat`,
+        match: NodeMatcher.whenEchoObject,
+        connector: (object, get) =>
+          Effect.gen(function* () {
+            const state = get(yield* Capability.get(AssistantCapabilities.State));
+            const currentChatState = state.currentChat[Obj.getDXN(object).toString()];
+            // If no state, continue to allow chat initialization.
+            if (!currentChatState) {
+              return [
+                {
+                  id: [Obj.getDXN(object).toString(), 'assistant-chat'].join(ATTENDABLE_PATH_SEPARATOR),
+                  type: PLANK_COMPANION_TYPE,
+                  data: 'assistant-chat',
+                  properties: {
+                    label: ['assistant chat label', { ns: meta.id }],
+                    icon: 'ph--sparkle--regular',
+                    position: 'hoist',
+                    disposition: 'hidden',
+                  },
+                },
+              ];
+            }
+
+            const db = Obj.getDatabase(object);
+            const currentChatDxn = DXN.tryParse(currentChatState);
+            const currentChatRef = currentChatDxn ? db?.makeRef(currentChatDxn) : undefined;
+            const currentChat = currentChatRef ? get(AtomObj.make(currentChatRef as Ref.Ref<Obj.Unknown>)) : undefined;
+            if (!Obj.isObject(currentChat)) {
+              return [];
+            }
+
+            return [
+              {
+                id: [Obj.getDXN(object).toString(), 'assistant-chat'].join(ATTENDABLE_PATH_SEPARATOR),
+                type: PLANK_COMPANION_TYPE,
+                data: currentChat,
+                properties: {
+                  label: ['assistant chat label', { ns: meta.id }],
+                  icon: 'ph--sparkle--regular',
+                  position: 'hoist',
+                  disposition: 'hidden',
+                },
+              },
+            ];
+          }),
+      }),
+
+      GraphBuilder.createExtension({
+        id: `${meta.id}/invocations`,
+        match: NodeMatcher.whenAny(
+          NodeMatcher.whenEchoTypeMatches(Sequence),
+          NodeMatcher.whenEchoTypeMatches(Prompt.Prompt),
+        ),
+        connector: (node) =>
+          Effect.succeed([
+            {
+              id: [node.id, 'invocations'].join(ATTENDABLE_PATH_SEPARATOR),
+              type: PLANK_COMPANION_TYPE,
+              data: 'invocations',
+              properties: {
+                label: ['invocations label', { ns: meta.id }],
+                icon: 'ph--clock-countdown--regular',
+                disposition: 'hidden',
+              },
+            },
+          ]),
+      }),
+    ]);
+
+    return Capability.contributes(Common.Capability.AppGraphBuilder, extensions);
+  }),
+);
+
+// TODO(burdon): Factor out.
+const getOrCreateChat = async (
+  invokePromise: OperationInvoker.OperationInvoker['invokePromise'],
+  db: Database.Database,
+): Promise<Assistant.Chat | undefined> => {
+  // TODO(wittjosiah): This should be possible with a single query.
+  const allChats = await db.query(Query.type(Assistant.Chat)).run();
+  const relatedChats = await db.query(Query.type(Assistant.Chat).sourceOf(Assistant.CompanionTo).source()).run();
+
+  const chats = allChats.filter((chat) => !relatedChats.includes(chat));
+  if (chats.length > 0) {
+    return chats.at(-1);
+  }
+
+  const { data } = await invokePromise(AssistantOperation.CreateChat, { db });
+  invariant(Obj.instanceOf(Assistant.Chat, data?.object));
+  await invokePromise(SpaceOperation.AddObject, { target: db, object: data.object });
+  return data.object;
+};

package/src/capabilities/app-graph-builder/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const AppGraphBuilder = Capability.lazy('AppGraphBuilder', () => import('./app-graph-builder'));

package/src/capabilities/blueprint-definition/blueprint-definition.ts

@@ -0,0 +1,60 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+
+import { Capability, Common } from '@dxos/app-framework';
+import {
+  Agent,
+  Discord,
+  DiscordBlueprint,
+  EntityExtraction,
+  Linear,
+  LinearBlueprint,
+  Research,
+  ResearchBlueprint,
+  WebSearchBlueprint,
+} from '@dxos/assistant-toolkit';
+import { type Blueprint } from '@dxos/blueprints';
+
+import { AssistantBlueprint } from '../../blueprints';
+
+export const createBlueprint: () => Blueprint.Blueprint = AssistantBlueprint.make;
+export { AssistantBlueprint };
+
+export type BlueprintCapabilities = [
+  Capability.Capability<typeof Common.Capability.Functions>,
+  Capability.Capability<typeof Common.Capability.BlueprintDefinition>,
+  Capability.Capability<typeof Common.Capability.Functions>,
+  Capability.Capability<typeof Common.Capability.BlueprintDefinition>,
+  Capability.Capability<typeof Common.Capability.Functions>,
+  Capability.Capability<typeof Common.Capability.BlueprintDefinition>,
+  Capability.Capability<typeof Common.Capability.Functions>,
+  Capability.Capability<typeof Common.Capability.BlueprintDefinition>,
+  Capability.Capability<typeof Common.Capability.Functions>,
+  Capability.Capability<typeof Common.Capability.BlueprintDefinition>,
+];
+
+export default Capability.makeModule<[], BlueprintCapabilities>(() =>
+  Effect.succeed([
+    Capability.contributes(Common.Capability.Functions, AssistantBlueprint.functions$),
+    Capability.contributes(Common.Capability.BlueprintDefinition, AssistantBlueprint.blueprint),
+
+    // TODO(burdon): Factor out.
+    Capability.contributes(Common.Capability.Functions, [Research.create, Research.research]),
+    Capability.contributes(Common.Capability.BlueprintDefinition, ResearchBlueprint),
+
+    // TODO(burdon): Factor out.
+    Capability.contributes(Common.Capability.Functions, [Agent.prompt, EntityExtraction.extract]),
+    Capability.contributes(Common.Capability.BlueprintDefinition, WebSearchBlueprint),
+
+    // TODO(burdon): Factor out.
+    Capability.contributes(Common.Capability.Functions, [Discord.fetch]),
+    Capability.contributes(Common.Capability.BlueprintDefinition, DiscordBlueprint),
+
+    // TODO(burdon): Factor out.
+    Capability.contributes(Common.Capability.Functions, [Linear.sync]),
+    Capability.contributes(Common.Capability.BlueprintDefinition, LinearBlueprint),
+  ]),
+);

package/src/capabilities/blueprint-definition/index.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const BlueprintDefinition = Capability.lazy('BlueprintDefinition', () => import('./blueprint-definition'));
+export { AssistantBlueprint } from './blueprint-definition';

package/src/capabilities/edge-model-resolver/edge-model-resolver.ts

@@ -0,0 +1,32 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as AnthropicClient from '@effect/ai-anthropic/AnthropicClient';
+import * as FetchHttpClient from '@effect/platform/FetchHttpClient';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+
+import { AnthropicResolver } from '@dxos/ai/resolvers';
+import { Capability, Common } from '@dxos/app-framework';
+
+export type EdgeModelResolverCapabilities = [Capability.Capability<typeof Common.Capability.AiModelResolver>];
+
+const edgeModelResolver = Capability.makeModule<[], EdgeModelResolverCapabilities>(() =>
+  Effect.succeed([
+    Capability.contributes(
+      Common.Capability.AiModelResolver,
+      AnthropicResolver.make().pipe(
+        Layer.provide(
+          AnthropicClient.layer({
+            // TODO(dmaretskyi): Read endpoint from config/settings.
+            apiUrl: 'https://ai-service.dxos.workers.dev/provider/anthropic',
+          }),
+        ),
+        Layer.provide(FetchHttpClient.layer),
+      ),
+    ),
+  ]),
+);
+
+export default edgeModelResolver;

package/src/capabilities/edge-model-resolver/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const EdgeModelResolver = Capability.lazy('EdgeModelResolver', () => import('./edge-model-resolver'));

package/src/capabilities/index.ts

@@ -2,18 +2,14 @@
 // Copyright 2025 DXOS.org
 //
 
-import { lazy } from '@dxos/app-framework';
-
-export const AiService = lazy(() => import('./ai-service'));
-export const AppGraphBuilder = lazy(() => import('./app-graph-builder'));
-export const AssistantState = lazy(() => import('./state'));
-export const BlueprintDefinition = lazy(() => import('./blueprint-definition'));
-export const EdgeModelResolver = lazy(() => import('./edge-model-resolver'));
-export const IntentResolver = lazy(() => import('./intent-resolver'));
-export const LocalModelResolver = lazy(() => import('./local-model-resolver'));
-export const ReactSurface = lazy(() => import('./react-surface'));
-export const Settings = lazy(() => import('./settings'));
-export const Toolkit = lazy(() => import('./toolkit'));
-
+export * from './ai-service';
+export * from './app-graph-builder';
 export * from './blueprint-definition';
-export * from './capabilities';
+export * from './edge-model-resolver';
+export * from './local-model-resolver';
+export * from './operation-resolver';
+export * from './react-surface';
+export * from './repair';
+export * from './settings';
+export * from './state';
+export * from './toolkit';

package/src/capabilities/local-model-resolver/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const LocalModelResolver = Capability.lazy('LocalModelResolver', () => import('./local-model-resolver'));

package/src/capabilities/local-model-resolver/local-model-resolver.ts

@@ -0,0 +1,37 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as OpenAiClient from '@effect/ai-openai/OpenAiClient';
+import * as FetchHttpClient from '@effect/platform/FetchHttpClient';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+
+import { LMStudioResolver } from '@dxos/ai/resolvers';
+import { Capability, Common } from '@dxos/app-framework';
+
+export type LocalModelResolverCapabilities = [Capability.Capability<typeof Common.Capability.AiModelResolver>];
+
+/**
+ * To start LM Studio server:
+ * ```bash
+ * ~/.lmstudio/bin/lms server start --cors
+ * ```
+ */
+const localModelResolver = Capability.makeModule<[], LocalModelResolverCapabilities>(() =>
+  Effect.succeed([
+    Capability.contributes(
+      Common.Capability.AiModelResolver,
+      LMStudioResolver.make().pipe(
+        Layer.provide(
+          OpenAiClient.layer({
+            apiUrl: LMStudioResolver.DEFAULT_LMSTUDIO_ENDPOINT,
+          }),
+        ),
+        Layer.provide(FetchHttpClient.layer),
+      ),
+    ),
+  ]),
+);
+
+export default localModelResolver;

package/src/capabilities/operation-resolver/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const OperationResolver = Capability.lazy('OperationResolver', () => import('./operation-resolver'));

package/src/capabilities/operation-resolver/operation-resolver.ts

@@ -0,0 +1,110 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+
+import { Capability, Common } from '@dxos/app-framework';
+import { AiContextBinder, AiConversation } from '@dxos/assistant';
+import { Agent } from '@dxos/assistant-toolkit';
+import { Blueprint, Prompt } from '@dxos/blueprints';
+import { type Queue } from '@dxos/client/echo';
+import { Filter, Obj, Ref, Type } from '@dxos/echo';
+import { TracingService, serializeFunction } from '@dxos/functions';
+import { invariant } from '@dxos/invariant';
+import { Operation, OperationResolver } from '@dxos/operation';
+import { AutomationCapabilities } from '@dxos/plugin-automation';
+import { ClientCapabilities } from '@dxos/plugin-client';
+import { Collection } from '@dxos/schema';
+import { type Message } from '@dxos/types';
+
+import { type AiChatServices, updateName } from '../../processor';
+import { Assistant, AssistantCapabilities, AssistantOperation } from '../../types';
+import { AssistantBlueprint, createBlueprint } from '../blueprint-definition/blueprint-definition';
+
+export default Capability.makeModule(
+  Effect.fnUntraced(function* () {
+    return Capability.contributes(Common.Capability.OperationResolver, [
+      OperationResolver.make({
+        operation: AssistantOperation.OnCreateSpace,
+        handler: Effect.fnUntraced(function* ({ space, rootCollection }) {
+          const chatCollection = Collection.makeManaged({ key: Assistant.Chat.typename });
+          const blueprintCollection = Collection.makeManaged({ key: Blueprint.Blueprint.typename });
+          const promptCollection = Collection.makeManaged({ key: Type.getTypename(Prompt.Prompt) });
+          Obj.change(rootCollection, (c) => {
+            c.objects.push(Ref.make(chatCollection), Ref.make(blueprintCollection), Ref.make(promptCollection));
+          });
+
+          // TODO(wittjosiah): Remove once function registry is avaiable.
+          space.db.add(serializeFunction(Agent.prompt));
+
+          // Create default chat.
+          const { object: chat } = yield* Operation.invoke(AssistantOperation.CreateChat, { db: space.db });
+          space.db.add(chat);
+        }),
+      }),
+      OperationResolver.make({
+        operation: AssistantOperation.CreateChat,
+        handler: Effect.fnUntraced(function* ({ db, name }) {
+          const registry = yield* Capability.get(Common.Capability.AtomRegistry);
+          const client = yield* Capability.get(ClientCapabilities.Client);
+          const space = client.spaces.get(db.spaceId);
+          invariant(space, 'Space not found');
+          const queue = space.queues.create();
+          const chat = Assistant.make({ name, queue: Ref.fromDXN(queue.dxn) });
+          space.db.add(chat);
+
+          // TODO(wittjosiah): This should be a space-level setting.
+          // TODO(burdon): Clone when activated. Copy-on-write for template.
+          const blueprints = yield* Effect.promise(() => db.query(Filter.type(Blueprint.Blueprint)).run());
+          let defaultBlueprint = blueprints.find((blueprint) => blueprint.key === AssistantBlueprint.Key);
+          if (!defaultBlueprint) {
+            defaultBlueprint = db.add(createBlueprint());
+          }
+
+          const binder = new AiContextBinder({ queue, registry });
+          yield* Effect.promise(() =>
+            binder.use((b: AiContextBinder) => b.bind({ blueprints: [Ref.make(defaultBlueprint!)] })),
+          );
+
+          return { object: chat };
+        }),
+      }),
+      OperationResolver.make({
+        operation: AssistantOperation.UpdateChatName,
+        handler: Effect.fnUntraced(function* ({ chat }) {
+          const registry = yield* Capability.get(Common.Capability.AtomRegistry);
+          const db = Obj.getDatabase(chat);
+          const queue = chat.queue.target as Queue<Message.Message>;
+          if (!db || !queue) {
+            return;
+          }
+
+          const runtimeResolver = yield* Capability.get(AutomationCapabilities.ComputeRuntime);
+          const runtime = yield* Effect.promise(() =>
+            runtimeResolver
+              .getRuntime(db.spaceId)
+              .runPromise(Effect.runtime<AiChatServices>().pipe(Effect.provide(TracingService.layerNoop))),
+          );
+
+          yield* Effect.promise(() =>
+            new AiConversation({ queue, registry }).use(async (conversation) =>
+              updateName(runtime, conversation, chat),
+            ),
+          );
+        }),
+      }),
+      OperationResolver.make({
+        operation: AssistantOperation.SetCurrentChat,
+        handler: Effect.fnUntraced(function* ({ companionTo, chat }) {
+          const companionToId = Obj.getDXN(companionTo).toString();
+          const chatId = chat && Obj.getDXN(chat).toString();
+          yield* Common.Capability.updateAtomValue(AssistantCapabilities.State, (current) => ({
+            ...current,
+            currentChat: { ...current.currentChat, [companionToId]: chatId },
+          }));
+        }),
+      }),
+    ]);
+  }),
+);

package/src/capabilities/react-surface/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const ReactSurface = Capability.lazy('ReactSurface', () => import('./react-surface'));

package/src/capabilities/react-surface/react-surface.tsx

@@ -0,0 +1,95 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+import React from 'react';
+
+import { Capability, Common } from '@dxos/app-framework';
+import { useSettingsState } from '@dxos/app-framework/react';
+import { Blueprint, Prompt } from '@dxos/blueprints';
+import { getSpace } from '@dxos/client/echo';
+import { Sequence } from '@dxos/conductor';
+import { InvocationTraceContainer } from '@dxos/devtools';
+import { Obj } from '@dxos/echo';
+import { Layout } from '@dxos/react-ui-mosaic';
+
+import {
+  AssistantSettings,
+  BlueprintArticle,
+  ChatCompanion,
+  ChatContainer,
+  ChatDialog,
+  PromptArticle,
+} from '../../components';
+import { ASSISTANT_DIALOG, meta } from '../../meta';
+import { Assistant } from '../../types';
+
+export default Capability.makeModule(() =>
+  Effect.succeed(
+    Capability.contributes(Common.Capability.ReactSurface, [
+      Common.createSurface({
+        id: `${meta.id}/plugin-settings`,
+        role: 'article',
+        filter: (data): data is { subject: Common.Capability.Settings } =>
+          Common.Capability.isSettings(data.subject) && data.subject.prefix === meta.id,
+        component: ({ data: { subject } }) => {
+          const { settings, updateSettings } = useSettingsState<Assistant.Settings>(subject.atom);
+          return <AssistantSettings settings={settings} onSettingsChange={updateSettings} />;
+        },
+      }),
+      Common.createSurface({
+        id: `${meta.id}/chat`,
+        role: 'article',
+        filter: (data): data is { subject: Assistant.Chat; variant: undefined } =>
+          Obj.instanceOf(Assistant.Chat, data.subject) && data.variant !== 'assistant-chat',
+        component: ({ data, role, ref }) => <ChatContainer role={role} subject={data.subject} ref={ref} />,
+      }),
+      // TODO(wittjosiah): This is flashing when chat changes.
+      Common.createSurface({
+        id: `${meta.id}/companion-chat`,
+        role: 'article',
+        filter: (data): data is { companionTo: Obj.Unknown; subject: Assistant.Chat | 'assistant-chat' } =>
+          Obj.isObject(data.companionTo) &&
+          (Obj.instanceOf(Assistant.Chat, data.subject) || data.subject === 'assistant-chat'),
+        component: ({ data, role, ref }) => <ChatCompanion role={role} data={data} ref={ref} />,
+      }),
+      Common.createSurface({
+        id: `${meta.id}/companion-invocations`,
+        role: 'article',
+        filter: (data): data is { companionTo: Sequence } =>
+          (Obj.instanceOf(Sequence, data.companionTo) || Obj.instanceOf(Prompt.Prompt, data.companionTo)) &&
+          data.subject === 'invocations',
+        component: ({ data, role }) => {
+          const space = getSpace(data.companionTo);
+          const queueDxn = space?.properties.invocationTraceQueue?.dxn;
+          // TODO(wittjosiah): Support invocation filtering for prompts.
+          const target = Obj.instanceOf(Prompt.Prompt, data.companionTo) ? undefined : data.companionTo;
+          return (
+            <Layout.Main role={role}>
+              <InvocationTraceContainer db={space?.db} queueDxn={queueDxn} target={target} detailAxis='block' />
+            </Layout.Main>
+          );
+        },
+      }),
+      Common.createSurface({
+        id: `${meta.id}/blueprint`,
+        role: 'article',
+        filter: (data): data is { subject: Blueprint.Blueprint } => Obj.instanceOf(Blueprint.Blueprint, data.subject),
+        component: ({ data }) => <BlueprintArticle subject={data.subject} />,
+      }),
+      Common.createSurface({
+        id: `${meta.id}/prompt`,
+        role: 'article',
+        filter: (data): data is { subject: Prompt.Prompt } => Obj.instanceOf(Prompt.Prompt, data.subject),
+        component: ({ data }) => <PromptArticle subject={data.subject} />,
+      }),
+      Common.createSurface({
+        id: ASSISTANT_DIALOG,
+        role: 'dialog',
+        filter: (data): data is { props: { chat: Assistant.Chat } } => data.component === ASSISTANT_DIALOG,
+        component: ({ data }) => <ChatDialog {...data.props} />,
+      }),
+    ]),
+  ),
+);

package/src/capabilities/repair/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Capability } from '@dxos/app-framework';
+
+export const Repair = Capability.lazy('Repair', () => import('./repair'));

package/src/capabilities/repair/repair.ts

@@ -0,0 +1,63 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+
+import { Capability } from '@dxos/app-framework';
+import { Blueprint, Prompt } from '@dxos/blueprints';
+import { Obj, Ref, Type } from '@dxos/echo';
+import { SpaceCapabilities } from '@dxos/plugin-space';
+import { type Space } from '@dxos/react-client/echo';
+import { Collection } from '@dxos/schema';
+
+import { Assistant } from '../../types';
+
+export default Capability.makeModule(() =>
+  Effect.succeed(
+    Capability.contributes(SpaceCapabilities.Repair, async ({ space }) => {
+      await ensureSystemCollections(space);
+    }),
+  ),
+);
+
+/**
+ * Ensure the root collection has system collections for AI Chats, Blueprints, and Prompts.
+ */
+const ensureSystemCollections = async (space: Space) => {
+  const rootCollection: Collection.Collection = await space.properties[Collection.Collection.typename]?.load();
+  if (!rootCollection) {
+    return;
+  }
+
+  const objects = await Promise.all(rootCollection.objects.map((ref) => ref.load()));
+  const chats = objects.find(
+    (object) => Obj.instanceOf(Collection.Managed, object) && object.key === Assistant.Chat.typename,
+  );
+  if (!chats) {
+    const chatsCollectionRef = Ref.make(Collection.makeManaged({ key: Assistant.Chat.typename }));
+    Obj.change(rootCollection, (c) => {
+      c.objects.push(chatsCollectionRef);
+    });
+  }
+
+  const blueprints = objects.find(
+    (object) => Obj.instanceOf(Collection.Managed, object) && object.key === Blueprint.Blueprint.typename,
+  );
+  if (!blueprints) {
+    const blueprintsCollectionRef = Ref.make(Collection.makeManaged({ key: Blueprint.Blueprint.typename }));
+    Obj.change(rootCollection, (c) => {
+      c.objects.push(blueprintsCollectionRef);
+    });
+  }
+
+  const prompts = objects.find(
+    (object) => Obj.instanceOf(Collection.Managed, object) && object.key === Type.getTypename(Prompt.Prompt),
+  );
+  if (!prompts) {
+    const promptsCollectionRef = Ref.make(Collection.makeManaged({ key: Type.getTypename(Prompt.Prompt) }));
+    Obj.change(rootCollection, (c) => {
+      c.objects.push(promptsCollectionRef);
+    });
+  }
+};