@dxos/assistant-toolkit 0.8.4-main.fcc0d83b33 → 0.8.4-staging.60fe92afc8

package/src/blueprints/memory/blueprint.test.ts

@@ -12,17 +12,16 @@ import { Database, Feed, Filter, Obj, Query } from '@dxos/echo';
 import { TestHelpers } from '@dxos/effect/testing';
 import { AgentService } from '@dxos/functions-runtime';
 import { AssistantTestLayer } from '@dxos/functions-runtime/testing';
-import { ObjectId } from '@dxos/keys';
+import { EntityId } from '@dxos/keys';
 
 import { Memory } from '../../types/Memory';
-import { WebSearchBlueprint, WebSearchToolkit } from '../websearch';
+import { WebSearchToolkit } from '../websearch';
 import MemoryBlueprint from './blueprint';
-import { MemoryHandlers } from './functions';
+import { MemoryHandlers } from './operations';
 
-ObjectId.dangerouslyDisableRandomness();
+EntityId.dangerouslyDisableRandomness();
 
 const TestLayer = AssistantTestLayer({
-  aiServicePreset: 'edge-remote',
   operationHandlers: MemoryHandlers,
   types: [Memory, Blueprint.Blueprint, Feed.Feed],
   blueprints: [MemoryBlueprint.make()],
@@ -30,11 +29,10 @@ const TestLayer = AssistantTestLayer({
 });
 
 const TestLayerWithWebSearch = AssistantTestLayer({
-  aiServicePreset: 'edge-remote',
   operationHandlers: MemoryHandlers,
   toolkits: [OpaqueToolkit.make(WebSearchToolkit, Layer.empty)],
   types: [Memory, Blueprint.Blueprint, Feed.Feed],
-  blueprints: [MemoryBlueprint.make(), WebSearchBlueprint.make()],
+  blueprints: [MemoryBlueprint.make()],
   tracing: 'pretty',
 });
 
@@ -48,7 +46,7 @@ describe('Memory Blueprint', () => {
         });
         yield* agent.submitPrompt('Remember that my favorite programming language is TypeScript.');
         yield* agent.waitForCompletion();
-        const memories = yield* Database.runQuery(Query.select(Filter.type(Memory)));
+        const memories = yield* Database.query(Query.select(Filter.type(Memory))).run;
         expect(memories.length).toBeGreaterThanOrEqual(1);
       },
       Effect.provide(TestLayer),
@@ -100,7 +98,7 @@ describe('Memory Blueprint', () => {
         });
         yield* agent.submitPrompt('Delete the memory about "Outdated fact".');
         yield* agent.waitForCompletion();
-        const memories = yield* Database.runQuery(Query.select(Filter.type(Memory)));
+        const memories = yield* Database.query(Query.select(Filter.type(Memory))).run;
         const found = memories.find((memory) => memory.title === 'Outdated fact');
         expect(found).toBeUndefined();
       },
