@dxos/assistant 0.8.4-main.b97322e → 0.8.4-main.bc674ce

package/src/session/format.test.ts

@@ -0,0 +1,124 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, expect, it } from '@effect/vitest';
+import * as Effect from 'effect/Effect';
+import * as Schema from 'effect/Schema';
+
+import { Blueprint, Template } from '@dxos/blueprints';
+import { Database, Obj, Query } from '@dxos/echo';
+import { TestHelpers } from '@dxos/effect/testing';
+import { defineFunction } from '@dxos/functions';
+import { Text } from '@dxos/schema';
+import { Organization } from '@dxos/types';
+import { trim } from '@dxos/util';
+
+import { createSystemPrompt } from '../templates/system';
+import { AssistantTestLayer } from '../testing';
+
+import { formatSystemPrompt } from './format';
+
+const organizationListFunction = defineFunction({
+  key: 'dxos.org/function/organization-list',
+  name: 'Organization List',
+  description: 'List organizations',
+  inputSchema: Schema.Struct({}),
+  outputSchema: Schema.Array(Schema.String),
+  handler: Effect.fnUntraced(function* () {
+    const organizations = yield* Database.Service.runQuery(Query.type(Organization.Organization));
+    return organizations.map((organization) => organization.name);
+  }),
+});
+
+const TestLayer = AssistantTestLayer({
+  types: [Text.Text, Organization.Organization, Blueprint.Blueprint],
+  functions: [organizationListFunction],
+});
+
+describe('format', () => {
+  it.effect(
+    'should format',
+    Effect.fnUntraced(
+      function* (_) {
+        const { db } = yield* Database.Service;
+        const object = db.add(
+          Obj.make(Organization.Organization, {
+            name: 'Test',
+            website: 'https://www.test.com',
+          }),
+        );
+
+        const blueprint = db.add(
+          Blueprint.make({
+            key: 'example.com/blueprint/test',
+            name: 'Test',
+            instructions: Template.make({
+              source: trim`
+                Test
+                This is the test blueprint.
+              `,
+            }),
+          }),
+        );
+
+        const output = yield* formatSystemPrompt({
+          system: createSystemPrompt({}),
+          blueprints: [blueprint],
+          objects: [object],
+        });
+
+        console.log(output);
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+  );
+
+  it.effect(
+    'should format with function input',
+    Effect.fnUntraced(
+      function* (_) {
+        const { db } = yield* Database.Service;
+        const object = db.add(
+          Obj.make(Organization.Organization, {
+            name: 'Test Org',
+            website: 'https://www.test.com',
+          }),
+        );
+
+        const blueprint = db.add(
+          Blueprint.make({
+            key: 'example.com/blueprint/test',
+            name: 'Test',
+            instructions: Template.make({
+              source: trim`
+                Organization List:
+                {{#each organizations}}
+                  - {{this}}
+                {{/each}}
+                `,
+              inputs: [
+                {
+                  name: 'organizations',
+                  kind: 'function',
+                  function: organizationListFunction.key,
+                },
+              ],
+            }),
+          }),
+        );
+
+        const output = yield* formatSystemPrompt({
+          system: createSystemPrompt({}),
+          blueprints: [blueprint],
+          objects: [object],
+        });
+
+        expect(output).to.include('Test Org');
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+  );
+});

package/src/session/format.ts

@@ -0,0 +1,143 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+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/keys';
+import { log } from '@dxos/log';
+import { type ContentBlock, Message } from '@dxos/types';
+import { trim } from '@dxos/util';
+
+import { AiAssistantError } from '../errors';
+
+import { ArtifactDiffResolver } from './artifact-diff';
+import { type AiSessionRunError, type AiSessionRunProps } from './session';
+
+/**
+ * Formats the system prompt.
+ */
+// TODO(burdon): Move to AiPreprocessor.
+export const formatSystemPrompt = ({
+  system,
+  blueprints = [],
+  objects = [],
+}: Pick<AiSessionRunProps<any>, 'system' | 'blueprints' | 'objects'>) =>
+  Effect.gen(function* () {
+    const blueprintDefs = yield* Function.pipe(
+      blueprints,
+      Effect.forEach((blueprint) => Effect.succeed(blueprint.instructions)),
+      Effect.flatMap(
+        Effect.forEach((template) =>
+          Effect.gen(function* () {
+            return trim`
+            <blueprint>
+              ${yield* Template.processTemplate(template)}
+            </blueprint>
+          `;
+          }),
+        ),
+      ),
+      Effect.map((blueprints) =>
+        blueprints.length > 0 ? ['## Blueprints Definitions', ...blueprints].join('\n\n') : undefined,
+      ),
+    );
+
+    const objectDefs = yield* Function.pipe(
+      objects,
+      Effect.forEach((object) =>
+        Effect.succeed(trim`
+          <object>
+            <dxn>${Obj.getDXN(object)}</dxn>
+            <typename>${Obj.getTypename(object)}</typename>
+          </object>
+        `),
+      ),
+      Effect.map((objects) => (objects.length > 0 ? ['## Context Objects', ...objects].join('\n\n') : undefined)),
+    );
+
+    return yield* Function.pipe(
+      Effect.succeed([system, blueprintDefs, objectDefs].filter((def): def is string => def !== undefined)),
+      Effect.map((parts) => parts.join('\n\n')),
+    );
+  });
+
+/**
+ * Formats the user prompt.
+ */
+// TODO(burdon): Move to AiPreprocessor.
+// TODO(burdon): Convert util below to `Effect.fn` (to preserve stack info)
+export const formatUserPrompt = ({
+  prompt,
+  history = [],
+}: Pick<AiSessionRunProps<any>, 'prompt' | 'history'>): Effect.Effect<Message.Message, AiSessionRunError> =>
+  Effect.gen(function* () {
+    const blocks: ContentBlock.Any[] = [];
+
+    // TODO(dmaretskyi): Evaluate other approaches as `serviceOption` isn't represented in the type system.
+    const artifactDiffResolver = yield* Effect.serviceOption(ArtifactDiffResolver);
+    if (Option.isSome(artifactDiffResolver)) {
+      const versions = gatherObjectVersions(history);
+      const artifactDiff = yield* Effect.tryPromise({
+        try: () =>
+          artifactDiffResolver.value.resolve(
+            [...versions.entries()].map(([id, version]) => ({ id, lastVersion: version })),
+          ),
+        catch: AiAssistantError.wrap({ message: 'Artifact diff resolution error' }),
+      });
+
+      log('version', { artifactDiff, versions });
+      for (const [id, { version }] of [...artifactDiff.entries()]) {
+        if (ObjectVersion.equals(version, versions.get(id)!)) {
+          artifactDiff.delete(id);
+          continue;
+        }
+
+        blocks.push({ _tag: 'anchor', objectId: id, version });
+      }
+
+      if (artifactDiff.size > 0) {
+        blocks.push(createArtifactUpdateBlock(artifactDiff));
+      }
+    }
+
+    return Obj.make(Message.Message, {
+      created: new Date().toISOString(),
+      sender: { role: 'user' },
+      blocks: [...blocks, { _tag: 'text', text: prompt }],
+    });
+  });
+
+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) {
+      if (block._tag === 'anchor') {
+        artifactIds.set(block.objectId, block.version as ObjectVersion);
+      }
+    }
+  }
+
+  return artifactIds;
+};
+
+const createArtifactUpdateBlock = (
+  artifactDiff: Map<ObjectId, { version: ObjectVersion; diff?: string }>,
+): ContentBlock.Any => {
+  return {
+    _tag: 'text',
+    // TODO(dmaretskyi): Does this need to be a special content-block?
+    disposition: 'artifact-update',
+    text: trim`
+      The following artifacts have been updated since the last message:
+      ${[...artifactDiff.entries()]
+        .map(([id, { diff }]) => `<changed-artifact id="${id}">${diff ? `\n${diff}` : ''}</changed-artifact>`)
+        .join('\n')}
+    `,
+  };
+};

package/src/session/index.ts

@@ -2,5 +2,10 @@
 // Copyright 2025 DXOS.org
 //
 
+export * from './artifact-diff';
+export * from './observer';
 export * from './session';
+export * from './toolkit';
 export * from './version-pin';
+
+export * as GenericToolkit from './GenericToolkit';

package/src/session/observer.ts

@@ -0,0 +1,61 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import type * as Tool from '@effect/ai/Tool';
+import * as Effect from 'effect/Effect';
+import * as Function from 'effect/Function';
+
+import { type AiParser, type ConsolePrinter } from '@dxos/ai';
+import { type Message } from '@dxos/types';
+
+/**
+ * Live observer of the generation process.
+ */
+export interface GenerationObserver<Tools extends Record<string, Tool.Any> = any>
+  extends AiParser.ParseResponseCallbacks<Tools> {
+  /**
+   * 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: Message.Message) => Effect.Effect<void>;
+}
+
+export const GenerationObserver = Object.freeze({
+  make: ({
+    onBegin = Function.constant(Effect.void),
+    onPart = Function.constant(Effect.void),
+    onBlock = Function.constant(Effect.void),
+    onEnd = Function.constant(Effect.void),
+    onMessage = Function.constant(Effect.void),
+  }: Partial<GenerationObserver> = {}): GenerationObserver => ({
+    onBegin,
+    onPart,
+    onBlock,
+    onEnd,
+    onMessage,
+  }),
+
+  noop: () => GenerationObserver.make(),
+
+  /**
+   * Debug printer to be used in unit-tests and browser devtools.
+   */
+  fromPrinter: (printer: ConsolePrinter) =>
+    GenerationObserver.make({
+      onBlock: (block) =>
+        Effect.sync(() => {
+          if (block.pending) {
+            return; // Only prints full blocks (better for unit-tests and browser devtools).
+          }
+          printer.printContentBlock(block);
+        }),
+      onMessage: (message) =>
+        Effect.sync(() => {
+          if (message.sender.role === 'assistant') {
+            return; // Skip assistant messages since they are printed in the `onBlock` callback.
+          }
+          printer.printMessage(message);
+        }),
+    }),
+});

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","isFailure":false,"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"}}}}]},{"parameters":{"model":"@anthropic/claude-opus-4-0","stream":true,"tools":[]},"prompt":{"content":[{"role":"user","content":"Hello world!","options":{}}]},"response":[{"type":"response-metadata","id":"msg_01QPPeofAkjiYa83LPCeSYfS","modelId":"claude-opus-4-20250514","timestamp":"1970-01-01T00:00:00.000Z","metadata":{}},{"type":"text-start","id":"0","metadata":{}},{"type":"text-delta","id":"0","delta":"Hello!","metadata":{}},{"type":"text-delta","id":"0","delta":" Welcome","metadata":{}},{"type":"text-delta","id":"0","delta":" to the worl","metadata":{}},{"type":"text-delta","id":"0","delta":"d!","metadata":{}},{"type":"text-delta","id":"0","delta":" 🌍 ","metadata":{}},{"type":"text-delta","id":"0","delta":"How","metadata":{}},{"type":"text-delta","id":"0","delta":" are you doing today? Is there anything","metadata":{}},{"type":"text-delta","id":"0","delta":" specific you'd like to chat about or","metadata":{}},{"type":"text-delta","id":"0","delta":" any questions","metadata":{}},{"type":"text-delta","id":"0","delta":" I can help you with?","metadata":{}},{"type":"text-end","id":"0","metadata":{}},{"type":"finish","reason":"stop","usage":{"inputTokens":10,"outputTokens":40,"totalTokens":50},"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":2,"server_tool_use":null,"service_tier":"standard"}}}}]},{"parameters":{"model":"@anthropic/claude-opus-4-0","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_01EaHNEDYN3Hvp5toY8hQGe6","modelId":"claude-opus-4-20250514","timestamp":"1970-01-01T00:00:00.000Z","metadata":{}},{"type":"text-start","id":"0","metadata":{}},{"type":"text-delta","id":"0","delta":"I'll help","metadata":{}},{"type":"text-delta","id":"0","delta":" you calculate ","metadata":{}},{"type":"text-delta","id":"0","delta":"10 + 30.","metadata":{}},{"type":"text-end","id":"0","metadata":{}},{"type":"tool-params-start","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","name":"Calculator","providerExecuted":false,"metadata":{}},{"type":"tool-params-delta","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","delta":"","metadata":{}},{"type":"tool-params-delta","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","delta":"{\"input\": ","metadata":{}},{"type":"tool-params-delta","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","delta":"\"10 + 3","metadata":{}},{"type":"tool-params-delta","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","delta":"0\"}","metadata":{}},{"type":"tool-params-end","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","metadata":{}},{"type":"tool-call","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","name":"Calculator","params":{"input":"10 + 30"},"providerExecuted":false,"metadata":{}},{"type":"finish","reason":"tool-calls","usage":{"inputTokens":394,"outputTokens":68,"totalTokens":462},"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":3,"server_tool_use":null,"service_tier":"standard"}}}}]},{"parameters":{"model":"@anthropic/claude-opus-4-0","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.","options":{}},{"type":"tool-call","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","name":"Calculator","params":{"input":"10 + 30"},"providerExecuted":false,"options":{}}],"options":{}},{"role":"tool","content":[{"type":"tool-result","id":"toolu_01X4qkrMwCisiUN8Cm4E4Voe","name":"Calculator","isFailure":false,"result":{"result":40},"options":{}}],"options":{}}]},"response":[{"type":"response-metadata","id":"msg_01RjnEkFT3iEUFCpfC4B28aS","modelId":"claude-opus-4-20250514","timestamp":"1970-01-01T00:00:00.000Z","metadata":{}},{"type":"text-start","id":"0","metadata":{}},{"type":"text-delta","id":"0","delta":"10 + 30 = ","metadata":{}},{"type":"text-delta","id":"0","delta":"40","metadata":{}},{"type":"text-end","id":"0","metadata":{}},{"type":"finish","reason":"stop","usage":{"inputTokens":478,"outputTokens":12,"totalTokens":490},"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":478,"output_tokens":8,"server_tool_use":null,"service_tier":"standard"}}}}]}]}

