@dxos/assistant 0.8.4-main.e098934 → 0.8.4-main.e8ec1fe

package/src/session/format.ts

@@ -2,21 +2,23 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Array, Effect, Option, pipe } from 'effect';
+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 { Template } from '@dxos/blueprints';
 import { Obj } from '@dxos/echo';
-import { ObjectVersion } from '@dxos/echo-db';
-import { type ObjectId } from '@dxos/echo-schema';
-import { DatabaseService } from '@dxos/functions';
+import { type ObjectId } from '@dxos/echo/internal';
+import { DatabaseService, ObjectVersion } from '@dxos/echo-db';
 import { log } from '@dxos/log';
-import { type ContentBlock, DataType } from '@dxos/schema';
+import { type ContentBlock, Message } from '@dxos/types';
 import { trim } from '@dxos/util';
 
 import { AiAssistantError } from '../errors';
 
 import { ArtifactDiffResolver } from './artifact-diff';
-import { type AiSessionRunParams } from './session';
+import { type AiSessionRunError, type AiSessionRunParams } from './session';
 
 /**
  * Formats the system prompt.
@@ -28,7 +30,7 @@ export const formatSystemPrompt = ({
   objects = [],
 }: Pick<AiSessionRunParams<any>, 'system' | 'blueprints' | 'objects'>) =>
   Effect.gen(function* () {
-    const blueprintDefs = yield* pipe(
+    const blueprintDefs = yield* Function.pipe(
       blueprints,
       Effect.forEach((blueprint) => Effect.succeed(blueprint.instructions)),
       Effect.flatMap(Effect.forEach((template) => DatabaseService.loadOption(template.source))),
@@ -47,7 +49,7 @@ export const formatSystemPrompt = ({
       ),
     );
 
-    const objectDefs = yield* pipe(
+    const objectDefs = yield* Function.pipe(
       objects,
       Effect.forEach((object) =>
         Effect.succeed(trim`
@@ -60,7 +62,7 @@ export const formatSystemPrompt = ({
       Effect.map((objects) => (objects.length > 0 ? ['## Context Objects', ...objects].join('\n\n') : undefined)),
     );
 
-    return yield* pipe(
+    return yield* Function.pipe(
       Effect.succeed([system, blueprintDefs, objectDefs].filter((def): def is string => def !== undefined)),
       Effect.map((parts) => parts.join('\n\n')),
     );
@@ -71,7 +73,10 @@ export const formatSystemPrompt = ({
  */
 // TODO(burdon): Move to AiPreprocessor.
 // TODO(burdon): Convert util below to `Effect.fn` (to preserve stack info)
