@ddwang/magnitude-core 0.3.1-ddwang.1 → 0.3.1-ddwang.2

package/dist/agent/index.js

@@ -2,9 +2,12 @@ import logger from '@/logger';
 import EventEmitter from "eventemitter3";
 import z from "zod";
 import { Observation } from '@/memory/observation';
-import { AgentError } from "@/agent/errors";
+import { ActionLimitError, AgentBusyError, AgentError } from "@/agent/errors";
+import { Operation, checkOperation, currentOperation, operationOptions, untilAborted } from '@/common/operation';
 import { AgentMemory } from "@/memory";
 import { taskActions } from "@/actions/taskActions";
+import { memoryActions } from '@/actions/memoryActions';
+import { NOTEBOOK_INSTRUCTIONS } from '@/memory/notebook';
 import { traceAsync } from "@/ai/baml_client";
 import { telemetrifyAgent } from '@/telemetry/events';
 import { isClaude } from '@/ai/util';
@@ -24,6 +27,7 @@ const DEFAULT_CONFIG = {
     },
     prompt: null,
     telemetry: true,
+    maxActions: Infinity,
 };
 export class Agent {
     // maybe remove conns/actions from options since stored sep
@@ -41,6 +45,9 @@ export class Agent {
     doneActing;
     _paused = false;
     _pauseResolve = null;
+    activeOperation;
+    idle = Promise.resolve();
+    stopped = false;
     latestTaskMemory; // | null = null;
     constructor(baseConfig = {}) {
         this.options = {
@@ -49,6 +56,9 @@ export class Agent {
             connectors: baseConfig.connectors ?? [],
             actions: [...(baseConfig.actions || DEFAULT_CONFIG.actions)],
         };
+        if (baseConfig.maxActions !== undefined && (!Number.isSafeInteger(baseConfig.maxActions) || baseConfig.maxActions < 1)) {
+            throw new Error('maxActions must be a positive integer');
+        }
         this.connectors = this.options.connectors;
         // Aggregate actions from connectors
         //const aggregatedActions = [...this.options.actions];
@@ -56,6 +66,11 @@ export class Agent {
         for (const connector of this.connectors) {
             this.actions.push(...(connector.getActionSpace ? connector.getActionSpace() : []));
         }
+        for (const action of memoryActions) {
+            if (this.actions.some(existing => existing.name === action.name))
+                throw new Error(`Action name '${action.name}' is reserved for task memory.`);
+            this.actions.push(action);
+        }
         // Deduplicate actions by name
         // TODO: maybe error instead, or automatically differentiate them?
         //this.options.actions = Array.from(new Map(aggregatedActions.map(actDef => [actDef.name, actDef])).values());
@@ -99,6 +114,9 @@ export class Agent {
         return connector;
     }
     async start() {
+        checkOperation();
+        if (this.busy)
+            throw new AgentBusyError();
         // Register telemetry if enabled - do on start instead of cons to prevent weird subclass event issues
         if (this.options.telemetry)
             telemetrifyAgent(this);
@@ -111,6 +129,7 @@ export class Agent {
                 await connector.onStart();
         }
         this.events.emit('start');
+        this.stopped = false;
         logger.info("Agent: All connectors started.");
         // logger.info("Making initial observations...");
         // await this._recordConnectorObservations();
@@ -127,7 +146,46 @@ export class Agent {
         }
         return actionDefinition;
     }
-    async exec(action, memory) {
+    /** True until underlying work settles, including after a cancelled caller returns. */
+    get busy() {
+        return this.activeOperation !== undefined;
+    }
+    whenIdle() {
+        return this.idle;
+    }
+    async runOperation(options, fn) {
+        const inherited = currentOperation();
+        if (inherited?.owner === this)
+            inherited.check();
+        if (this.busy)
+            throw new AgentBusyError();
+        if (this.stopped)
+            throw new AgentError('Agent is stopped; call start() before using it');
+        const operation = new Operation(this, options);
+        this.activeOperation = operation;
+        let markIdle;
+        this.idle = new Promise(resolve => { markIdle = resolve; });
+        const worker = operation.run(async () => {
+            operation.check();
+            try {
+                return await fn();
+            }
+            finally {
+                operation.check(); // Preserve the cancellation/deadline cause through downstream errors.
+            }
+        });
+        const settled = worker.finally(() => {
+            operation.dispose();
+            this.activeOperation = undefined;
+            markIdle();
+        });
+        return untilAborted(settled, operation.signal);
+    }
+    async exec(action, memory, options = {}) {
+        return this.runOperation(options, () => this._exec(action, memory));
+    }
+    async _exec(action, memory) {
+        checkOperation();
         /**
          * Execute an action that belongs to this Agent's action space.
          * Provide memory to record the action taken, its results, and any connector observations to that memory.
@@ -145,35 +203,54 @@ export class Agent {
         if (!parsed.success) {
             throw new AgentError(`Generated action '${action.variant}' violates input schema: ${parsed.error.message}`, { adaptable: true });
         }
+        const memoryOnly = memoryActions.includes(actionDefinition);
+        if (!memoryOnly)
+            for (const connector of this.connectors) {
+                await connector.beforeAction?.(action, operationOptions());
+                checkOperation();
+            }
         this.events.emit('actionStarted', action);
-        const data = await actionDefinition.resolver({ input: parsed.data, agent: this });
+        checkOperation();
+        const data = await actionDefinition.resolver({ input: parsed.data, agent: this, memory, ...operationOptions() });
+        checkOperation();
         this.events.emit('actionDone', action);
+        checkOperation();
         if (memory) {
             // Record action taken
-            memory.recordObservation(Observation.fromActionTaken(actionDefinition.name, JSON.stringify(action)));
+            memory.recordObservation(Observation.fromActionTaken(actionDefinition.name, JSON.stringify(action), memoryOnly ? { type: 'notebook-write' } : undefined));
             // Record results of action
             if (data) {
-                memory.recordObservation(Observation.fromActionResult(actionDefinition.name, data));
+                memory.recordObservation(Observation.fromActionResult(actionDefinition.name, data, memoryOnly ? { type: 'notebook-result', limit: 1 } : undefined));
             }
             // Collect and record observations from connectors
-            await this._recordConnectorObservations(memory);
+            if (memoryOnly)
+                this.events.emit('observationsRecorded'); // Checkpoint notes without another browser capture.
+            else
+                await this._recordConnectorObservations(memory);
         }
+        return data;
     }
     async _recordConnectorObservations(memory) {
+        checkOperation();
         for (const connector of this.connectors) {
             // could do Promise.all if matters
-            const connObservations = connector.collectObservations ? await connector.collectObservations() : [];
+            const connObservations = connector.collectObservations ? await connector.collectObservations(operationOptions()) : [];
+            checkOperation();
             //observations.push(...connObservations);
             for (const obs of connObservations) {
                 memory.recordObservation(obs);
             }
         }
+        this.events.emit('observationsRecorded');
     }
     get memory() {
         //if (!this.latestTaskMemory) throw new Error("No memory available");
         return this.latestTaskMemory;
     }
     async act(taskOrSteps, options = {}) {
+        return this.runOperation(options, () => this._runAct(taskOrSteps, options));
+    }
+    async _runAct(taskOrSteps, options) {
         const instructions = [
             ...(this.options.prompt ? [this.options.prompt] : []),
             ...(options.prompt ? [options.prompt] : []),
@@ -185,8 +262,10 @@ export class Agent {
             // trace overall task
             await (traceAsync('multistep', async (steps, options) => {
                 for (const step of steps) {
+                    checkOperation();
                     this.events.emit('actStarted', step, options);
                     await this._traceAct(step, taskMemory, options);
+                    checkOperation();
                     this.events.emit('actDone', step, options);
                 }
             })(steps, options));
@@ -196,6 +275,7 @@ export class Agent {
             const task = taskOrSteps;
             this.events.emit('actStarted', task, options);
             await this._traceAct(task, taskMemory, options);
+            checkOperation();
             this.events.emit('actDone', task, options);
         }
     }
@@ -205,12 +285,15 @@ export class Agent {
             await this._act(task, memory, options);
         })(task));
     }
-    async _buildContext(memory) {
-        const messages = await memory.render();
+    async _buildContext(memory, options) {
+        checkOperation();
+        const messages = await memory.render(options);
+        checkOperation();
         const connectorInstructions = [];
         for (const connector of this.connectors) {
             if (connector.getInstructions) {
-                const instructions = await connector.getInstructions();
+                const instructions = await connector.getInstructions(operationOptions());
+                checkOperation();
                 if (instructions) {
                     connectorInstructions.push({
                         connectorId: connector.id,
@@ -227,7 +310,12 @@ export class Agent {
         };
     }
     async _act(description, memory, options = {}) {
+        checkOperation();
         this.doneActing = false;
+        for (const connector of this.connectors) {
+            connector.onTaskStart?.(operationOptions());
+            checkOperation();
+        }
         logger.info(`Act: ${description}`);
         // for now simply add data to task
         let dataContentParts = [];
@@ -247,20 +335,29 @@ export class Agent {
         //const testData = convertOptionsToTestData(options);
         // Initialize task memory and record initial observations
         // Combine any agent-level and task-level instructions
+        checkOperation();
         this.latestTaskMemory = memory;
         // record initial observations
         logger.info("Making initial observations...");
         await this._recordConnectorObservations(memory);
         logger.info("Initial observations recorded");
+        let actionCount = 0;
         while (true) {
+            checkOperation();
+            if (actionCount >= this.options.maxActions)
+                throw new ActionLimitError(this.options.maxActions);
             // Removed direct screenshot/tabState access here; it's part of memoryContext via connectors
             logger.info(`Creating partial recipe`);
             let reasoning = "";
             let actions = [];
+            let memoryUpdates = [];
             try {
+                this.events.emit('planningStarted');
                 const memoryContext = await this._buildContext(memory);
+                memoryContext.connectorInstructions.unshift({ connectorId: 'task_memory', instructions: NOTEBOOK_INSTRUCTIONS });
                 await retryOnError(async () => {
-                    ({ reasoning, actions } = await this.models.partialAct(memoryContext, description, dataContentParts, this.actions));
+                    ({ reasoning, actions, memory_updates: memoryUpdates } = await this.models.partialAct(memoryContext, description, dataContentParts, this.actions));
+                    checkOperation();
                     if (actions.length === 0) {
                         // Empty action list behavior - default wait else ... err? what if not in action space?
                         //actions.push()
@@ -278,6 +375,7 @@ export class Agent {
                 });
             }
             catch (error) {
+                checkOperation();
                 logger.error(`Error planning actions: ${error instanceof Error ? error.message : String(error)}`);
                 /**
                  * (1) Failure to conform to JSON
@@ -293,13 +391,24 @@ export class Agent {
             logger.info({ reasoning, actions }, `Partial recipe created`);
             // Could be emitted in memory and bubbled up instead of recordThought was called in more places
             this.events.emit('thought', reasoning);
+            checkOperation();
             memory.recordThought(reasoning);
-            // Execute partial recipe
-            for (const action of actions) {
+            // Persist the review using the existing audited, budgeted note action.
+            // Empty reviews are free; each attempted write consumes one action.
+            const batch = [...memoryUpdates.map(note => ({ variant: 'memory:note', ...note })), ...actions];
+            for (const action of batch) {
                 await this._waitIfPaused();
+                checkOperation();
                 if (this.doneActing)
                     break;
-                await this.exec(action, memory);
+                if (actionCount >= this.options.maxActions)
+                    throw new ActionLimitError(this.options.maxActions);
+                const result = await this._exec(action, memory);
+                actionCount++;
+                // Preserve the current page when an update fails. Successful
+                // earlier writes remain; the next plan sees the failure result.
+                if (action.variant === 'memory:note' && result?.saved === false)
+                    break;
                 // const postActionScreenshot = await this.screenshot();
                 // const actionDescriptor: ActionDescriptor = { ...action, screenshot: postActionScreenshot.image } as ActionDescriptor;
                 // this.events.emit('action', actionDescriptor);
@@ -318,28 +427,42 @@ export class Agent {
         //this.events.emit('stepSuccess');
         //this.currentTaskMemory = null;
     }
-    async query(query, schema) {
-        // Record observations in case no act() was used beforehand
-        await this._recordConnectorObservations(this.latestTaskMemory);
-        const memoryContext = await this._buildContext(this.memory); //this.memory.buildContext(this.connectors);
-        return await this.models.query(memoryContext, query, schema);
+    async query(query, schema, options = {}) {
+        return this.runOperation(options, async () => {
+            // Record observations in case no act() was used beforehand
+            await this._recordConnectorObservations(this.latestTaskMemory);
+            const memoryContext = await this._buildContext(this.memory, options);
+            return await this.models.query(memoryContext, query, schema);
+        });
     }
     async queueDone() {
+        checkOperation();
         this.doneActing = true;
     }
     async _waitIfPaused() {
+        checkOperation();
         if (!this._paused)
             return;
         this.events.emit('pause');
+        checkOperation();
+        if (!this._paused)
+            return; // A pause listener may have resumed synchronously.
         logger.info("Agent: Paused");
-        await new Promise((resolve) => {
-            this._pauseResolve = resolve;
-        });
+        try {
+            await untilAborted(new Promise((resolve) => {
+                this._pauseResolve = resolve;
+            }), currentOperation()?.signal);
+        }
+        finally {
+            this._pauseResolve = null;
+        }
     }
     pause() {
+        checkOperation();
         this._paused = true;
     }
     resume() {
+        checkOperation();
         this._paused = false;
         if (this._pauseResolve) {
             this._pauseResolve();
@@ -352,13 +475,18 @@ export class Agent {
         return this._paused;
     }
     async stop() {
+        checkOperation();
         /**
          * Stop the agent and close the browser context.
          * May be called asynchronously and interrupt an agent in the middle of a action sequence.
          */
+        this.stopped = true;
+        this.activeOperation?.cancel('Agent stopped');
         this.doneActing = true;
         if (this._paused) {
-            this.resume(); // unblock so loop can see doneActing and exit
+            this._paused = false;
+            this._pauseResolve?.();
+            this._pauseResolve = null;
         }
         logger.info("Agent: Stopping connectors...");
         for (const connector of this.connectors) {

package/dist/ai/baml_client/async_client.d.ts

@@ -7,8 +7,8 @@ $ yarn add @boundaryml/baml
 $ pnpm add @boundaryml/baml
 
 *************************************************************************************************/
-import type { BamlRuntime, BamlCtxManager, ClientRegistry, Image, Collector } from "@boundaryml/baml";
-import { BamlStream } from "@boundaryml/baml";
+import type { BamlRuntime, BamlCtxManager, Image, FunctionLog } from "@boundaryml/baml";
+import { BamlStream, Collector, ClientRegistry } from "@boundaryml/baml";
 import type { RecursivePartialNull as MovedRecursivePartialNull } from "./types";
 import type { partial_types } from "./partial_types";
 import type * as types from "./types";
@@ -16,14 +16,20 @@ import type TypeBuilder from "./type_builder";
 import { AsyncHttpRequest, AsyncHttpStreamRequest } from "./async_request";
 import { LlmResponseParser, LlmStreamParser } from "./parser";
 /**
- * @deprecated Use RecursivePartialNull from 'baml_client/types' instead.
- */
+* @deprecated Use RecursivePartialNull from 'baml_client/types' instead.
+*/
 export type RecursivePartialNull<T> = MovedRecursivePartialNull<T>;
-type BamlCallOptions = {
+type TickReason = "Unknown";
+type BamlCallOptions<WatchersT = never> = {
     tb?: TypeBuilder;
     clientRegistry?: ClientRegistry;
+    client?: string;
     collector?: Collector | Collector[];
     env?: Record<string, string | undefined>;
+    tags?: Record<string, string>;
+    signal?: AbortSignal;
+    onTick?: (reason: TickReason, log: FunctionLog | null) => void;
+    watchers?: WatchersT;
 };
 export declare class BamlAsyncClient {
     private runtime;
@@ -41,33 +47,18 @@ export declare class BamlAsyncClient {
     get streamRequest(): AsyncHttpStreamRequest;
     get parse(): LlmResponseParser;
     get parseStream(): LlmStreamParser;
-    CreatePartialRecipe(memory: types.AgentContext, task: string, taskData: (Image | string)[], includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions): Promise<types.PartialRecipe>;
-    ExtractData(instructions: string, screenshot: Image, domContent: string, includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions): Promise<types.ExtractedData>;
-    QueryMemory(memory: types.AgentContext, query: string, includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions): Promise<types.QueryResponse>;
+    CreatePartialRecipe(memory: types.AgentContext, task: string, taskData: (Image | string)[], includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions<never>): Promise<types.PartialRecipe>;
+    ExtractData(instructions: string, screenshot: Image, domContent: string, includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions<never>): Promise<types.ExtractedData>;
+    QueryMemory(memory: types.AgentContext, query: string, includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions<never>): Promise<types.QueryResponse>;
 }
 declare class BamlStreamClient {
     private runtime;
     private ctxManager;
     private bamlOptions;
     constructor(runtime: BamlRuntime, ctxManager: BamlCtxManager, bamlOptions?: BamlCallOptions);
-    CreatePartialRecipe(memory: types.AgentContext, task: string, taskData: (Image | string)[], includeClaudeSpoof: boolean, __baml_options__?: {
-        tb?: TypeBuilder;
-        clientRegistry?: ClientRegistry;
-        collector?: Collector | Collector[];
-        env?: Record<string, string | undefined>;
-    }): BamlStream<partial_types.PartialRecipe, types.PartialRecipe>;
-    ExtractData(instructions: string, screenshot: Image, domContent: string, includeClaudeSpoof: boolean, __baml_options__?: {
-        tb?: TypeBuilder;
-        clientRegistry?: ClientRegistry;
-        collector?: Collector | Collector[];
-        env?: Record<string, string | undefined>;
-    }): BamlStream<partial_types.ExtractedData, types.ExtractedData>;
-    QueryMemory(memory: types.AgentContext, query: string, includeClaudeSpoof: boolean, __baml_options__?: {
-        tb?: TypeBuilder;
-        clientRegistry?: ClientRegistry;
-        collector?: Collector | Collector[];
-        env?: Record<string, string | undefined>;
-    }): BamlStream<partial_types.QueryResponse, types.QueryResponse>;
+    CreatePartialRecipe(memory: types.AgentContext, task: string, taskData: (Image | string)[], includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions<never>): BamlStream<partial_types.PartialRecipe, types.PartialRecipe>;
+    ExtractData(instructions: string, screenshot: Image, domContent: string, includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions<never>): BamlStream<partial_types.ExtractedData, types.ExtractedData>;
+    QueryMemory(memory: types.AgentContext, query: string, includeClaudeSpoof: boolean, __baml_options__?: BamlCallOptions<never>): BamlStream<partial_types.QueryResponse, types.QueryResponse>;
 }
 export declare const b: BamlAsyncClient;
 export {};