package/src/session/session.test.ts

@@ -2,15 +2,19 @@
 // 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 { AiServiceTestingPreset } from '@dxos/ai/testing';
 import { Obj, Type } from '@dxos/echo';
+import { TestHelpers } from '@dxos/effect/testing';
 import { log } from '@dxos/log';
 
+import { AssistantTestLayer } from '../testing';
+
 import { AiSession } from './session';
 
 // Define a calendar event artifact schema.
@@ -20,7 +24,7 @@ const CalendarEventSchema = Schema.Struct({
   endTime: Schema.String,
   description: Schema.String,
 }).pipe(
-  Type.Obj({
+  Type.object({
     typename: 'example.com/type/CalendarEvent',
     version: '0.1.0',
   }),
@@ -28,8 +32,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({
@@ -41,7 +45,7 @@ class TestToolkit extends AiToolkit.make(
     }),
     failure: Schema.Never,
   }),
-) {}
+);
 
 // Tool handlers.
 const toolkitLayer = TestToolkit.toLayer({
@@ -52,7 +56,7 @@ const toolkitLayer = TestToolkit.toLayer({
         const sanitizedInput = input.replace(/[^0-9+\-*/().\s]/g, '');
         log.info('calculate', { sanitizedInput });
 
-        // eslint-disable-next-line no-new-func, @typescript-eslint/no-implied-eval
+        // eslint-disable-next-line @typescript-eslint/no-implied-eval
         return Function(`"use strict"; return (${sanitizedInput})`)();
       })();
 
@@ -60,182 +64,51 @@ const toolkitLayer = TestToolkit.toLayer({
     }),
 });
 