-export const formatUserPrompt = ({ prompt, history = [] }: Pick<AiSessionRunParams<any>, 'prompt' | 'history'>) =>
+export const formatUserPrompt = ({
+  prompt,
+  history = [],
+}: Pick<AiSessionRunParams<any>, 'prompt' | 'history'>): Effect.Effect<Message.Message, AiSessionRunError> =>
   Effect.gen(function* () {
     const blocks: ContentBlock.Any[] = [];
 
@@ -102,14 +107,14 @@ export const formatUserPrompt = ({ prompt, history = [] }: Pick<AiSessionRunPara
       }
     }
 
-    return Obj.make(DataType.Message, {
+    return Obj.make(Message.Message, {
       created: new Date().toISOString(),
       sender: { role: 'user' },
       blocks: [...blocks, { _tag: 'text', text: prompt }],
     });
   });
 
-const gatherObjectVersions = (messages: DataType.Message[]): Map<ObjectId, ObjectVersion> => {
+const gatherObjectVersions = (messages: Message.Message[]): Map<ObjectId, ObjectVersion> => {
   const artifactIds = new Map<ObjectId, ObjectVersion>();
   for (const message of messages) {
     for (const block of message.blocks) {

package/src/session/index.ts

@@ -7,3 +7,5 @@ export * from './observer';
 export * from './session';
 export * from './toolkit';
 export * from './version-pin';
+
+export * as GenericToolkit from './GenericToolkit';

package/src/session/observer.ts

@@ -2,10 +2,11 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Effect, Function } from 'effect';
+import * as Effect from 'effect/Effect';
+import * as Function from 'effect/Function';
 
 import { type AiParser, type ConsolePrinter } from '@dxos/ai';
-import { type DataType } from '@dxos/schema';
+import { type Message } from '@dxos/types';
 
 /**
  * Live observer of the generation process.
@@ -15,7 +16,7 @@ export interface GenerationObserver extends AiParser.ParseResponseCallbacks {
    * Complete messages fired during the session, both from the model and from the user.
    * This message will contain all message blocks emitted through the `onBlock` callback.
    */
-  onMessage: (message: DataType.Message) => Effect.Effect<void>;
+  onMessage: (message: Message.Message) => Effect.Effect<void>;
 }
 
 export const GenerationObserver = Object.freeze({

package/src/session/session.conversations.json

@@ -0,0 +1 @@
+{"conversations":[{"parameters":{"model":"@anthropic/claude-3-5-sonnet-20241022","stream":true,"tools":[]},"prompt":{"content":[{"role":"user","content":"Hello world!","options":{}}]},"response":[{"type":"response-metadata","id":"msg_01Bm496JGgDMgC9YwdYTFFHo","modelId":"claude-3-5-sonnet-20241022","timestamp":"1970-01-01T00:00:00.000Z","metadata":{}},{"type":"text-start","id":"0","metadata":{}},{"type":"text-delta","id":"0","delta":"Hi","metadata":{}},{"type":"text-delta","id":"0","delta":"!","metadata":{}},{"type":"text-delta","id":"0","delta":" How can","metadata":{}},{"type":"text-delta","id":"0","delta":" I help you today?","metadata":{}},{"type":"text-end","id":"0","metadata":{}},{"type":"finish","reason":"stop","usage":{"inputTokens":10,"outputTokens":12,"totalTokens":22},"metadata":{"anthropic":{"usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":10,"output_tokens":1,"server_tool_use":null,"service_tier":"standard"}}}}]},{"parameters":{"model":"@anthropic/claude-3-5-sonnet-20241022","stream":true,"tools":[{"name":"Calculator","description":"Basic calculator tool","inputSchema":{"type":"object","required":["input"],"properties":{"input":{"type":"string","description":"The calculation to perform."}},"additionalProperties":false}}]},"prompt":{"content":[{"role":"user","content":"What is 10 + 30?","options":{}}]},"response":[{"type":"response-metadata","id":"msg_01UNTbjtTHEpdFgGunu2k8Br","modelId":"claude-3-5-sonnet-20241022","timestamp":"1970-01-01T00:00:00.000Z","metadata":{}},{"type":"text-start","id":"0","metadata":{}},{"type":"text-delta","id":"0","delta":"I","metadata":{}},{"type":"text-delta","id":"0","delta":"'ll help","metadata":{}},{"type":"text-delta","id":"0","delta":" you calculate","metadata":{}},{"type":"text-delta","id":"0","delta":" ","metadata":{}},{"type":"text-delta","id":"0","delta":"10 + 30","metadata":{}},{"type":"text-delta","id":"0","delta":" using the","metadata":{}},{"type":"text-delta","id":"0","delta":" Calculator tool.","metadata":{}},{"type":"text-end","id":"0","metadata":{}},{"type":"tool-params-start","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","name":"Calculator","providerExecuted":false,"metadata":{}},{"type":"tool-params-delta","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","delta":"","metadata":{}},{"type":"tool-params-delta","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","delta":"{\"input\": \"","metadata":{}},{"type":"tool-params-delta","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","delta":"10 + 30","metadata":{}},{"type":"tool-params-delta","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","delta":"\"}","metadata":{}},{"type":"tool-params-end","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","metadata":{}},{"type":"tool-call","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","name":"Calculator","params":{"input":"10 + 30"},"providerExecuted":false,"metadata":{}},{"type":"finish","reason":"tool-calls","usage":{"inputTokens":394,"outputTokens":73,"totalTokens":467},"metadata":{"anthropic":{"usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":394,"output_tokens":1,"server_tool_use":null,"service_tier":"standard"}}}}]},{"parameters":{"model":"@anthropic/claude-3-5-sonnet-20241022","stream":true,"tools":[{"name":"Calculator","description":"Basic calculator tool","inputSchema":{"type":"object","required":["input"],"properties":{"input":{"type":"string","description":"The calculation to perform."}},"additionalProperties":false}}]},"prompt":{"content":[{"role":"user","content":"What is 10 + 30?","options":{}},{"role":"assistant","content":[{"type":"text","text":"I'll help you calculate 10 + 30 using the Calculator tool.","options":{}},{"type":"tool-call","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","name":"Calculator","params":{"input":"10 + 30"},"providerExecuted":false,"options":{}}],"options":{}},{"role":"tool","content":[{"type":"tool-result","id":"toolu_01LTGUgfihCQtVchm1b37xHQ","name":"Calculator","result":{"result":40},"options":{}}],"options":{}}]},"response":[{"type":"response-metadata","id":"msg_01AKio3GheJGn38BhvAp33rs","modelId":"claude-3-5-sonnet-20241022","timestamp":"1970-01-01T00:00:00.000Z","metadata":{}},{"type":"text-start","id":"0","metadata":{}},{"type":"text-delta","id":"0","delta":"The result","metadata":{}},{"type":"text-delta","id":"0","delta":" of 10 +","metadata":{}},{"type":"text-delta","id":"0","delta":" 30 is ","metadata":{}},{"type":"text-delta","id":"0","delta":"40.","metadata":{}},{"type":"text-end","id":"0","metadata":{}},{"type":"finish","reason":"stop","usage":{"inputTokens":482,"outputTokens":18,"totalTokens":500},"metadata":{"anthropic":{"usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":482,"output_tokens":3,"server_tool_use":null,"service_tier":"standard"}}}}]}]}

package/src/session/session.test.ts

@@ -2,13 +2,17 @@
 // Copyright 2025 DXOS.org
 //
 
-import { AiTool, AiToolkit } from '@effect/ai';
+import * as Tool from '@effect/ai/Tool';
+import * as Toolkit from '@effect/ai/Toolkit';
 import { describe, it } from '@effect/vitest';
-import { Effect, Layer, Schema } from 'effect';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+import * as Schema from 'effect/Schema';
 
-import { AiService, ToolExecutionService, ToolResolverService } from '@dxos/ai';
+import { AiService, MemoizedAiService, ToolExecutionService, ToolResolverService } from '@dxos/ai';
 import { AiServiceTestingPreset } from '@dxos/ai/testing';
 import { Obj, Type } from '@dxos/echo';
+import { TestHelpers } from '@dxos/effect';
 import { TracingService } from '@dxos/functions';
 import { log } from '@dxos/log';
 
@@ -29,8 +33,8 @@ const CalendarEventSchema = Schema.Struct({
 
 type CalendarEvent = Schema.Schema.Type<typeof CalendarEventSchema>;
 
-class TestToolkit extends AiToolkit.make(
-  AiTool.make('Calculator', {
+const TestToolkit = Toolkit.make(
+  Tool.make('Calculator', {
     description: 'Basic calculator tool',
     parameters: {
       input: Schema.String.annotations({
@@ -42,7 +46,7 @@ class TestToolkit extends AiToolkit.make(
     }),
     failure: Schema.Never,
   }),
-) {}
+);
 
 // Tool handlers.
 const toolkitLayer = TestToolkit.toLayer({
@@ -61,175 +65,52 @@ const toolkitLayer = TestToolkit.toLayer({
     }),
 });
 
-describe.runIf(process.env.DX_RUN_SLOW_TESTS)('AiSession', () => {
-  it.effect('no tools', () =>
-    Effect.gen(function* () {
-      const session = new AiSession({ operationModel: 'configured' });
-      const response = yield* session.run({
-        prompt: 'Hello world!',
-        history: [],
-      });
-      log.info('response', { response });
-    }).pipe(
-      Effect.provide(
-        Layer.mergeAll(
-          toolkitLayer,
-          AiService.model('@anthropic/claude-3-5-sonnet-20241022').pipe(
-            Layer.provideMerge(ToolResolverService.layerEmpty),
-            Layer.provideMerge(ToolExecutionService.layerEmpty),
-            Layer.provideMerge(AiServiceTestingPreset('direct')),
-            Layer.provideMerge(TracingService.layerNoop),
-          ),
-        ),
-      ),
-    ),
-  );
+const TestLayer = Layer.mergeAll(
+  toolkitLayer,
+  AiService.model('@anthropic/claude-3-5-sonnet-20241022'),
+  TracingService.layerNoop,
+  ToolResolverService.layerEmpty,
+  ToolExecutionService.layerEmpty,
+).pipe(
+  //
+  Layer.provideMerge(MemoizedAiService.layerTest()),
+  Layer.provide(AiServiceTestingPreset('direct')),
+);
 
-  it.effect('calculator', () =>
-    Effect.gen(function* () {
-      const session = new AiSession({ operationModel: 'configured' });
-      const toolkit = yield* TestToolkit;
-      const response = yield* session.run({
-        toolkit,
-        prompt: 'What is 10 + 20?',
-        history: [],
-      });
-      log.info('response', { response });
-    }).pipe(
-      Effect.provide(
-        Layer.mergeAll(
-          TracingService.layerNoop,
-          AiService.model('@anthropic/claude-3-5-sonnet-20241022').pipe(
-            Layer.provideMerge(AiServiceTestingPreset('direct')),
-          ),
-          ToolResolverService.layerEmpty,
-          ToolExecutionService.layerEmpty,
-          toolkitLayer,
-        ),
-      ),
+describe('AiSession', () => {
+  it.effect(
+    'no tools',
+    Effect.fnUntraced(
+      function* (_) {
+        const session = new AiSession({ operationModel: 'configured' });
+        const response = yield* session.run({
+          prompt: 'Hello world!',
+          history: [],
+        });
+        log.info('response', { response });
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
     ),
   );
 
-  // TODO(dmaretskyi): Revive test.
-  /*
-  it.skip('create calendar itinerary', { timeout: 60_000 }, async () => {
-    // overrides: { model: 'llama3.1:8b' },
-    const session = new AiSession({ operationModel: 'configured' });
-    const objects = new Set<string>();
-
-    // Define calendar artifact.
-    const calendarArtifact = defineArtifact({
-      id: 'artifact:dxos.org/plugin/calendar',
-      name: 'Calendar',
-      instructions: 'Use this to create and query calendar events.',
-      schema: CalendarEventSchema,
-      tools: [
-        createTool('calendar', {
-          name: 'query',
-          description: 'Query the calendar for events',
-          schema: Schema.Struct({}),
-          execute: async () => {
-            return ToolResult.Success(CALENDAR_EVENTS);
-          },
-        }),
-      ],
-    });
-
-    const tableArtifact = defineArtifact({
-      id: 'artifact:dxos.org/plugin/table',
-      name: 'Table',
-      instructions: 'Use this to create and manage tables. Each table has a unique id.',
-      schema: Schema.Struct({}),
-      tools: [
-        createTool('table', {
-          name: 'create',
-          description: 'Create a table',
-          schema: Schema.Struct({
-            data: Schema.Array(Schema.Any).annotations({ description: 'Array of data payloads to add as rows' }),
-          }),
-          execute: async ({ data }) => {
-            log('create table', { data });
-            const id = DXN.fromLocalObjectId(ObjectId.random()).toString();
-            objects.add(id);
-            // TODO(dmaretskyi): consider xml for refs instead of @dxn:echo:@:XXXXX
-            return ToolResult.Success(`table @${id}`);
-          },
-        }),
-      ],
-    });
-
-    const mapArtifact = defineArtifact({
-      id: 'artifact:dxos.org/plugin/map',
-      name: 'Map',
-      instructions:
-        'Use this to create and manage maps. Maps source data from tables. Table id is required to create a map.',
-      schema: Schema.Struct({}),
-      tools: [
-        createTool('map', {
-          name: 'create',
-          description: 'Create a map',
-          schema: Schema.Struct({
-            source: ArtifactId.annotations({
-              description: 'The table that will be used as the source of the map',
-            }),
-          }),
-          execute: async ({ source }) => {
-            // TODO(dmaretskyi): Use effect-schema decode instead of manual parsing.
-            const sourceId = ArtifactId.toDXN(source);
-            if (!objects.has(sourceId.toString())) {
-              return ToolResult.Error(`table id=${source} not found`);
-            }
-            log('create map', { sourceId });
-            const id = DXN.fromLocalObjectId(ObjectId.random()).toString();
-            objects.add(id);
-            // TODO(dmaretskyi): consider xml for refs instead of @dxn:echo:@:XXXXX
-            return ToolResult.Success(`map @${id}`);
-          },
-        }),
-      ],
-    });
-
-    const scriptArtifact = defineArtifact({
-      id: 'artifact:dxos.org/plugin/script',
-      name: 'Script',
-      instructions: 'Use this to create and manage scripts',
-      schema: Schema.Struct({}),
-      tools: [],
-    });
-
-    // session.streamEvent.on((event) => {
-    //   printStreamEvent(event);
-    // });
-
-    const printer = new ConsolePrinter();
-    session.message.on((message) => printer.printMessage(message));
-    session.userMessage.on((message) => printer.printMessage(message));
-    session.block.on((block) => printer.printContentBlock(block));
-
-    // session.update.on((update) => {
-    //   log('update', { update });
-    // });
-
-    // Test creating an itinerary
-    const response = await session.run({
-      client: aiClient,
-      tools: [],
-      artifacts: [calendarArtifact, tableArtifact, mapArtifact, scriptArtifact],
-      requiredArtifactIds: [calendarArtifact.id, tableArtifact.id, mapArtifact.id, scriptArtifact.id],
-      history: [],
-      generationOptions: {
-        model: '@anthropic/claude-3-5-haiku-20241022',
+  it.effect(
+    'calculator',
+    Effect.fnUntraced(
+      function* (_) {
+        const session = new AiSession({ operationModel: 'configured' });
+        const toolkit = yield* TestToolkit;
+        const response = yield* session.run({
+          toolkit,
+          prompt: 'What is 10 + 30?',
+          history: [],
+        });
+        log.info('response', { response });
       },
-      prompt: 'create a table and map for a travel itinerary based on events in my calendar',
-      toolResolver: new ToolRegistry([]),
-    });
-
-    log('result', {
-      objects,
-      finalMessage: response.at(-1),
-    });
-  });
-  */
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+  );
 });
 
 // Travel to rome, florence, livorno, siena, madrid for conferences
@@ -266,43 +147,3 @@ const _CALENDAR_EVENTS: CalendarEvent[] = [
     description: 'Travel to Madrid',
   }),
 ];
-
-// const printStreamEvent = (event: GenerationStreamEvent) => {
-//   switch (event.type) {
-//     case 'message_start': {
-//       process.stdout.write(`${event.message.role.toUpperCase()}\n\n`);
-//       for (const content of event.message.content) {
-//         printContentBlock(content);
-//       }
-//       break;
-//     }
-//     case 'content_block_start': {
-//       printContentBlock(event.content);
-//       break;
-//     }
-//     case 'content_block_delta': {
-//       switch (event.delta.type) {
-//         case 'text_delta': {
-//           process.stdout.write(event.delta.text);
-//           break;
-//         }
-//         case 'input_json_delta': {
-//           process.stdout.write(event.delta.partial_json);
-//           break;
-//         }
-//       }
-//       break;
-//     }
-//     case 'content_block_stop': {
-//       process.stdout.write('\n');
-//       break;
-//     }
-//     case 'message_delta': {
-//       break;
-//     }
-//     case 'message_stop': {
-//       process.stdout.write('\n\n');
-//       break;
-//     }
-//   }
-// };

package/src/session/session.ts

@@ -2,14 +2,20 @@
 // Copyright 2025 DXOS.org
 //
 
-import { type AiError, AiLanguageModel, type AiTool, type AiToolkit } from '@effect/ai';
-import { Chunk, Effect, type Schema, Stream } from 'effect';
+import type * as AiError from '@effect/ai/AiError';
+import * as LanguageModel from '@effect/ai/LanguageModel';
+import type * as Tool from '@effect/ai/Tool';
+import type * as Toolkit from '@effect/ai/Toolkit';
+import * as Chunk from 'effect/Chunk';
+import * as Effect from 'effect/Effect';
+import type * as Schema from 'effect/Schema';
+import * as Stream from 'effect/Stream';
 
 import {
-  type AiInputPreprocessingError,
   AiParser,
   AiPreprocessor,
   type AiToolNotFoundError,
+  type PromptPreprocessingError,
   type ToolExecutionService,
   type ToolResolverService,
   callTool,
@@ -20,29 +26,29 @@ import { todo } from '@dxos/debug';
 import { Obj } from '@dxos/echo';
 import { TracingService } from '@dxos/functions';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { Message } from '@dxos/types';
 
 import { type AiAssistantError } from '../errors';
 
-import { mapAiError } from './error-handling';
 import { formatSystemPrompt, formatUserPrompt } from './format';
 import { GenerationObserver } from './observer';
 
-export type AiSessionRunError = AiError.AiError | AiInputPreprocessingError | AiToolNotFoundError | AiAssistantError;
+export type AiSessionRunError = AiError.AiError | PromptPreprocessingError | AiToolNotFoundError | AiAssistantError;
 
 export type AiSessionRunRequirements =
-  | AiLanguageModel.AiLanguageModel
+  | LanguageModel.LanguageModel
   | ToolExecutionService
   | ToolResolverService
   | TracingService;
 
-export type AiSessionRunParams<Tools extends AiTool.Any> = {
+export type AiSessionRunParams<Tools extends Record<string, Tool.Any>> = {
   prompt: string;
+  // TODO(wittjosiah): Rename to systemPrompt.
   system?: string;
-  history?: DataType.Message[];
+  history?: Message.Message[];
   objects?: Obj.Any[];
   blueprints?: Blueprint.Blueprint[];
-  toolkit?: AiToolkit.ToHandler<Tools>;
+  toolkit?: Toolkit.WithHandler<Tools>;
   observer?: GenerationObserver;
 };
 
@@ -64,14 +70,14 @@ export class AiSession {
 
   /** Message history from queue. */
   // TODO(burdon): Evolve into supporting a git-like graph of messages.
-  private _history: DataType.Message[] = [];
+  private _history: Message.Message[] = [];
 
   /** Pending messages for this session (incl. the current prompt). */
-  private _pending: DataType.Message[] = [];
+  private _pending: Message.Message[] = [];
 
   constructor(private readonly _options: AiSessionOptions = {}) {}
 
-  run = <Tools extends AiTool.Any>({
+  run = <Tools extends Record<string, Tool.Any>>({
     prompt,
     system: systemTemplate,
     history = [],
@@ -79,13 +85,13 @@ export class AiSession {
     blueprints = [],
     toolkit,
     observer = GenerationObserver.noop(),
-  }: AiSessionRunParams<Tools>): Effect.Effect<DataType.Message[], AiSessionRunError, AiSessionRunRequirements> =>
+  }: AiSessionRunParams<Tools>): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements> =>
     Effect.gen(this, function* () {
       this._history = [...history];
       this._pending = [];
       const pending = this._pending;
 
-      const submitMessage = Effect.fnUntraced(function* (message: DataType.Message) {
+      const submitMessage = Effect.fnUntraced(function* (message: Message.Message) {
         pending.push(message);
         yield* observer.onMessage(message);
         yield* TracingService.emitConverationMessage(message);
@@ -109,21 +115,16 @@ export class AiSession {
           objects: objects?.length ?? 0,
         });
 
-        const prompt = yield* AiPreprocessor.preprocessAiInput([...this._history, ...this._pending]);
+        const prompt = yield* AiPreprocessor.preprocessPrompt([...this._history, ...this._pending], { system });
 
         // Execute the stream request.
-        const blocks = yield* AiLanguageModel.streamText({
+        // logDump('prompt', Prompt.Prompt.pipe(Schema.encodeSync)(prompt));
+        const blocks = yield* LanguageModel.streamText({
           prompt,
-          system,
           toolkit,
           disableToolCallResolution: true,
         }).pipe(
-          Stream.catchTag(
-            'AiError',
-            Effect.fnUntraced(function* (err) {
-              return yield* Effect.fail(yield* mapAiError(err));
-            }),
-          ),
+          // TOOD(dmaretskyi): Error mapping.
           AiParser.parseResponse({
             onBegin: () => observer.onBegin(),
             onBlock: (block) => observer.onBlock(block),
@@ -136,7 +137,7 @@ export class AiSession {
 
         // Create the response message.
         const response = yield* submitMessage(
-          Obj.make(DataType.Message, {
+          Obj.make(Message.Message, {
             created: new Date().toISOString(),
             sender: { role: 'assistant' },
             blocks,
@@ -166,10 +167,12 @@ export class AiSession {
         });
 
         // Add to queue and continue loop.
+        // TODO(wittjosiah): Sometimes tool error results are added to the queue before the tool agent statuses.
+        //   This results in a broken execution graph.
         yield* submitMessage(
-          Obj.make(DataType.Message, {
+          Obj.make(Message.Message, {
             created: new Date().toISOString(),
-            sender: { role: 'user' },
+            sender: { role: 'tool' },
             blocks: toolResults,
           }),
         );
@@ -185,7 +188,7 @@ export class AiSession {
   // TODO(burdon): Implement or remove.
   async runStructured<S extends Schema.Schema.AnyNoContext>(
     _schema: S,
-    _options: AiSessionRunParams<AiTool.Any>,
+    _options: AiSessionRunParams<any>,
   ): Promise<Schema.Schema.Type<S>> {
     return todo();
     // const parser = structuredOutputParser(schema);
@@ -199,3 +202,12 @@ export class AiSession {
 
 const createSnippet = (text: string, len = 32) =>
   text.length <= len * 2 ? text : [text.slice(0, len), '...', text.slice(-len)].join('');
+
+// TODO(dmaretskyi): Extract as a general util.
+const logDump = (message: string, data: unknown) => {
+  // eslint-disable-next-line @typescript-eslint/no-require-imports
+  const { writeFileSync } = require('node:fs');
+  const path = `/tmp/log-data-${Date.now()}.json`;
+  writeFileSync(path, JSON.stringify(data, null, 2));
+  log.info(message, { path });
+};

package/src/session/toolkit.ts

@@ -2,15 +2,17 @@
 // Copyright 2025 DXOS.org
 //
 
-import { type AiTool, AiToolkit } from '@effect/ai';
-import { type Context, Effect } from 'effect';
+import type * as Tool from '@effect/ai/Tool';
+import * as Toolkit from '@effect/ai/Toolkit';
+import type * as Context from 'effect/Context';
+import * as Effect from 'effect/Effect';
 
-import { ToolExecutionService, type ToolId, ToolResolverService } from '@dxos/ai';
+import { type AiToolNotFoundError, ToolExecutionService, type ToolId, ToolResolverService } from '@dxos/ai';
 import { type Blueprint } from '@dxos/blueprints';
-import { isNotFalsy } from '@dxos/util';
+import { isTruthy } from '@dxos/util';
 
-export type ToolkitParams<Tools extends AiTool.Any> = {
-  toolkit?: AiToolkit.AiToolkit<Tools>;
+export type ToolkitParams<Tools extends Record<string, Tool.Any>> = {
+  toolkit?: Toolkit.Toolkit<Tools>;
   toolIds?: ToolId[];
   blueprints?: Blueprint.Blueprint[];
 };
@@ -18,25 +20,29 @@ export type ToolkitParams<Tools extends AiTool.Any> = {
 /**
  * Build a combined toolkit from the blueprint tools and the provided toolkit.
  */
-export const createToolkit = <Tools extends AiTool.Any>({
+export const createToolkit = <Tools extends Record<string, Tool.Any> = {}>({
   toolkit: toolkitParam,
   toolIds = [],
   blueprints = [],
-}: ToolkitParams<Tools>) =>
+}: ToolkitParams<Tools>): Effect.Effect<
+  Toolkit.WithHandler<any>,
+  AiToolNotFoundError,
+  ToolResolverService | ToolExecutionService | Tool.HandlersFor<Tools>
+> =>
   Effect.gen(function* () {
     const blueprintToolkit = yield* ToolResolverService.resolveToolkit([
       ...blueprints.flatMap(({ tools }) => tools),
       ...toolIds,
     ]);
 
-    const blueprintToolkitHandler: Context.Context<AiTool.ToHandler<AiTool.Any>> = yield* blueprintToolkit.toContext(
+    const blueprintToolkitHandler: Context.Context<Tool.HandlersFor<any>> = yield* blueprintToolkit.toContext(
       ToolExecutionService.handlersFor(blueprintToolkit),
     );
 
-    const toolkit = AiToolkit.merge(...[toolkitParam, blueprintToolkit].filter(isNotFalsy));
-    return yield* toolkit.pipe(Effect.provide(blueprintToolkitHandler)) as Effect.Effect<
-      AiToolkit.ToHandler<any>,
+    const toolkit = Toolkit.merge(...[toolkitParam, blueprintToolkit].filter(isTruthy));
+    return yield* toolkit.pipe(Effect.provide(blueprintToolkitHandler)) as any as Effect.Effect<
+      Toolkit.WithHandler<any>,
       never,
-      AiTool.ToHandler<Tools>
+      Tool.HandlersFor<Tools>
     >;
   });