@dxos/plugin-assistant 0.8.4-main.406dc2a → 0.8.4-main.52d7546f51

package/src/components/ChatCompanion.tsx

@@ -3,134 +3,224 @@
 //
 
 import * as Array from 'effect/Array';
-import * as Effect from 'effect/Effect';
 import * as Function from 'effect/Function';
 import * as Option from 'effect/Option';
-import React, { useMemo } from 'react';
+import React, { forwardRef, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
-import { Capabilities, createIntent, useCapabilities, useIntentDispatcher } from '@dxos/app-framework';
+import { useAtomCapability, useCapabilities, useOperationInvoker } from '@dxos/app-framework/ui';
+import { AppCapabilities } from '@dxos/app-toolkit';
+import { Chat } from '@dxos/assistant-toolkit';
 import { Blueprint } from '@dxos/blueprints';
 import { getSpace } from '@dxos/client/echo';
-import { Filter, Obj, Query, Ref } from '@dxos/echo';
-import { DatabaseService } from '@dxos/functions';
-import { SpaceAction } from '@dxos/plugin-space/types';
+import { DXN, Filter, Obj, Query, Ref } from '@dxos/echo';
+import { SpaceOperation } from '@dxos/plugin-space/types';
 import { useQuery } from '@dxos/react-client/echo';
 import { useAsyncEffect } from '@dxos/react-ui';
 
-import { ChatContainer } from '../components';
+import { ChatContainer, type ChatEvent } from '../components';
 import { useBlueprintRegistry, useContextBinder } from '../hooks';
-import { Assistant, AssistantAction } from '../types';
+import { Assistant, AssistantCapabilities, AssistantOperation } from '../types';
 
 export type ChatCompanionProps = {
   role?: string;
-  data: { companionTo: Obj.Any; subject: Assistant.Chat | 'assistant-chat' };
+  data: {
+    subject: Chat.Chat | 'assistant-chat';
+    companionTo: Obj.Unknown;
+  };
 };
 
-export const ChatCompanion = ({ role, data }: ChatCompanionProps) => {
-  const companionTo = data.companionTo;
-  const space = getSpace(companionTo);
-  const chat = data.subject === 'assistant-chat' ? undefined : data.subject;
-  const blueprintRegistry = useBlueprintRegistry();
-  const binder = useContextBinder(chat);
-  const { dispatch } = useIntentDispatcher();
-
-  // Initialize companion chat if it doesn't exist.
-  useAsyncEffect(async () => {
-    if (chat || !space) {
-      return;
+export const ChatCompanion = forwardRef<HTMLDivElement, ChatCompanionProps>(
+  ({ role, data }: ChatCompanionProps, forwardedRef) => {
+    const { invokePromise } = useOperationInvoker();
+    const blueprintRegistry = useBlueprintRegistry();
+    const companionTo = data.companionTo;
+
+    const space = getSpace(companionTo);
+    const [chat, setChat] = useState(data.subject === 'assistant-chat' ? undefined : data.subject);
+
+    // Watch the state atom directly to detect when plus button clears the chat selection.
+    const state = useAtomCapability(AssistantCapabilities.State);
+    const companionToId = Obj.getDXN(companionTo).toString();
+    const currentChatState = state?.currentChat[companionToId];
+
+    // Track if initial setup has been done (to distinguish initial mount from plus button click).
+    // Reset when companionTo changes.
+    const initialSetupDoneRef = useRef(false);
+    const prevCompanionToIdRef = useRef(companionToId);
+    if (prevCompanionToIdRef.current !== companionToId) {
+      initialSetupDoneRef.current = false;
+      prevCompanionToIdRef.current = companionToId;
     }
 
-    // TODO(burdon): Garbage collection of queues?
-    // TODO(wittjosiah): Figure out how to prevent companion chats from showing up in the navtree.
-    await Effect.gen(function* () {
-      const { objects } = yield* DatabaseService.runQuery(
-        Query.select(Filter.ids(companionTo.id)).targetOf(Assistant.CompanionTo).source(),
-      );
-
-      // TODO(wittjosiah): This should be the default sort order.
-      let nextChat = objects.toSorted((a, b) => a.id.localeCompare(b.id)).at(-1);
-      if (!nextChat) {
-        const { object } = yield* dispatch(createIntent(AssistantAction.CreateChat, { space }));
-        nextChat = object;
-        yield* dispatch(createIntent(SpaceAction.AddObject, { object: nextChat, target: space, hidden: true }));
-        yield* dispatch(
-          createIntent(SpaceAction.AddRelation, {
-            space,
-            schema: Assistant.CompanionTo,
-            source: nextChat,
-            target: data.companionTo,
-          }),
-        );
+    // Sync local chat state with the state atom.
+    // When currentChatState is undefined (plus button clicked), reset local chat.
+    // When data.subject is a Chat object, use it directly.
+    // When currentChatState is a different chat, fetch and set it.
+    useEffect(() => {
+      // If state says no chat selected AND current chat is persisted, reset local state (handles plus button click).
+      // Don't reset in-memory chats - they're expected to have currentChatState === undefined.
+      if (!currentChatState && chat && getSpace(chat)) {
+        setChat(undefined);
+      } else if (data.subject !== 'assistant-chat') {
+        // If data.subject is a Chat object (resolved from graph), use it.
+        setChat(data.subject);
+      } else if (currentChatState && space) {
+        // currentChatState is set but graph couldn't resolve it - fetch manually.
+        const currentChatDxnStr = chat ? Obj.getDXN(chat).toString() : undefined;
+        if (currentChatState !== currentChatDxnStr) {
+          // Parse DXN and fetch the chat object from the database.
+          const parsedDxn = DXN.tryParse(currentChatState);
+          if (parsedDxn) {
+            const chatRef = space.db.makeRef(parsedDxn);
+            const resolvedChat = chatRef?.target;
+            if (Obj.instanceOf(Assistant.Chat, resolvedChat)) {
+              setChat(resolvedChat);
+            }
+          }
+        }
       }
+    }, [currentChatState, data.subject, space, chat]);
 
-      yield* dispatch(createIntent(AssistantAction.SetCurrentChat, { companionTo, chat: nextChat }));
-    }).pipe(Effect.provide(DatabaseService.layer(space.db)), Effect.runPromise);
-  }, [space, chat, companionTo]);
-
-  const metadata = useCapabilities(Capabilities.Metadata);
-  const blueprintKeys = useMemo(
-    () =>
-      Function.pipe(
-        metadata,
-        Array.findFirst(
-          (capability): capability is { id: string; metadata: { blueprints?: string[] } } =>
-            capability.id === Obj.getTypename(companionTo),
-        ),
-        Option.flatMap((c) => Option.fromNullable(c.metadata.blueprints)),
-        Option.getOrElse(() => [] as string[]),
-      ),
-    [metadata, companionTo],
-  );
-  const existingBlueprints = useQuery(space, Filter.type(Blueprint.Blueprint));
-  const pluginBlueprints = useMemo(
-    () => existingBlueprints.filter((blueprint) => blueprintKeys.includes(blueprint.key)),
-    [existingBlueprints, blueprintKeys],
-  );
-
-  // Initialize related blueprints that are not already in the space.
-  useAsyncEffect(async () => {
-    if (!space) {
-      return;
-    }
+    const chatQueue = space && chat ? space.queues.get(chat.queue.dxn) : undefined;
+    const binder = useContextBinder(chatQueue);
 
-    // NOTE: This must be run instead of using the useQuery result to avoid duplicates.
-    const { objects: existingBlueprints } = await space.db.query(Filter.type(Blueprint.Blueprint)).run();
-    for (const key of blueprintKeys) {
-      const existingBlueprint = existingBlueprints.find((blueprint) => blueprint.key === key);
-      if (existingBlueprint) {
-        continue;
+    // Initialize companion chat if it doesn't exist, but don't add it to the space immediately.
+    useAsyncEffect(async () => {
+      if (!space || chat) {
+        return;
       }
 
-      const blueprint = blueprintRegistry.getByKey(key);
-      if (!blueprint) {
-        continue;
+      // Only query for existing chats on initial mount, not when plus button is clicked.
+      if (!initialSetupDoneRef.current) {
+        // Query for existing companion chats linked to this object.
+        const existingChats = await space.db
+          .query(Query.select(Filter.id(companionTo.id)).targetOf(Chat.CompanionTo).source())
+          .run();
+
+        initialSetupDoneRef.current = true;
+
+        // Use existing chat if found on initial mount.
+        if (existingChats.length > 0) {
+          const existingChat = existingChats.at(-1) as Assistant.Chat;
+          setChat(existingChat);
+          await invokePromise(AssistantOperation.SetCurrentChat, { companionTo, chat: existingChat });
+          return;
+        }
       }
 
-      space.db.add(Obj.clone(blueprint));
-    }
-  }, [space, blueprintRegistry, blueprintKeys]);
+      // Create chat in-memory only - it will be added to space on first message.
+      const { data: createResult } = await invokePromise(AssistantOperation.CreateChat, {
+        db: space.db,
+        addToSpace: false,
+      });
+      setChat(createResult?.object);
+    }, [chat, space, companionTo, invokePromise]);
+
+    // Add chat to space and create relation when user submits the first message.
+    const handleEvent = useCallback(
+      async (event: ChatEvent) => {
+        if (!chat || !space) {
+          return;
+        }
+
+        // If chat is not in space yet, persist it on first submit.
+        if (event.type === 'submit' && !getSpace(chat)) {
+          await invokePromise(SpaceOperation.AddObject, {
+            object: chat,
+            target: space.db,
+            hidden: true,
+          });
+          await invokePromise(SpaceOperation.AddRelation, {
+            db: space.db,
+            schema: Chat.CompanionTo,
+            source: chat,
+            target: companionTo,
+          });
+          await invokePromise(AssistantOperation.SetCurrentChat, {
+            companionTo,
+            chat,
+          });
+        }
+      },
+      [chat, space, companionTo, invokePromise],
+    );
+
+    const metadata = useCapabilities(AppCapabilities.Metadata);
+    const blueprintKeys = useMemo(
+      () =>
+        Function.pipe(
+          metadata,
+          Array.findFirst(
+            (
+              capability,
+            ): capability is {
+              id: string;
+              metadata: { blueprints?: string[] };
+            } => capability.id === Obj.getTypename(companionTo),
+          ),
+          Option.flatMap((c) => Option.fromNullable(c.metadata.blueprints)),
+          Option.getOrElse(() => [] as string[]),
+        ),
+      [metadata, companionTo],
+    );
+    const existingBlueprints = useQuery(space?.db, Filter.type(Blueprint.Blueprint));
+    const pluginBlueprints = useMemo(
+      () => existingBlueprints.filter((blueprint) => blueprintKeys.includes(blueprint.key)),
+      [existingBlueprints, blueprintKeys],
+    );
+
+    // Initialize related blueprints that are not already in the space.
+    useAsyncEffect(async () => {
+      if (!space) {
+        return;
+      }
 
-  useAsyncEffect(async () => {
-    if (!binder) {
-      return;
-    }
+      // NOTE: This must be run instead of using the useQuery result to avoid duplicates.
+      const existingBlueprints = await space.db.query(Filter.type(Blueprint.Blueprint)).run();
+      for (const key of blueprintKeys) {
+        const existingBlueprint = existingBlueprints.find((blueprint) => blueprint.key === key);
+        if (existingBlueprint) {
+          continue;
+        }
 
-    if (pluginBlueprints.length > 0) {
-      await binder.bind({ blueprints: pluginBlueprints.map((blueprint) => Ref.make(blueprint)) });
-    }
+        const blueprint = blueprintRegistry.getByKey(key);
+        if (!blueprint) {
+          continue;
+        }
 
-    if (Obj.instanceOf(Blueprint.Blueprint, companionTo)) {
-      await binder.bind({ blueprints: [Ref.make(companionTo)] });
-    } else {
-      await binder.bind({ objects: [Ref.make(companionTo)] });
-    }
-  }, [binder, companionTo, pluginBlueprints]);
+        space.db.add(Obj.clone(blueprint, { deep: true }));
+      }
+    }, [space, blueprintRegistry, blueprintKeys]);
 
-  if (!chat) {
-    return null;
-  }
+    useAsyncEffect(async () => {
+      if (!binder?.isOpen) {
+        return;
+      }
 
-  return <ChatContainer role={role} chat={chat} companionTo={companionTo} />;
-};
+      if (pluginBlueprints.length > 0) {
+        await binder.bind({
+          blueprints: pluginBlueprints.map((blueprint) => Ref.make(blueprint)),
+        });
+      }
+
+      if (Obj.instanceOf(Blueprint.Blueprint, companionTo)) {
+        await binder.bind({ blueprints: [Ref.make(companionTo)] });
+      } else {
+        await binder.bind({ objects: [Ref.make(companionTo)] });
+      }
+    }, [binder, companionTo, blueprintKeys]);
+
+    return (
+      <ChatContainer
+        role={role}
+        space={space}
+        subject={chat}
+        companionTo={companionTo}
+        onEvent={handleEvent}
+        ref={forwardedRef}
+      />
+    );
+  },
+);
 
 export default ChatCompanion;

package/src/components/ChatContainer.tsx

@@ -2,50 +2,67 @@
 // Copyright 2025 DXOS.org
 //
 
-import React from 'react';
+import React, { forwardRef } from 'react';
 
-import { Capabilities, useCapability } from '@dxos/app-framework';
-import { getSpace } from '@dxos/client/echo';
+import { useAtomCapability } from '@dxos/app-framework/ui';
+import { type SurfaceComponentProps } from '@dxos/app-toolkit/ui';
+import { type Space, getSpace } from '@dxos/client/echo';
 import { type Obj } from '@dxos/echo';
-import { StackItem } from '@dxos/react-ui-stack';
+import { Layout } from '@dxos/react-ui';
 
 import { useBlueprintRegistry, useChatProcessor, useChatServices, useOnline, usePresets } from '../hooks';
-import { meta } from '../meta';
-import { type Assistant } from '../types';
-
-import { Chat } from './Chat';
-import { Toolbar } from './Toolbar';
-
-export type ChatContainerProps = {
-  chat: Assistant.Chat;
-  companionTo?: Obj.Any;
-  role?: string;
-};
-
-export const ChatContainer = ({ chat, companionTo }: ChatContainerProps) => {
-  const space = getSpace(chat);
-  const settings = useCapability(Capabilities.SettingsStore).getStore<Assistant.Settings>(meta.id)?.value;
-  const services = useChatServices({ space, chat });
+import { AssistantCapabilities, type ChatType } from '../types';
+
+import { Chat as ChatComponent, type ChatRootProps } from './Chat';
+
+export type ChatContainerProps = SurfaceComponentProps<
+  ChatType.Chat | undefined,
+  {
+    space?: Space;
+    companionTo?: Obj.Unknown;
+  } & Pick<ChatRootProps, 'onEvent'>
+>;
+
+export const ChatContainer = forwardRef<HTMLDivElement, ChatContainerProps>((props, forwardedRef) => {
+  const { role, subject: chat, space: spaceProp, companionTo, onEvent } = props;
+  const space = spaceProp ?? getSpace(chat);
+  const settings = useAtomCapability(AssistantCapabilities.Settings);
+  const services = useChatServices({ id: space?.id, chat });
   const [online, setOnline] = useOnline();
   const { preset, ...chatProps } = usePresets(online);
   const blueprintRegistry = useBlueprintRegistry();
-  const processor = useChatProcessor({ chat, preset, services, blueprintRegistry, settings });
+  const processor = useChatProcessor({
+    space,
+    chat,
+    preset,
+    services,
+    blueprintRegistry,
+    settings,
+  });
 
   if (!processor) {
     return null;
   }
 
   return (
-    <StackItem.Content toolbar={!!companionTo}>
-      {!!companionTo && <Toolbar chat={chat} companionTo={companionTo} />}
-      <Chat.Root classNames='container-max-width' chat={chat} processor={processor}>
-        <Chat.Thread />
-        <div className='p-2'>
-          <Chat.Prompt {...chatProps} outline preset={preset?.id} online={online} onOnlineChange={setOnline} />
-        </div>
-      </Chat.Root>
-    </StackItem.Content>
+    <Layout.Main toolbar role={role} ref={forwardedRef}>
+      <ChatComponent.Root db={space?.db} chat={chat} processor={processor} onEvent={onEvent}>
+        <ChatComponent.Toolbar companionTo={companionTo} />
+        <ChatComponent.Viewport classNames='container-max-width'>
+          <ChatComponent.Thread />
+          <div role='none' className='p-4'>
+            <ChatComponent.Prompt
+              {...chatProps}
+              outline
+              preset={preset?.id}
+              online={online}
+              onOnlineChange={setOnline}
+            />
+          </div>
+        </ChatComponent.Viewport>
+      </ChatComponent.Root>
+    </Layout.Main>
   );
-};
+});
 
 export default ChatContainer;

package/src/components/ChatDialog.tsx

@@ -4,31 +4,38 @@
 
 import React, { useCallback, useState } from 'react';
 
-import { Capabilities, useCapability } from '@dxos/app-framework';
-import { getSpace } from '@dxos/client/echo';
+import { useAtomCapability } from '@dxos/app-framework/ui';
+import { type Chat as ChatTypes } from '@dxos/assistant-toolkit';
+import { Obj } from '@dxos/echo';
 import { useTranslation } from '@dxos/react-ui';
 import { ChatDialog as NaturalChatDialog } from '@dxos/react-ui-chat';
 
 import { useBlueprintRegistry, useChatProcessor, useChatServices, useOnline, usePresets } from '../hooks';
 import { meta } from '../meta';
-import { type Assistant } from '../types';
+import { AssistantCapabilities } from '../types';
 
 import { Chat, type ChatRootProps } from './Chat';
 
 export type ChatDialogProps = {
-  chat?: Assistant.Chat;
+  chat?: ChatTypes.Chat;
 };
 
 export const ChatDialog = ({ chat }: ChatDialogProps) => {
   const { t } = useTranslation(meta.id);
 
-  const space = getSpace(chat);
-  const settings = useCapability(Capabilities.SettingsStore).getStore<Assistant.Settings>(meta.id)?.value;
-  const services = useChatServices({ space, chat });
+  const db = chat && Obj.getDatabase(chat);
+  const settings = useAtomCapability(AssistantCapabilities.Settings);
+  const services = useChatServices({ id: db?.spaceId, chat });
   const [online, setOnline] = useOnline();
   const { preset, ...chatProps } = usePresets(online);
   const blueprintRegistry = useBlueprintRegistry();
-  const processor = useChatProcessor({ chat, preset, services, blueprintRegistry, settings });
+  const processor = useChatProcessor({
+    chat,
+    preset,
+    services,
+    blueprintRegistry,
+    settings,
+  });
 
   // TODO(burdon): Refocus when open.
   const [open, setOpen] = useState(true);

package/src/components/ChatProgress/ChatProgress.tsx

@@ -4,26 +4,12 @@
 
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
+import { type Chat } from '@dxos/assistant-toolkit';
 import { addEventListener } from '@dxos/async';
-import { Flex, ProgressBar, type ProgressBarProps, TextCrawl, useExecutionGraph } from '@dxos/react-ui-components';
-
-import { type Assistant } from '../../types';
-
-// TODO(burdon): Reset after each session.
-// TODO(burdon): Show after delay.
-// TODO(burdon): Test with errors.
-// TODO(burdon): Filter out empty trace messages.
-
-// {
-//   "id": "01K3TE3C2VAXETP500QKDHK3C4_block_3",
-//   "branch": "main",
-//   "parents": [],
-//   "icon": "ph--robot--regular",
-//   "message": "\n"
-// }
+import { ProgressBar, type ProgressBarProps, TextCrawl, useExecutionGraph } from '@dxos/react-ui-components';
 
 export type ChatProgressProps = {
-  chat: Assistant.Chat;
+  chat: Chat.Chat;
 };
 
 export const ChatProgress = ({ chat }: ChatProgressProps) => {
@@ -59,9 +45,9 @@ export const ChatProgress = ({ chat }: ChatProgressProps) => {
   }, [nodes.length]);
 
   return (
-    <Flex column tabIndex={0} ref={ref} classNames='outline-none'>
+    <div role='none' tabIndex={0} ref={ref} className='flex flex-col outline-none'>
       <ProgressBar nodes={nodes} index={index} onSelect={handleSelect} />
-      <TextCrawl lines={lines} index={index} autoAdvance classNames='pis-4 text-sm text-subdued' />
-    </Flex>
+      <TextCrawl classNames='pis-4 text-sm text-description' lines={lines} autoAdvance />
+    </div>
   );
 };

package/src/components/ChatPrompt/ChatActions.tsx

@@ -5,7 +5,7 @@
 import React, { type PropsWithChildren } from 'react';
 
 import { IconButton, type ThemedClassName, useTranslation } from '@dxos/react-ui';
-import { mx } from '@dxos/react-ui-theme';
+import { mx } from '@dxos/ui-theme';
 
 import { meta } from '../../meta';
 import { type ChatEvent } from '../Chat';
@@ -29,10 +29,9 @@ export const ChatActions = ({ classNames, children, microphone, recording, proce
       <IconButton
         disabled={!processing}
         variant='ghost'
-        size={5}
         icon='ph--x--regular'
         iconOnly
-        label={t('button cancel processing')}
+        label={t('cancel processing button')}
         onClick={() => onEvent?.({ type: 'cancel' })}
       />
 
@@ -41,17 +40,24 @@ export const ChatActions = ({ classNames, children, microphone, recording, proce
           disabled={!processing}
           classNames={mx(recording && 'bg-primary-500')}
           variant='ghost'
-          size={5}
           icon='ph--microphone--regular'
           iconOnly
           noTooltip
-          label={t('button microphone')}
+          label={t('microphone button')}
           onMouseDown={() => onEvent?.({ type: 'record-start' })}
           onMouseUp={() => onEvent?.({ type: 'record-stop' })}
           onTouchStart={() => onEvent?.({ type: 'record-start' })}
           onTouchEnd={() => onEvent?.({ type: 'record-stop' })}
         />
       )}
+
+      <IconButton
+        variant='ghost'
+        icon='ph--wrench--regular'
+        iconOnly
+        label={t('debug button')}
+        onClick={() => onEvent?.({ type: 'toggle-debug' })}
+      />
     </div>
   );
 };