@@ -110,18 +108,19 @@ describe('Memory Blueprint', () => {
     { timeout: 60_000 },
   );
 
-  it.effect(
+  // TODO(dmaretskyi): Flaky. The model does not reliably call save-memory after a provider-executed web search.
+  it.effect.skip(
     'natural: saves memories from a conversation with web search',
     Effect.fnUntraced(
       function* (_) {
         const agent = yield* AgentService.createSession({
-          blueprints: [MemoryBlueprint.make(), WebSearchBlueprint.make()],
+          blueprints: [MemoryBlueprint.make()],
         });
         yield* agent.submitPrompt(
           "I'm going to LA next week. Find me some good hotels and remember the recommendations.",
         );
         yield* agent.waitForCompletion();
-        const memories = yield* Database.runQuery(Query.select(Filter.type(Memory)));
+        const memories = yield* Database.query(Query.select(Filter.type(Memory))).run;
         expect(memories.length).toBeGreaterThanOrEqual(1);
         const hasLAMemory = memories.some(
           (memory) =>

package/src/blueprints/memory/blueprint.ts

@@ -7,7 +7,7 @@ import { Ref } from '@dxos/echo';
 import { Text } from '@dxos/schema';
 import { trim } from '@dxos/util';
 
-import { QueryMemories, SaveMemory, DeleteMemory } from './functions';
+import { QueryMemories, SaveMemory, DeleteMemory } from './operations/definitions';
 
 const BLUEPRINT_KEY = 'org.dxos.blueprint.memory';
 
@@ -28,7 +28,7 @@ const make = () =>
     description: 'Persistent memory storage and retrieval.',
     agentCanEnable: true,
     instructions: {
-      source: Ref.make(Text.make(instructions)),
+      source: Ref.make(Text.make({ content: instructions })),
     },
     tools: Blueprint.toolDefinitions({ operations: [SaveMemory, QueryMemories, DeleteMemory] }),
   });

package/src/blueprints/memory/index.ts

@@ -3,4 +3,4 @@
 //
 
 export { default as MemoryBlueprint } from './blueprint';
-export { QueryMemories, SaveMemory, DeleteMemory, MemoryHandlers } from './functions';
+export { MemoryHandlers, MemoryOperations } from './operations';

package/src/blueprints/memory/{functions → operations}/definitions.ts

@@ -6,15 +6,17 @@ import * as Schema from 'effect/Schema';
 
 import { Operation } from '@dxos/compute';
 import { Database, Ref } from '@dxos/echo';
+import { DXN } from '@dxos/keys';
 
 import { Memory } from '../../../types/Memory';
 
 export const QueryMemories = Operation.make({
   meta: {
-    key: 'org.dxos.function.memory.query',
+    key: DXN.make('org.dxos.function.memory.query'),
     name: 'Query memories',
     description:
       'Search for stored memories using full-text search. Returns memories matching the query terms. Use this to recall previously saved knowledge, facts, or preferences.',
+    icon: 'ph--brain--regular',
   },
   input: Schema.Struct({
     text: Schema.optional(
@@ -36,10 +38,11 @@ export const QueryMemories = Operation.make({
 
 export const SaveMemory = Operation.make({
   meta: {
-    key: 'org.dxos.function.memory.save',
+    key: DXN.make('org.dxos.function.memory.save'),
     name: 'Save memory',
     description:
       'Saves a new memory to the database. Use this to persist knowledge, facts, preferences, or any information that should be remembered across conversations.',
+    icon: 'ph--brain--regular',
   },
   input: Schema.Struct({
     title: Schema.String.annotations({
@@ -56,9 +59,10 @@ export const SaveMemory = Operation.make({
 
 export const DeleteMemory = Operation.make({
   meta: {
-    key: 'org.dxos.function.memory.delete',
+    key: DXN.make('org.dxos.function.memory.delete'),
     name: 'Delete memory',
     description: 'Deletes a memory from the database. Use this to remove outdated or incorrect memories.',
+    icon: 'ph--trash--regular',
   },
   input: Schema.Struct({
     memory: Ref.Ref(Memory),

package/src/blueprints/memory/{functions → operations}/index.ts

@@ -4,7 +4,7 @@
 
 import { OperationHandlerSet } from '@dxos/compute';
 
-export * from './definitions';
+export * as MemoryOperations from './definitions';
 
 export const MemoryHandlers = OperationHandlerSet.lazy(
   () => import('./save'),

package/src/blueprints/memory/{functions → operations}/query.ts

@@ -25,7 +25,7 @@ export default QueryMemories.pipe(
       query = query.limit(limit);
 
       yield* Database.flush();
-      const results = yield* Database.runQuery(query);
+      const results = yield* Database.query(query).run;
       return results.map((obj) => Entity.toJSON(obj));
     }),
   ),

package/src/blueprints/planning/blueprint.ts

@@ -4,7 +4,7 @@
 
 import { Blueprint } from '@dxos/compute';
 
-import { UpdateTasks } from './functions';
+import { UpdateTasks } from './operations/update-tasks';
 
 const BLUEPRINT_KEY = 'org.dxos.blueprint.planning';
 
@@ -12,7 +12,7 @@ const make = () =>
   Blueprint.make({
     key: BLUEPRINT_KEY,
     name: 'Planning',
-    description: 'Plans and tracks complex tasks with artifact management.',
+    description: 'Plans and tracks complex tasks using artifacts.',
     tools: Blueprint.toolDefinitions({ operations: [UpdateTasks] }),
   });
 

package/src/blueprints/planning/index.ts

@@ -3,4 +3,4 @@
 //
 
 export { default as PlanningBlueprint } from './blueprint';
-export { UpdateTasks as PlanningUpdateTasks, PlanningHandlers } from './functions';
+export { PlanningHandlers, PlanningOperations } from './operations';

package/src/blueprints/planning/{functions → operations}/index.ts

@@ -4,6 +4,6 @@
 
 import { OperationHandlerSet } from '@dxos/compute';
 
-export * from './definitions';
+export * as PlanningOperations from './update-tasks';
 
 export const PlanningHandlers = OperationHandlerSet.lazy(() => import('./update-tasks'));

package/src/blueprints/planning/{functions/definitions.ts → operations/update-tasks.md}

@@ -1,41 +1,38 @@
-//
-// Copyright 2026 DXOS.org
-//
+## PLANNING TOOL
 
-import * as Schema from 'effect/Schema';
+This tool maintains an organized task list.
+Use this to track progress, break down objectives, and ensure thoroughness.
+If you are asked to create a plan use this tool instead of creating a new document.
+After creating initial tasks, update them silently without announcing changes to the user.
+Important: Do not show or summarize the contents of the task list unless the user asks for this.
 
-import { AiContextService } from '@dxos/assistant';
-import { Operation } from '@dxos/compute';
-import { Database } from '@dxos/echo';
-import { trim } from '@dxos/util';
+### CORE USAGE PRINCIPLES
 
-import { Plan } from '../../../types';
+Create and manage tasks for: multi-step objectives requiring 3+ distinct actions, complex projects needing careful sequencing,
+user requests for task organization, multiple deliverables provided together, new instructions (capture as tasks immediately with new IDs),
+completed work (mark `done` and add follow-ups), and active work (mark as `in-progress`, limit one at a time).
 
-const INSTRUCTIONS = trim`
-TASK MANAGEMENT TOOL - USAGE GUIDELINES
+Skip task management for: single straightforward actions, simple requests achievable in 1-2 steps,
+informational queries, quick lookups or clarifications, and avoid creating verification tasks unless requested.
 
-This tool maintains an organized task list during work sessions to track progress, break down objectives, and ensure thoroughness. After creating initial tasks, update them silently without announcing changes to the user.
+### TOOL SPECIFICATION
 
-=== CORE USAGE PRINCIPLES ===
+`update-tasks` requires an array of task objects. Each task object contains:
 
-Create and manage tasks for: multi-step objectives requiring 3+ distinct actions, complex projects needing careful sequencing, user requests for task organization, multiple deliverables provided together, new instructions (capture as tasks immediately with new IDs), completed work (mark done and add follow-ups), and active work (mark as 'in-progress', limit one at a time).
-
-Skip task management for: single straightforward actions, simple requests achievable in 1-2 steps, informational queries, quick lookups or clarifications, and avoid creating verification tasks unless requested.
-
-=== TOOL SPECIFICATION ===
-
-update-tasks requires an array of task objects. Each task object contains:
 - id (string, required): unique identifier like "task_1" or "research_sources"
 - title (string, optional): update to clarify or refine task description
-- status (string, optional): 'todo' | 'in-progress' | 'done'
+- status (string, optional): `todo` | `in-progress` | `done`
 
-Task status meanings: 'todo' means not yet started, 'in-progress' means currently being worked on, 'done' means completed successfully.
+Task status meanings: 'todo' means not yet started, `in-progress` means currently being worked on, `done` means completed successfully.
 
-=== OPERATIONAL GUIDELINES ===
+### OPERATIONAL GUIDELINES
 
-Update tasks in real-time as work progresses. Mark tasks 'done' immediately upon completion. Maintain only ONE task with 'in-progress' status at a time. Complete current tasks before starting new ones. Use specific, actionable task titles. Break complex work into manageable logical pieces. Batch task updates with other actions when possible for efficiency.
+Update tasks in realtime as work progresses. Mark tasks `done` immediately upon completion.
+Maintain only ONE task with `in-progress` status at a time. Complete current tasks before starting new ones.
+Use specific, actionable task titles. Break complex work into manageable logical pieces.
+Batch task updates with other actions when possible for efficiency.
 
-=== USAGE EXAMPLES ===
+### USAGE EXAMPLES
 
 <example type="research_project">
 <user_message>I need to research sustainable packaging options and write a report comparing costs and environmental impact.</user_message>
@@ -110,28 +107,10 @@ Creates tasks:
 <reasoning>One-step action with immediate completion, no organizational benefit from tasks.</reasoning>
 </example>
 
-=== BEST PRACTICES ===
-
-For task creation: use descriptive unique IDs reflecting the work, start first task as 'in-progress', batch initial creation with beginning work. For progress tracking: update status immediately upon completion, keep only one 'in-progress' task unless parallel work is natural, add follow-up tasks as they emerge. For task breakdown: aim for reasonably-scoped tasks, group related small actions into logical units, split tasks requiring different approaches.
-
-When uncertain whether to use task management, err on the side of creating tasks. Proactive organization demonstrates thoroughness and ensures comprehensive work completion.
-`;
-
-const TaskProps = Schema.Struct({
-  id: Plan.TaskId,
-  title: Schema.String,
-  status: Schema.Literal('todo', 'in-progress', 'done'),
-});
-
-export const UpdateTasks = Operation.make({
-  meta: {
-    key: 'org.dxos.function.planning.update-tasks',
-    name: 'Update tasks',
-    description: INSTRUCTIONS,
-  },
-  input: Schema.Struct({
-    tasks: Schema.Array(TaskProps),
-  }),
-  output: Schema.Any,
-  services: [AiContextService, Database.Service],
-});
+### BEST PRACTICES
+
+For task creation: use descriptive unique IDs reflecting the work, start first task as `in-progress`, batch initial creation with beginning work.
+For progress tracking: update status immediately upon completion, keep only one `in-progress` task unless parallel work is natural, add follow-up tasks as they emerge. For task breakdown: aim for reasonably-scoped tasks, group related small actions into logical units, split tasks requiring different approaches.
+
+When uncertain whether to use task management, err on the side of creating tasks.
+Proactive organization demonstrates thoroughness and ensures comprehensive work completion.

package/src/blueprints/planning/operations/update-tasks.test.ts

@@ -0,0 +1,55 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import { describe, it } from '@effect/vitest';
+import * as Effect from 'effect/Effect';
+
+import { Blueprint, Operation } from '@dxos/compute';
+import { Database, Feed, Obj } from '@dxos/echo';
+import { TestHelpers } from '@dxos/effect/testing';
+import { AssistantTestLayer } from '@dxos/functions-runtime/testing';
+import { invariant } from '@dxos/invariant';
+import { EntityId } from '@dxos/keys';
+
+import { Agent, Chat, Plan } from '../../../types';
+import PlanningBlueprint from '../blueprint';
+import { PlanningHandlers } from './index';
+import { UpdateTasks } from './update-tasks';
+
+EntityId.dangerouslyDisableRandomness();
+
+const TestLayer = AssistantTestLayer({
+  operationHandlers: PlanningHandlers,
+  types: [Agent.Agent, Plan.Plan, Chat.Chat, Chat.CompanionTo, Blueprint.Blueprint, Feed.Feed],
+  blueprints: [PlanningBlueprint.make()],
+  disableLlmMemoization: true,
+});
+
+describe('UpdateTasks', () => {
+  it.scoped(
+    'adds tasks to the plan',
+    Effect.fnUntraced(
+      function* ({ expect }) {
+        const agent = yield* Agent.makeInitialized(
+          { name: 'Planner', instructions: 'Test.' },
+          PlanningBlueprint.make(),
+        );
+        yield* Database.flush();
+
+        const chatFeed = agent.chat?.target?.feed?.target;
+        invariant(chatFeed, 'Agent chat feed not found.');
+
+        yield* Operation.invoke(UpdateTasks, {
+          tasks: [{ id: Plan.TaskId.make('task-1'), title: 'Hello', status: 'todo' }],
+        }).pipe(Effect.provide(Operation.withInvocationOptions({ conversation: Obj.getURI(chatFeed) })));
+
+        const plan = yield* Database.load(agent.plan);
+        expect(plan.tasks).toHaveLength(1);
+        expect(plan.tasks[0]).toMatchObject({ id: 'task-1', title: 'Hello', status: 'todo' });
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+  );
+});

package/src/blueprints/planning/operations/update-tasks.ts

@@ -0,0 +1,71 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+import * as Schema from 'effect/Schema';
+
+import { AiContext } from '@dxos/assistant';
+import { Operation } from '@dxos/compute';
+import { Database, Obj } from '@dxos/echo';
+import { DXN } from '@dxos/keys';
+import { trim } from '@dxos/util';
+
+import { Plan, Agent } from '../../../types';
+import INSTRUCTIONS from './update-tasks.md?raw';
+
+// Omit `delegated` (and `chat`) from the LLM-facing schema: delegation is set only by the
+// delegation tool, never by ordinary planning, and keeping it out leaves the tool schema unchanged.
+const SimpleTask = Plan.Task.omit('chat', 'delegated');
+
+export const UpdateTasks = Operation.make({
+  meta: {
+    key: DXN.make('org.dxos.function.planning.updateTasks'),
+    name: 'Update tasks',
+    description: INSTRUCTIONS,
+    icon: 'ph--check-square-offset--regular',
+  },
+  input: Schema.Struct({
+    tasks: Schema.Array(SimpleTask),
+  }),
+  output: Schema.Any,
+  services: [AiContext.Service, Database.Service],
+});
+
+/**
+ * Updates the planning document (Agent.plan) with the given tasks.
+ */
+export default UpdateTasks.pipe(
+  Operation.withHandler(
+    Effect.fn(function* ({ tasks: newTasks }) {
+      const agent = yield* Agent.getFromChatContext;
+      // TODO(burdon): How to specify requirements/preconditions before calling?
+      // TODO(burdon): How to report non-technical error?
+      const plan = yield* Database.load(agent.plan);
+
+      Obj.update(plan, (plan) => {
+        for (const task of newTasks) {
+          const existingTask = plan.tasks.find((t) => t.id === task.id);
+          if (existingTask) {
+            existingTask.title = task.title;
+            existingTask.status = task.status;
+          } else {
+            plan.tasks.push({
+              id: task.id,
+              title: task.title,
+              status: task.status,
+            });
+          }
+        }
+      });
+
+      return trim`
+        You must update the task status to 'in-progress' when you start and 'done' when complete.
+        Current plan updated:
+        <plan>
+          ${Plan.formatPlan(plan)}
+        </plan>
+      `;
+    }) as any,
+  ),
+);

package/src/blueprints/testing.ts

@@ -7,19 +7,25 @@ import * as Effect from 'effect/Effect';
 import { pipe } from 'effect/Function';
 
 import { ConsolePrinter } from '@dxos/ai';
-import { AiContextService, type AiSession, type AiSessionRunProps, GenerationObserver } from '@dxos/assistant';
-import type { Definition } from '@dxos/compute';
+import { AiContext, AiRequest, AiSession, GenerationObserver } from '@dxos/assistant';
+import type { Blueprint } from '@dxos/compute';
 import { Database, Ref } from '@dxos/echo';
 import { log } from '@dxos/log';
 
-export type TestStep = Pick<AiSessionRunProps, 'prompt' | 'system'> & {
+export type TestStep = Pick<AiSession.RunProps, 'prompt' | 'system'> & {
   test?: () => Promise<void>;
 };
 
 /**
  * Runs the prompt steps, calling the test function after each step.
  */
-export const runSteps = Effect.fn(function* (session: AiSession, steps: TestStep[]) {
+export const runSteps: (
+  session: AiSession.Session,
+  steps: TestStep[],
+) => Effect.Effect<void, AiRequest.RunError, AiRequest.RunRequirements> = Effect.fn(function* (
+  session: AiSession.Session,
+  steps: TestStep[],
+) {
   for (const { test, ...props } of steps) {
     yield* session.createRequest({
       ...props,
@@ -37,8 +43,8 @@ export const runSteps = Effect.fn(function* (session: AiSession, steps: TestStep
  * Binds blueprints from the blueprint definitions.
  */
 // TODO(dmaretskyi): Potentially the agent will auto-bind the blueprints.
-export const addBlueprints = Effect.fnUntraced(function* (blueprints: Definition[]) {
-  yield* AiContextService.bindContext({
+export const addBlueprints = Effect.fnUntraced(function* (blueprints: Blueprint.Definition[]) {
+  yield* AiContext.Service.bindContext({
     blueprints: yield* pipe(
       blueprints,
       Arr.map((blueprint) => blueprint.make()),