-describe.runIf(process.env.DX_RUN_SLOW_TESTS)('AiSession', () => {
+const TestLayer = Layer.mergeAll(
+  AssistantTestLayer({
+    types: [CalendarEventSchema],
+  }),
+  toolkitLayer,
+);
+
+describe('AiSession', () => {
   it.effect(
     'no tools',
-    Effect.fn(
-      function* ({ expect }) {
+    Effect.fnUntraced(
+      function* (_) {
         const session = new AiSession({ operationModel: 'configured' });
-
         const response = yield* session.run({
           prompt: 'Hello world!',
           history: [],
         });
         log.info('response', { response });
       },
-      Effect.provide(
-        AiService.model('@anthropic/claude-3-5-sonnet-20241022').pipe(
-          Layer.provideMerge(ToolResolverService.layerEmpty),
-          Layer.provideMerge(ToolExecutionService.layerEmpty),
-          Layer.provideMerge(AiServiceTestingPreset('direct')),
-        ),
-      ),
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
     ),
   );
 
-  it.effect.only(
+  it.effect(
     'calculator',
-    Effect.fn(
-      function* ({ expect }) {
+    Effect.fnUntraced(
+      function* (_) {
         const session = new AiSession({ operationModel: 'configured' });
-
+        const toolkit = yield* TestToolkit;
         const response = yield* session.run({
-          prompt: 'What is 10 + 20?',
+          toolkit,
+          prompt: 'What is 10 + 30?',
           history: [],
-          toolkit: TestToolkit,
         });
         log.info('response', { response });
       },
-      Effect.provide(
-        Layer.mergeAll(
-          toolkitLayer,
-          ToolResolverService.layerEmpty,
-          ToolExecutionService.layerEmpty,
-          AiService.model('@anthropic/claude-3-5-sonnet-20241022').pipe(
-            Layer.provideMerge(AiServiceTestingPreset('direct')),
-          ),
-        ),
-      ),
+      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',
-      },
-      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),
-    });
-  });
-  */
 });
 
 // Travel to rome, florence, livorno, siena, madrid for conferences
-// eslint-disable-next-line no-unused-vars
+
 const _CALENDAR_EVENTS: CalendarEvent[] = [
   Obj.make(CalendarEventSchema, {
     title: 'Exploring Ancient Ruins in Rome',
@@ -268,43 +141,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;
-//     }
-//   }
-// };