@dxos/plugin-assistant 0.8.4-main.fcfe5033a5 → 0.8.4-staging.60fe92afc8

package/src/containers/TracePanel/TracePanel.stories.tsx

@@ -2,206 +2,268 @@
 // Copyright 2026 DXOS.org
 //
 
-import { Registry } from '@effect-atom/atom';
-import * as KeyValueStore from '@effect/platform/KeyValueStore';
 import { type Meta, type StoryObj } from '@storybook/react-vite';
 import * as Effect from 'effect/Effect';
-import * as Layer from 'effect/Layer';
-import * as ManagedRuntime from 'effect/ManagedRuntime';
-import * as Schema from 'effect/Schema';
-import React, { useCallback, useEffect, useRef } from 'react';
+import React, { PropsWithChildren, type ReactNode, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 
 import { withPluginManager } from '@dxos/app-framework/testing';
-import { AgentRequestBegin, AgentRequestEnd, CompleteBlock } from '@dxos/assistant';
-import { Database, Feed } from '@dxos/echo';
-import { createFeedServiceLayer } from '@dxos/echo-db';
-import { Process, ServiceResolver, Trace } from '@dxos/functions';
-import { FeedTraceSink, ProcessManager } from '@dxos/functions-runtime';
-import { ObjectId } from '@dxos/keys';
-import { OperationHandlerSet } from '@dxos/operation';
-import { ClientPlugin } from '@dxos/plugin-client';
+import { useProcessManagerRuntime } from '@dxos/app-framework/ui';
+import { addEventListener } from '@dxos/async';
+import { Process, Trace } from '@dxos/compute';
+import { ProcessManager } from '@dxos/compute-runtime';
+import { Feed, Filter, Query } from '@dxos/echo';
+import { FeedTraceSink } from '@dxos/functions-runtime';
+import { log } from '@dxos/log';
+import { AutomationPlugin } from '@dxos/plugin-automation/testing';
+import { ClientPlugin } from '@dxos/plugin-client/testing';
 import { initializeIdentity } from '@dxos/plugin-client/testing';
 import { corePlugins } from '@dxos/plugin-testing';
-import { type Space, useSpaces } from '@dxos/react-client/echo';
-import { IconButton, Panel, Toolbar } from '@dxos/react-ui';
-import { Loading, withLayout, withTheme } from '@dxos/react-ui/testing';
-
-import { translations } from '../../translations';
+import { useQuery, useSpaces } from '@dxos/react-client/echo';
+import { IconButton, Panel, ScrollContainer, Toolbar } from '@dxos/react-ui';
+import { Timeline } from '@dxos/react-ui-components';
+import { withLayout, withTheme } from '@dxos/react-ui/testing';
+
+import { buildExecutionGraph } from '#execution-graph';
+import { AssistantPlugin } from '#plugin';
+import { translations } from '#translations';
+
+import subAgentFixture from '../../execution-graph/testing/sub-agent-delegation.json';
+// TODO(dmaretskyi): testing.ts module shadows the ./testing dir.
+import { initClientFromSpaceSnapshot } from '../../testing/snapshot';
+import { PLAYBACK_INTERVAL_MS, STEP_STORAGE_KEY, SimulatedAgent, useLocalStorageNumber } from './testing';
 import { TracePanel } from './TracePanel';
 
-//
-// Helpers.
-//
+type BaseStoryProps = PropsWithChildren<{
+  toolbar: ReactNode;
+}>;
 
-const delay = (ms: number) => Effect.promise(() => new Promise<void>((resolve) => setTimeout(resolve, ms)));
+const BaseStory = ({ children, toolbar }: BaseStoryProps) => (
+  <Panel.Root>
+    <Panel.Toolbar asChild>{toolbar}</Panel.Toolbar>
+    <Panel.Content asChild>{children}</Panel.Content>
+  </Panel.Root>
+);
 
-//
-// Simulated agent scenarios.
-//
+const DefaultStory = () => {
+  const [space] = useSpaces();
+  const runtime = useProcessManagerRuntime();
 
-const agentScenarios: { prompt: string; tools: string[] }[] = [
-  {
-    prompt: 'Create an organization called "Cyberdyne Systems"',
-    tools: ['list-schemas', 'create-object'],
-  },
-  {
-    prompt: 'Search for all organizations and persons',
-    tools: ['list-schemas', 'query', 'query'],
-  },
-  {
-    prompt: 'Create a person named "John Connor"',
-    tools: ['create-object'],
-  },
-];
+  const handleStart = useCallback(() => {
+    if (!runtime) {
+      return;
+    }
 
-/**
- * Write a trace event and yield to the event loop so FeedTraceSink flushes it to ECHO before continuing.
- */
-const writeAndFlush = <T,>(eventType: Trace.EventType<T>, payload: T) =>
-  Effect.gen(function* () {
-    yield* Trace.write(eventType, payload);
-    yield* Effect.promise(() => new Promise((resolve) => setTimeout(resolve, 100)));
-  });
-
-const SimulatedAgent = Process.make(
-  {
-    key: 'org.dxos.testing.process.agent',
-    input: Schema.Number,
-    output: Schema.Void,
-    services: [Trace.TraceService],
-  },
-  (ctx) =>
-    Effect.succeed({
-      onInput: (scenarioIndex: number) =>
-        Effect.gen(function* () {
-          const scenario = agentScenarios[scenarioIndex % agentScenarios.length];
-          const messageId = ObjectId.random();
-
-          // Agent begins processing.
-          yield* writeAndFlush(AgentRequestBegin, {});
-
-          // User message.
-          yield* writeAndFlush(CompleteBlock, {
-            messageId,
-            role: 'user',
-            block: { _tag: 'text', text: scenario.prompt },
-          });
-
-          // Simulate tool calls sequentially with delays.
-          for (const toolName of scenario.tools) {
-            const toolCallId = ObjectId.random();
-
-            // Tool call event.
-            yield* writeAndFlush(CompleteBlock, {
-              messageId,
-              role: 'assistant',
-              block: { _tag: 'toolCall', toolCallId, name: toolName, input: '{}', providerExecuted: false },
-            });
-
-            // Simulate tool execution time.
-            yield* delay(1_000 + Math.random() * 3_000);
-
-            // Tool result event.
-            yield* writeAndFlush(CompleteBlock, {
-              messageId,
-              role: 'assistant',
-              block: { _tag: 'toolResult', toolCallId, name: toolName, providerExecuted: false },
-            });
-          }
-
-          // Agent completes.
-          yield* writeAndFlush(AgentRequestEnd, {});
-          ctx.succeed();
-        }).pipe(Effect.orDie),
-      onAlarm: () => Effect.void,
-      onChildEvent: () => Effect.void,
-    }),
-);
+    void runtime.runPromise(
+      Effect.gen(function* () {
+        const manager = yield* ProcessManager.Service;
+        const handle = yield* manager.spawn(SimulatedAgent);
+        yield* handle.submitInput(Math.floor(Math.random() * 1_000));
+        log.info('submitInput', { handle });
+      }),
+    );
+  }, [runtime]);
+
+  // TODO(burdon): Implement.
+  const handleStop = useCallback(
+    (process: Process.Info) => {
+      log.info('stop', { process });
+    },
+    [runtime],
+  );
 
-//
-// Story runtime — lightweight ProcessManager + FeedTraceSink.
-//
+  return (
+    <BaseStory
+      toolbar={
+        <Toolbar.Root>
+          <IconButton icon='ph--plus--regular' label='Start Agent' onClick={handleStart} />
+        </Toolbar.Root>
+      }
+    >
+      <TracePanel space={space} attendableId={space.id} onProcessTerminate={handleStop} />
+    </BaseStory>
+  );
+};
 
-type StoryRuntime = ManagedRuntime.ManagedRuntime<ProcessManager.ProcessManagerService, never>;
+const SnapshotStory = () => {
+  const [space] = useSpaces();
+  const [feed] = useQuery(space?.db, FeedTraceSink.query);
 
-const createStoryRuntime = (space: Space): StoryRuntime => {
-  const dbLayer = Database.layer(space.db);
-  const feedServiceLayer = createFeedServiceLayer(space.queues);
-  const traceSinkLayer = FeedTraceSink.layerLive.pipe(Layer.provide(dbLayer), Layer.provide(feedServiceLayer));
+  // All messages.
+  const allMessages = useQuery(
+    space?.db,
+    feed ? Query.select(Filter.everything()).from(feed) : Query.select(Filter.nothing()),
+  );
 
-  const layer = ProcessManager.layer().pipe(
-    Layer.provide(Layer.succeed(ServiceResolver.ServiceResolver, ServiceResolver.empty)),
-    Layer.provide(KeyValueStore.layerMemory),
-    Layer.provide(OperationHandlerSet.provide(OperationHandlerSet.empty)),
-    Layer.provideMerge(Registry.layer),
-    Layer.provide(traceSinkLayer),
+  // Sort by first event timestamp so playback order is chronological regardless of query ordering.
+  const sortedMessages = useMemo(
+    () => [...allMessages].sort((a, b) => (a.events[0]?.timestamp ?? 0) - (b.events[0]?.timestamp ?? 0)),
+    [allMessages],
   );
 
-  return ManagedRuntime.make(layer);
+  return <TimelinePlayback messages={sortedMessages} />;
 };
 
-//
-// Story component.
-//
+// Raw Trace.Message[] captured from a live sub-agent delegation via `dxosDumpTrace()` (see
+// TracePanel). External JSON → typed at this boundary; `buildExecutionGraph` only reads
+// `meta`/`events`, so the plain data shape is sufficient.
+const subAgentMessages = subAgentFixture as unknown as Trace.Message[];
 
-const DefaultStory = () => {
-  const [space] = useSpaces();
-  const runtimeRef = useRef<StoryRuntime | null>(null);
-  const invokeCounterRef = useRef(0);
+const FixtureStory = () => {
+  const sortedMessages = useMemo(
+    () => [...subAgentMessages].sort((a, b) => (a.events[0]?.timestamp ?? 0) - (b.events[0]?.timestamp ?? 0)),
+    [],
+  );
+
+  return <TimelinePlayback messages={sortedMessages} />;
+};
+
+/**
+ * Step/playback over a fixed list of trace messages — renders the execution-graph Timeline at each
+ * step so layout (e.g. concurrent sub-agent lanes) can be inspected as events "stream" in.
+ */
+const TimelinePlayback = ({ messages: sortedMessages }: { messages: readonly Trace.Message[] }) => {
+  const total = sortedMessages.length;
+  const [step, setStep, stepHydrated] = useLocalStorageNumber(STEP_STORAGE_KEY, 0);
+  const [playing, setPlaying] = useState(false);
 
   useEffect(() => {
-    if (!space) {
-      return;
+    setStep((current) => Math.min(Math.max(current, 0), total));
+  }, [total, setStep]);
+
+  // Start with all messages loaded once the snapshot first arrives; afterwards the user controls `step`.
+  // Skip the auto-init if we restored a value from localStorage so we don't clobber the user's last position.
+  const hasInitializedRef = useRef(false);
+  useEffect(() => {
+    if (!hasInitializedRef.current && total > 0) {
+      hasInitializedRef.current = true;
+      if (!stepHydrated) {
+        setStep(total);
+      }
     }
+  }, [total, stepHydrated, setStep]);
 
-    const runtime = createStoryRuntime(space);
-    runtimeRef.current = runtime;
-    return () => {
-      runtimeRef.current = null;
-      void runtime.dispose();
-    };
-  }, [space]);
+  const visibleMessages = useMemo(() => sortedMessages.slice(0, step), [sortedMessages, step]);
+  const { commits, branches } = useMemo(
+    () => buildExecutionGraph({ traceMessages: visibleMessages }),
+    [visibleMessages],
+  );
 
-  const handleRunAgent = useCallback(() => {
-    const runtime = runtimeRef.current;
-    if (!runtime) {
+  // Auto-play steps forward until we reach the end.
+  useEffect(() => {
+    if (!playing) {
       return;
     }
-
-    const scenarioIndex = invokeCounterRef.current++;
-    void runtime.runPromise(
-      Effect.gen(function* () {
-        const manager = yield* ProcessManager.ProcessManagerService;
-        const handle = yield* manager.spawn(SimulatedAgent);
-        yield* handle.submitInput(scenarioIndex);
-      }),
-    );
+    const id = window.setInterval(() => {
+      setStep((current) => {
+        if (current >= total) {
+          setPlaying(false);
+          return current;
+        }
+        return current + 1;
+      });
+    }, PLAYBACK_INTERVAL_MS);
+    return () => window.clearInterval(id);
+  }, [playing, total]);
+
+  // Handlers.
+  const handleNext = useCallback(() => setStep((current) => Math.min(current + 1, total)), [total]);
+  const handlePrev = useCallback(() => setStep((current) => Math.max(current - 1, 0)), []);
+  const handleReset = useCallback(() => {
+    setStep(0);
+    setPlaying(false);
   }, []);
+  const handleShowAll = useCallback(() => {
+    setStep(total);
+    setPlaying(false);
+  }, [total]);
+  const handleTogglePlay = useCallback(() => setPlaying((previous) => !previous), []);
 
-  if (!space) {
-    return <Loading />;
-  }
+  // Keyboard shortcuts: ←/h step back, →/l step forward, space play/pause, r reset, e/End show all.
+  useEffect(() => {
+    return addEventListener(window, 'keydown', (event: KeyboardEvent) => {
+      const target = event.target as HTMLElement | null;
+      if (target?.tagName === 'INPUT' || target?.tagName === 'TEXTAREA' || target?.isContentEditable) {
+        return;
+      }
+
+      switch (event.key) {
+        case 'ArrowRight':
+        case 'l':
+          event.preventDefault();
+          handleNext();
+          break;
+        case 'ArrowLeft':
+        case 'h':
+          event.preventDefault();
+          handlePrev();
+          break;
+        case ' ':
+          event.preventDefault();
+          handleTogglePlay();
+          break;
+        case 'r':
+          event.preventDefault();
+          handleReset();
+          break;
+        case 'e':
+        case 'End':
+          event.preventDefault();
+          handleShowAll();
+          break;
+      }
+    });
+  }, [handleNext, handlePrev, handleTogglePlay, handleReset, handleShowAll]);
 
   return (
-    <Panel.Root>
-      <Panel.Toolbar asChild>
-        <Toolbar.Root>
-          <IconButton icon='ph--plus--regular' label='Start Agent' onClick={handleRunAgent} />
-        </Toolbar.Root>
-      </Panel.Toolbar>
-      <Panel.Content asChild>
-        <TracePanel space={space} />
-      </Panel.Content>
-    </Panel.Root>
+    <ScrollContainer.Root pin>
+      <BaseStory
+        toolbar={
+          <Toolbar.Root>
+            <IconButton iconOnly icon='ph--skip-back--regular' label='Reset (R)' onClick={handleReset} />
+            <IconButton iconOnly icon='ph--caret-left--regular' label='Step back (← / H)' onClick={handlePrev} />
+            <IconButton
+              iconOnly
+              icon={playing ? 'ph--pause--regular' : 'ph--play--regular'}
+              label={playing ? 'Pause (Space)' : 'Play (Space)'}
+              onClick={handleTogglePlay}
+            />
+            <IconButton iconOnly icon='ph--caret-right--regular' label='Step forward (→ / L)' onClick={handleNext} />
+            <IconButton iconOnly icon='ph--skip-forward--regular' label='Show all (E / End)' onClick={handleShowAll} />
+            <Toolbar.Text className='text-right text-sm tabular-nums opacity-70'>
+              {step} / {total}
+            </Toolbar.Text>
+          </Toolbar.Root>
+        }
+      >
+        <ScrollContainer.Content thin>
+          <ScrollContainer.Viewport>
+            <Timeline branches={branches} commits={commits} showTimestamp />
+          </ScrollContainer.Viewport>
+          <ScrollContainer.ScrollDownButton />
+          <ScrollContainer.Fade />
+        </ScrollContainer.Content>
+      </BaseStory>
+    </ScrollContainer.Root>
   );
 };
 
 const meta = {
   title: 'plugins/plugin-assistant/containers/TracePanel',
+  parameters: {
+    layout: 'fullscreen',
+    translations,
+  },
+} satisfies Meta;
+
+export default meta;
+
+type Story = StoryObj<typeof meta>;
+
+export const Default: Story = {
   render: DefaultStory,
   decorators: [
     withTheme(),
-    withLayout({ layout: 'column' }),
+    withLayout({ layout: 'column', classNames: 'w-(--dx-complementary-sidebar-size) overflow-hidden' }),
     withPluginManager({
       plugins: [
         ...corePlugins(),
@@ -214,17 +276,72 @@ const meta = {
               yield* Effect.promise(() => space.waitUntilReady());
             }),
         }),
+        AssistantPlugin(),
+        AutomationPlugin(),
       ],
     }),
   ],
-  parameters: {
-    layout: 'fullscreen',
-    translations,
-  },
-} satisfies Meta;
+};
 
-export default meta;
+export const WithSubAgentFixture: Story = {
+  render: FixtureStory,
+  decorators: [withTheme(), withLayout({ layout: 'column', classNames: 'w-(--dx-complementary-sidebar-size)' })],
+};
 
-type Story = StoryObj<typeof meta>;
+export const WithSnapshot: Story = {
+  render: SnapshotStory,
+  decorators: [
+    withTheme(),
+    withLayout({ layout: 'column', classNames: 'w-(--dx-complementary-sidebar-size)' }),
+    withPluginManager({
+      plugins: [
+        ...corePlugins(),
+        ClientPlugin({
+          types: [Feed.Feed, Trace.Message],
+          onClientInitialized: initClientFromSpaceSnapshot(() => import('../../testing/data/trace-timeline.dx.json')),
+        }),
+        AutomationPlugin(),
+      ],
+    }),
+  ],
+};
+
+export const WithRemoteSnapshot: Story = {
+  render: SnapshotStory,
+  decorators: [
+    withTheme(),
+    withLayout({ layout: 'column', classNames: 'w-(--dx-complementary-sidebar-size)' }),
+    withPluginManager({
+      plugins: [
+        ...corePlugins(),
+        ClientPlugin({
+          types: [Feed.Feed, Trace.Message],
+          onClientInitialized: initClientFromSpaceSnapshot(
+            () => import('../../testing/data/trace-timeline-remote.dx.json'),
+          ),
+        }),
+        AutomationPlugin(),
+      ],
+    }),
+  ],
+};
 
-export const Default: Story = {};
+export const WithRemoteMultipleSnapshot: Story = {
+  render: SnapshotStory,
+  decorators: [
+    withTheme(),
+    withLayout({ layout: 'column', classNames: 'w-(--dx-complementary-sidebar-size)' }),
+    withPluginManager({
+      plugins: [
+        ...corePlugins(),
+        ClientPlugin({
+          types: [Feed.Feed, Trace.Message],
+          onClientInitialized: initClientFromSpaceSnapshot(
+            () => import('../../testing/data/trace-timeline-multiple.dx.json'),
+          ),
+        }),
+        AutomationPlugin(),
+      ],
+    }),
+  ],
+};