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

package/src/components/Chat/Chat.tsx

@@ -10,23 +10,23 @@ import * as Array from 'effect/Array';
 import * as Option from 'effect/Option';
 import React, { type PropsWithChildren, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
+import { type Chat as ChatModule } from '@dxos/assistant-toolkit';
 import { Event } from '@dxos/async';
-import { type Database, Obj } from '@dxos/echo';
+import { type Database, Filter, Obj } from '@dxos/echo';
 import { useVoiceInput } from '@dxos/plugin-transcription';
-import { useQueue } from '@dxos/react-client/echo';
+import { type Queue, useQuery } from '@dxos/react-client/echo';
 import { useIdentity } from '@dxos/react-client/halo';
 import { Input, type ThemedClassName, useDynamicRef, useTranslation } from '@dxos/react-ui';
 import { ChatEditor, type ChatEditorController, type ChatEditorProps } from '@dxos/react-ui-chat';
 import { type MarkdownStreamController } from '@dxos/react-ui-components';
-import { MenuProvider, ToolbarMenu } from '@dxos/react-ui-menu';
+import { Menu, MenuRootProps } from '@dxos/react-ui-menu';
 import { Message } from '@dxos/types';
-import { mx } from '@dxos/ui-theme';
+import { composable, composableProps, mx } from '@dxos/ui-theme';
 import { isTruthy } from '@dxos/util';
 
-import { useChatToolbarActions } from '../../hooks';
-import { meta } from '../../meta';
+import { useChatToolbarActions } from '#hooks';
+import { meta } from '#meta';
 import { type AiChatProcessor } from '../../processor';
-import { type Assistant } from '../../types';
 import {
   ChatActions,
   type ChatActionsProps,
@@ -49,7 +49,7 @@ type ChatContextValue = {
   debug?: boolean;
   event: Event<ChatEvent>;
   db?: Database.Database;
-  chat?: Assistant.Chat;
+  chat?: ChatModule.Chat;
   messages: Message.Message[];
   processor: AiChatProcessor;
 };
@@ -62,22 +62,22 @@ export const [ChatContextProvider, useChatContext] = createContext<ChatContextVa
 
 type ChatRootProps = PropsWithChildren<
   Pick<ChatContextValue, 'db' | 'chat' | 'processor'> & {
+    queue?: Queue;
     onEvent?: (event: ChatEvent) => void;
   }
 >;
 
-const ChatRoot = ({ children, chat, processor, onEvent, ...props }: ChatRootProps) => {
+const ChatRoot = ({ children, chat, queue, processor, onEvent, ...props }: ChatRootProps) => {
   const [debug, setDebug] = useState(false);
   const pending = useAtomValue(processor.messages);
   const streaming = useAtomValue(processor.streaming);
   const lastPrompt = useRef<string | undefined>(undefined);
 
   // Messages.
-  const queue = useQueue<Message.Message>(chat?.queue.dxn);
+  const storedMessages = useQuery(queue, Filter.type(Message.Message));
   const messages = useMemo(() => {
-    const queueMessages = queue?.objects?.filter(Obj.instanceOf(Message.Message)) ?? [];
-    return Array.dedupeWith([...queueMessages, ...pending], ({ id: a }, { id: b }) => a === b);
-  }, [queue?.objects, pending]);
+    return Array.dedupeWith([...storedMessages, ...pending], ({ id: a }, { id: b }) => a === b);
+  }, [storedMessages, pending]);
 
   // Events.
   const event = useMemo(() => new Event<ChatEvent>(), []);
@@ -85,7 +85,26 @@ const ChatRoot = ({ children, chat, processor, onEvent, ...props }: ChatRootProp
     return event.on((ev) => {
       switch (ev.type) {
         case 'toggle-debug': {
-          setDebug((current) => !current);
+          setDebug((debug) => !debug);
+
+          // Dump state to console.
+          queueMicrotask(async () => {
+            const objects = processor.context.getObjects();
+            const blueprints = processor.context.getBlueprints();
+            const tools = await processor.getTools();
+            const system = await processor.getSystemPrompt();
+            // eslint-disable-next-line no-console
+            console.log('Chat processor state:', { objects, blueprints });
+            // eslint-disable-next-line no-console
+            console.log(`
+              ==== System Prompt ====
+              ${system}
+              ==== Tools ====
+              ${Object.values(tools)
+                .map((tool) => JSON.stringify(tool, null, 2))
+                .join('\n')}
+            `);
+          });
           break;
         }
 
@@ -143,24 +162,30 @@ ChatRoot.displayName = 'Chat.Root';
 // Viewport
 //
 
-type ChatViewportProps = ThemedClassName<PropsWithChildren>;
+const CHAT_VIEWPORT_NAME = 'Chat.Viewport';
 
-const ChatViewport = ({ classNames, children }: ChatViewportProps) => {
+type ChatViewportProps = {};
+
+const ChatViewport = composable<HTMLDivElement, ChatViewportProps>(({ children, ...props }, forwardedRef) => {
   return (
-    <div role='none' className={mx('flex flex-col bs-full is-full', classNames)}>
+    <div {...composableProps(props, { classNames: 'dx-expander flex flex-col' })} ref={forwardedRef}>
       {children}
     </div>
   );
-};
+});
+
+ChatViewport.displayName = CHAT_VIEWPORT_NAME;
 
 //
 // Thread
 //
 
+const CHAT_THREAD_NAME = 'Chat.Thread';
+
 type ChatThreadProps = Omit<NaturalChatThreadProps, 'identity' | 'messages' | 'tools'>;
 
 const ChatThread = (props: ChatThreadProps) => {
-  const { debug, event, messages, processor } = useChatContext(ChatThread.displayName);
+  const { debug, event, messages, processor } = useChatContext(CHAT_THREAD_NAME);
   const identity = useIdentity();
   const error = useAtomValue(processor.error).pipe(Option.getOrUndefined);
 
@@ -206,12 +231,14 @@ const ChatThread = (props: ChatThreadProps) => {
   );
 };
 
-ChatThread.displayName = 'Chat.Thread';
+ChatThread.displayName = CHAT_THREAD_NAME;
 
 //
 // Prompt
 //
 
+const CHAT_PROMPT_NAME = 'Chat.Prompt';
+
 type ChatPromptProps = ThemedClassName<
   {
     outline?: boolean;
@@ -238,7 +265,7 @@ const ChatPrompt = ({
   onOnlineChange,
 }: ChatPromptProps) => {
   const { t } = useTranslation(meta.id);
-  const { db, processor, event } = useChatContext(ChatPrompt.displayName);
+  const { db, processor, event } = useChatContext(CHAT_PROMPT_NAME);
 
   const error = useAtomValue(processor.error).pipe(Option.getOrUndefined);
   const streaming = useAtomValue(processor.streaming);
@@ -337,9 +364,9 @@ const ChatPrompt = ({
     <div
       role='group'
       className={mx(
-        'flex flex-col is-full density-fine',
+        'flex flex-col w-full dx-density-fine',
         outline &&
-          'bg-groupSurface border border-subduedSeparator transition transition-border [&:has(.cm-content:focus)]:border-separator rounded',
+          'bg-group-surface border border-subdued-separator transition transition-border [&:has(.cm-content:focus)]:border-separator rounded-sm',
         classNames,
       )}
     >
@@ -349,8 +376,8 @@ const ChatPrompt = ({
           ref={editorRef}
           autoFocus
           lineWrapping
-          classNames='col-span-2 pbs-0.5'
-          placeholder={placeholder ?? t('prompt placeholder')}
+          classNames='col-span-2 pt-0.5'
+          placeholder={placeholder ?? t('prompt.placeholder')}
           extensions={extensions}
           onSubmit={handleSubmit}
         />
@@ -381,8 +408,8 @@ const ChatPrompt = ({
             {/* TODO(burdon): Move offline switch into dialog. */}
             {online !== undefined && (
               <Input.Root>
-                <Input.Label srOnly>{t('online switch label')}</Input.Label>
-                <Input.Switch classNames='mli-2' checked={online} onCheckedChange={onOnlineChange} />
+                <Input.Label srOnly>{t('online-switch.label')}</Input.Label>
+                <Input.Switch classNames='mx-2' checked={online} onCheckedChange={onOnlineChange} />
               </Input.Root>
             )}
           </ChatActions>
@@ -392,29 +419,32 @@ const ChatPrompt = ({
   );
 };
 
-ChatPrompt.displayName = 'Chat.Prompt';
+ChatPrompt.displayName = CHAT_PROMPT_NAME;
 
 //
 // Toolbar
 //
 
-type ChatToolbarProps = ThemedClassName<{ companionTo?: Obj.Unknown }>;
-
-const ChatToolbar = ({ classNames, companionTo }: ChatToolbarProps) => {
-  const { chat } = useChatContext(ChatToolbar.displayName);
-  const menu = useChatToolbarActions({ chat, companionTo });
+const CHAT_TOOLBAR_NAME = 'Chat.Toolbar';
 
-  return (
-    <MenuProvider
-      {...menu}
-      attendableId={companionTo ? Obj.getDXN(companionTo).toString() : chat ? Obj.getDXN(chat).toString() : ''}
-    >
-      <ToolbarMenu classNames={classNames} textBlockWidth />
-    </MenuProvider>
-  );
+type ChatToolbarProps = Pick<MenuRootProps, 'attendableId'> & {
+  companionTo?: Obj.Unknown;
 };
 
-ChatToolbar.displayName = 'Chat.Toolbar';
+const ChatToolbar = composable<HTMLDivElement, ChatToolbarProps>(
+  ({ attendableId, companionTo, ...props }, forwardedRef) => {
+    const { chat } = useChatContext(CHAT_TOOLBAR_NAME);
+    const menuActions = useChatToolbarActions({ chat, companionTo });
+
+    return (
+      <Menu.Root {...menuActions} attendableId={attendableId}>
+        <Menu.Toolbar {...composableProps(props)} ref={forwardedRef} />
+      </Menu.Root>
+    );
+  },
+);
+
+ChatToolbar.displayName = CHAT_TOOLBAR_NAME;
 
 //
 // Chat

package/src/components/ChatProgress/ChatProgress.tsx

@@ -5,16 +5,18 @@
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
 import { addEventListener } from '@dxos/async';
+import { type Queue } from '@dxos/echo-db';
 import { ProgressBar, type ProgressBarProps, TextCrawl, useExecutionGraph } from '@dxos/react-ui-components';
 
-import { type Assistant } from '../../types';
-
 export type ChatProgressProps = {
-  chat: Assistant.Chat;
+  traceQueue?: Queue;
 };
 
-export const ChatProgress = ({ chat }: ChatProgressProps) => {
-  const { commits } = useExecutionGraph(chat.traceQueue?.target, true);
+/**
+ * Displays execution progress from the space's invocation trace queue.
+ */
+export const ChatProgress = ({ traceQueue }: ChatProgressProps) => {
+  const { commits } = useExecutionGraph(traceQueue, true);
   const nodes = useMemo(() => commits.map((commit) => ({ id: commit.id, message: commit.message })), [commits]);
   const lines = useMemo(() => commits.map((commit) => commit.message), [commits]);
 
@@ -46,9 +48,9 @@ export const ChatProgress = ({ chat }: ChatProgressProps) => {
   }, [nodes.length]);
 
   return (
-    <div role='none' tabIndex={0} ref={ref} className='flex flex-col outline-none'>
+    <div role='none' tabIndex={0} ref={ref} className='flex flex-col outline-hidden'>
       <ProgressBar nodes={nodes} index={index} onSelect={handleSelect} />
-      <TextCrawl classNames='pis-4 text-sm text-description' lines={lines} autoAdvance />
+      <TextCrawl classNames='ps-4 text-sm text-description' lines={lines} autoAdvance />
     </div>
   );
 };

package/src/components/ChatPrompt/ChatActions.tsx

@@ -7,7 +7,7 @@ import React, { type PropsWithChildren } from 'react';
 import { IconButton, type ThemedClassName, useTranslation } from '@dxos/react-ui';
 import { mx } from '@dxos/ui-theme';
 
-import { meta } from '../../meta';
+import { meta } from '#meta';
 import { type ChatEvent } from '../Chat';
 
 export type ChatActionsProps = ThemedClassName<
@@ -23,7 +23,7 @@ export const ChatActions = ({ classNames, children, microphone, recording, proce
   const { t } = useTranslation(meta.id);
 
   return (
-    <div className={mx('flex items-center mie-1', classNames)}>
+    <div className={mx('flex items-center me-1', classNames)}>
       {children}
 
       <IconButton
@@ -31,7 +31,7 @@ export const ChatActions = ({ classNames, children, microphone, recording, proce
         variant='ghost'
         icon='ph--x--regular'
         iconOnly
-        label={t('cancel processing button')}
+        label={t('cancel-processing.button')}
         onClick={() => onEvent?.({ type: 'cancel' })}
       />
 
@@ -43,13 +43,21 @@ export const ChatActions = ({ classNames, children, microphone, recording, proce
           icon='ph--microphone--regular'
           iconOnly
           noTooltip
-          label={t('microphone button')}
+          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>
   );
 };

package/src/components/ChatPrompt/ChatOptions.tsx

@@ -2,26 +2,23 @@
 // Copyright 2025 DXOS.org
 //
 
+import * as Option from 'effect/Option';
 import React, { type JSX, useMemo, useState } from 'react';
 
 import { type AiContextBinder } from '@dxos/assistant';
 import { type Blueprint } from '@dxos/blueprints';
 import { type Database, Filter, Obj, Type } from '@dxos/echo';
+import { Annotation } from '@dxos/echo';
 import { useQuery } from '@dxos/react-client/echo';
 import { IconButton, Popover, Select, useTranslation } from '@dxos/react-ui';
-import { Listbox, SearchList, useSearchListResults } from '@dxos/react-ui-searchlist';
+import { Listbox, SearchList, useSearchListResults } from '@dxos/react-ui-search';
 import { Tabs } from '@dxos/react-ui-tabs';
+import { getStyles, mx } from '@dxos/ui-theme';
 
-import {
-  useActiveBlueprints,
-  useBlueprintHandlers,
-  useBlueprints,
-  useContextObjects,
-  useFilteredTypes,
-} from '../../hooks';
-import { meta } from '../../meta';
+import { useActiveBlueprints, useBlueprintHandlers, useBlueprints, useContextObjects, useFilteredTypes } from '#hooks';
+import { meta } from '#meta';
 
-const panelClassNames = 'is-[calc(100dvw-.5rem)] sm:is-max md:is-72 max-is-[--text-content]';
+const panelClassNames = 'w-[calc(100dvw-.5rem)] sm:w-max md:w-72 max-w-document-width';
 
 export type ChatOptionsProps = {
   db: Database.Database;
@@ -42,7 +39,7 @@ export const ChatOptions = ({ db, context, blueprintRegistry, presets, preset, o
     <div role='none' className='flex p-2'>
       <Popover.Root>
         <Popover.Trigger asChild>
-          <IconButton variant='ghost' icon='ph--plus--regular' iconOnly label={t('context objects button')} />
+          <IconButton variant='ghost' icon='ph--plus--regular' iconOnly label={t('context-objects.button')} />
         </Popover.Trigger>
         <Popover.Portal>
           <Popover.Content side='top' classNames={panelClassNames}>
@@ -60,27 +57,34 @@ export const ChatOptions = ({ db, context, blueprintRegistry, presets, preset, o
             variant='ghost'
             icon='ph--sliders-horizontal--regular'
             iconOnly
-            label={t('context settings button')}
+            label={t('context-settings.button')}
           />
         </Popover.Trigger>
         <Popover.Portal>
           <Popover.Content side='top' classNames={panelClassNames}>
             <Popover.Viewport>
               <Tabs.Root orientation='horizontal' defaultValue='blueprints' defaultActivePart='list' tabIndex={-1}>
-                <Tabs.Viewport classNames='max-bs-[--radix-popover-content-available-height] grid grid-rows-[1fr_min-content] [&_[cmdk-root]]:contents [&_[role="tabpanel"]]:grid [&_[role="tabpanel"]]:grid-rows-[1fr_min-content] [&_[role="listbox"]]:min-bs-0 [&_[role="listbox"]]:overflow-y-auto [&_[role="tabpanel"]]:min-bs-0 [&_[role="tabpanel"]]:pli-cardSpacingChrome [&_[role="tabpanel"][data-state="active"]]:order-first [&_[role="tabpanel"][data-state="inactive"]]:hidden'>
-                  <Tabs.Tabpanel value='blueprints' tabIndex={-1} classNames='dx-focus-ring-inset'>
+                <Tabs.Viewport
+                  classNames={mx(
+                    'max-h-(--radix-popover-content-available-height) grid grid-rows-[1fr_min-content]',
+                    '[&_[cmdk-root]]:contents [&_[role="tabpanel"]]:grid [&_[role="tabpanel"]]:grid-rows-[1fr_min-content]',
+                    '[&_[role="listbox"]]:min-h-0 [&_[role="listbox"]]:overflow-y-auto',
+                    '[&_[role="tabpanel"]]:min-h-0 [&_[role="tabpanel"]]:px-form-chrome [&_[role="tabpanel"][data-state="active"]]:order-first [&_[role="tabpanel"][data-state="inactive"]]:hidden',
+                  )}
+                >
+                  <Tabs.Panel value='blueprints' tabIndex={-1} classNames='dx-focus-ring-inset'>
                     <BlueprintsPanel blueprintRegistry={blueprintRegistry} db={db} context={context} />
-                  </Tabs.Tabpanel>
-                  <Tabs.Tabpanel value='model' tabIndex={-1} classNames='dx-focus-ring-inset !pli-0'>
+                  </Tabs.Panel>
+                  <Tabs.Panel value='model' tabIndex={-1} classNames='dx-focus-ring-inset px-0!'>
                     <ModelsPanel presets={presets} preset={preset} onPresetChange={onPresetChange} />
-                  </Tabs.Tabpanel>
-                  <Tabs.Tablist classNames='sm:overflow-x-hidden justify-center p-[--dx-cardSpacingChrome] border-bs border-subduedSeparator order-last'>
+                  </Tabs.Panel>
+                  <Tabs.Tablist classNames='justify-center p-form-chrome border-y border-subdued-separator order-last'>
                     <Tabs.IconTab
                       value='blueprints'
                       icon='ph--blueprint--regular'
-                      label={t('blueprints in context title')}
+                      label={t('blueprints-in-context.title')}
                     />
-                    <Tabs.IconTab value='model' label={t('chat model title')} icon='ph--cpu--regular' />
+                    <Tabs.IconTab value='model' label={t('chat-model.title')} icon='ph--cpu--regular' />
                   </Tabs.Tablist>
                 </Tabs.Viewport>
               </Tabs.Root>
@@ -111,7 +115,7 @@ const BlueprintsPanel = ({
 
   return (
     <SearchList.Root onSearch={handleSearch}>
-      <SearchList.Content classNames='plb-cardSpacingChrome'>
+      <SearchList.Content classNames='py-form-chrome'>
         <SearchList.Viewport>
           {results.map((blueprint) => {
             const isActive = activeBlueprints.has(blueprint.key);
@@ -128,7 +132,7 @@ const BlueprintsPanel = ({
           })}
         </SearchList.Viewport>
       </SearchList.Content>
-      <SearchList.Input placeholder={t('search placeholder')} classNames='mbe-cardSpacingChrome' autoFocus />
+      <SearchList.Input placeholder={t('search.placeholder')} classNames='mb-form-chrome' autoFocus />
     </SearchList.Root>
   );
 };
@@ -164,7 +168,7 @@ const ObjectsPanel = ({ db, context }: Pick<ChatOptionsProps, 'db' | 'context'>)
       const typename = Type.getTypename(type);
       return {
         typename,
-        label: t('typename label', { ns: typename, defaultValue: typename }),
+        label: t('typename.label', { ns: typename, defaultValue: typename }),
       };
     });
 
@@ -190,16 +194,23 @@ const ObjectsPanel = ({ db, context }: Pick<ChatOptionsProps, 'db' | 'context'>)
 
   return (
     <SearchList.Root onSearch={handleSearch}>
-      <SearchList.Content classNames='p-cardSpacingChrome [&:has([cmdk-list-sizer]:empty)]:plb-0'>
+      <SearchList.Content classNames='p-form-chrome [&:has([cmdk-list-sizer]:empty)]:py-0'>
         <SearchList.Viewport>
           {results.length ? (
             results.map((object) => {
               const isActive = contextObjects.findIndex((obj) => obj.id === object.id) !== -1;
+              const { icon, hue } = Option.fromNullable(Obj.getSchema(object)).pipe(
+                Option.flatMap(Annotation.IconAnnotation.get),
+                Option.getOrElse(() => ({ icon: DEFAULT_OBJECT_ICON, hue: undefined as string | undefined })),
+              );
+              const styles = hue ? getStyles(hue) : undefined;
               return (
                 <SearchList.Item
                   classNames='flex items-center overflow-hidden'
                   key={object.id}
                   value={object.id}
+                  icon={icon}
+                  iconClassNames={styles?.surfaceText}
                   label={Obj.getLabel(object) ?? Obj.getTypename(object) ?? object.id}
                   checked={isActive}
                   onSelect={() => onUpdateObject?.(Obj.getDXN(object), !isActive)}
@@ -207,19 +218,19 @@ const ObjectsPanel = ({ db, context }: Pick<ChatOptionsProps, 'db' | 'context'>)
               );
             })
           ) : (
-            <SearchList.Item value='__empty__' label={t('no results')} />
+            <SearchList.Item value='__empty__' label={t('no-results.message')} />
           )}
         </SearchList.Viewport>
       </SearchList.Content>
 
-      <div role='none' className='grid grid-cols-[min-content_1fr] gap-2 pli-cardSpacingChrome mbe-cardSpacingChrome'>
+      <div role='none' className='grid grid-cols-[min-content_1fr] gap-2 px-form-chrome mb-form-chrome'>
         <Select.Root value={typename === ANY ? undefined : typename} onValueChange={setTypename}>
-          <Select.TriggerButton density='fine' placeholder={t('type filter placeholder')} />
+          <Select.TriggerButton placeholder={t('type-filter.placeholder')} />
           <Select.Portal>
             <Select.Content>
               <Select.ScrollUpButton />
               <Select.Viewport>
-                <Select.Option value={ANY}>{t('any type filter label')}</Select.Option>
+                <Select.Option value={ANY}>{t('any-type-filter.label')}</Select.Option>
                 {typenames.map(({ typename, label }) => (
                   <Select.Option key={typename} value={typename}>
                     {label}
@@ -231,8 +242,11 @@ const ObjectsPanel = ({ db, context }: Pick<ChatOptionsProps, 'db' | 'context'>)
             </Select.Content>
           </Select.Portal>
         </Select.Root>
-        <SearchList.Input placeholder={t('search placeholder')} classNames='mbe-0' autoFocus />
+        <SearchList.Input placeholder={t('search.placeholder')} classNames='mb-0' autoFocus />
       </div>
     </SearchList.Root>
   );
 };
+
+// TODO(dmaretskyi): Extract those somewhere else.
+const DEFAULT_OBJECT_ICON = 'ph--cube--regular';

package/src/components/ChatPrompt/ChatReferences.tsx

@@ -2,15 +2,16 @@
 // Copyright 2025 DXOS.org
 //
 
+import * as Option from 'effect/Option';
 import React from 'react';
 
 import { type AiContextBinder } from '@dxos/assistant';
-import { type Database, Obj } from '@dxos/echo';
-import { IconButton, type Label, type ThemedClassName, toLocalizedString, useTranslation } from '@dxos/react-ui';
-import { mx } from '@dxos/ui-theme';
+import { Annotation, type Database, Obj } from '@dxos/echo';
+import { Icon, IconButton, type Label, type ThemedClassName, toLocalizedString, useTranslation } from '@dxos/react-ui';
+import { getStyles, mx } from '@dxos/ui-theme';
 
-import { useContextObjects } from '../../hooks';
-import { meta } from '../../meta';
+import { useContextObjects } from '#hooks';
+import { meta } from '#meta';
 
 export type ChatReferencesProps = ThemedClassName<{
   context: AiContextBinder;
@@ -26,15 +27,21 @@ export const ChatReferences = ({ classNames, context, db }: ChatReferencesProps)
       {objects.map((obj) => {
         const dxn = Obj.getDXN(obj);
         const typename = Obj.getTypename(obj);
-        const label: Label = Obj.getLabel(obj) ?? (typename ? ['object name placeholder', { ns: typename }] : obj.id);
+        const label: Label = Obj.getLabel(obj) ?? (typename ? ['object-name.placeholder', { ns: typename }] : obj.id);
+        const { icon, hue } = Option.fromNullable(Obj.getSchema(obj)).pipe(
+          Option.flatMap(Annotation.IconAnnotation.get),
+          Option.getOrElse(() => ({ icon: DEFAULT_OBJECT_ICON, hue: undefined as string | undefined })),
+        );
+        const styles = hue ? getStyles(hue) : undefined;
         return (
-          <li key={dxn.toString()} className='dx-tag plb-0 pis-2 flex items-center' data-hue='neutral'>
+          <li key={dxn.toString()} className='dx-tag py-0 ps-2 flex items-center gap-1' data-hue='neutral'>
+            <Icon icon={icon} size={4} classNames={styles?.surfaceText} />
             {toLocalizedString(label, t)}
             <IconButton
               icon='ph--x--bold'
               iconOnly
               variant='ghost'
-              label={t('remove object in context label')}
+              label={t('remove-object-in-context.label')}
               classNames='p-0 hover:bg-transparent'
               size={3}
               onClick={() => onUpdateObject?.(dxn, false)}
@@ -45,3 +52,6 @@ export const ChatReferences = ({ classNames, context, db }: ChatReferencesProps)
     </ul>
   );
 };
+
+// TODO(dmaretskyi): Extract those somewhere else.
+const DEFAULT_OBJECT_ICON = 'ph--cube--regular';

package/src/components/ChatThread/ChatThread.stories.tsx

@@ -15,14 +15,14 @@ import { faker } from '@dxos/random';
 import { useQueue, useSpace } from '@dxos/react-client/echo';
 import { withClientProvider } from '@dxos/react-client/testing';
 import { Popover } from '@dxos/react-ui';
-import { withLayout, withTheme } from '@dxos/react-ui/testing';
+import { Card } from '@dxos/react-ui';
+import { Loading, withLayout, withTheme } from '@dxos/react-ui/testing';
 import { MarkdownStream } from '@dxos/react-ui-components';
 import { EditorPreviewProvider, useEditorPreview } from '@dxos/react-ui-editor';
-import { Card } from '@dxos/react-ui-mosaic';
-import { render } from '@dxos/storybook-utils';
+
 import { type Message, Organization, Person } from '@dxos/types';
 
-import { createMessageGenerator } from '../../testing';
+import { createMessageGenerator } from '#testing';
 import { translations } from '../../translations';
 
 import { ChatThread, type ChatThreadProps } from './ChatThread';
@@ -33,9 +33,9 @@ faker.seed(1);
 
 type MessageGenerator = Effect.Effect<void, never, Database.Service | ContextQueueService>;
 
-type StoryProps = { generator?: MessageGenerator[]; delay?: number; wait?: boolean } & ChatThreadProps;
+type DefaultStoryProps = { generator?: MessageGenerator[]; delay?: number; wait?: boolean } & ChatThreadProps;
 
-const DefaultStory = ({ generator = [], delay = 0, wait, ...props }: StoryProps) => {
+const DefaultStory = ({ generator = [], delay = 0, wait, ...props }: DefaultStoryProps) => {
   const space = useSpace();
   const queueDxn = useMemo(() => space?.queues.create().dxn, [space]);
   const queue = useQueue<Message.Message>(queueDxn);
@@ -56,7 +56,7 @@ const DefaultStory = ({ generator = [], delay = 0, wait, ...props }: StoryProps)
           }
         }
         setDone(true);
-      }).pipe(Effect.provide(Layer.mergeAll(Database.Service.layer(space.db), ContextQueueService.layer(queue)))),
+      }).pipe(Effect.provide(Layer.mergeAll(Database.layer(space.db), ContextQueueService.layer(queue)))),
     );
 
     return () => {
@@ -65,11 +65,11 @@ const DefaultStory = ({ generator = [], delay = 0, wait, ...props }: StoryProps)
   }, [space, queue, generator]);
 
   if (wait && !done) {
-    return null;
+    return <Loading data={{ wait, done }} />;
   }
 
   return (
-    <EditorPreviewProvider onLookup={async ({ label, ref }) => ({ label, text: ref })}>
+    <EditorPreviewProvider onLookup={async ({ dxn, label }) => ({ label, text: dxn })}>
       <ChatThread {...props} messages={queue?.objects} />
       <PreviewCard />
     </EditorPreviewProvider>
@@ -95,11 +95,11 @@ const PreviewCard = () => {
 };
 
 const meta = {
-  title: 'plugins/plugin-assistant/ChatThread',
+  title: 'plugins/plugin-assistant/components/ChatThread',
   component: ChatThread,
-  render: render(DefaultStory),
+  render: DefaultStory,
   decorators: [
-    withTheme,
+    withTheme(),
     withLayout({ layout: 'column' }),
     withClientProvider({
       createIdentity: true,
@@ -111,11 +111,11 @@ const meta = {
     layout: 'fullscreen',
     translations,
   },
-} satisfies Meta<StoryProps>;
+} satisfies Meta<DefaultStoryProps>;
 
 export default meta;
 
-type Story = StoryObj<StoryProps>;
+type Story = StoryObj<DefaultStoryProps>;
 
 export const Default: Story = {
   args: {
@@ -135,7 +135,7 @@ export const Delayed: Story = {
 
 export const Raw: Story = {
   render: () => (
-    <div className='contents' style={{ '--user-fill': 'var(--dx-amberFill)' } as CSSProperties}>
+    <div className='contents' style={{ '--user-fill': 'var(--color-amber-fill)' } as CSSProperties}>
       <MarkdownStream content={TEXT} />
     </div>
   ),
@@ -143,7 +143,7 @@ export const Raw: Story = {
 
 export const Static: Story = {
   render: () => (
-    <div className='contents' style={{ '--user-fill': 'var(--dx-amberFill)' } as CSSProperties}>
+    <div className='contents' style={{ '--user-fill': 'var(--color-amber-fill)' } as CSSProperties}>
       <MarkdownStream content={TEXT} registry={componentRegistry} />
     </div>
   ),