@orkestrel/agent 0.0.4 → 0.0.6

package/dist/src/core/index.cjs

@@ -2272,10 +2272,13 @@ var Agent = class {
 		return this.stream(options).result;
 	}
 	stream(options) {
-		const timeout = this.#timeoutMs === void 0 ? void 0 : (0, _orkestrel_timeout.createTimeout)({ ms: this.#timeoutMs });
+		const timeoutMs = options?.timeout ?? this.#timeoutMs;
+		const timeout = timeoutMs === void 0 ? void 0 : (0, _orkestrel_timeout.createTimeout)({ ms: timeoutMs });
 		timeout?.start();
-		this.#budget?.start();
-		const abort = (0, _orkestrel_abort.createAbort)({ signal: this.#parents(timeout) });
+		const budget = options?.budget ?? this.#budget;
+		budget?.start();
+		const limit = options?.limit ?? this.#limit;
+		const abort = (0, _orkestrel_abort.createAbort)({ signal: this.#parents(timeout, budget, options?.signal) });
 		this.#runs.add(abort);
 		this.#status = "running";
 		this.#emitter.emit("start", this.id);
@@ -2283,11 +2286,12 @@ var Agent = class {
 			content: "",
 			thinking: void 0,
 			usage: void 0,
-			partial: false
+			partial: false,
+			exhausted: false
 		};
 		const channel = new Channel();
 		const settled = (0, _orkestrel_workflow.createDeferred)();
-		this.#pump(abort, outcome, timeout, channel, settled, options?.think);
+		this.#pump(abort, outcome, timeout, channel, settled, options?.think, options?.schema, limit, budget);
 		settled.promise.catch(() => {});
 		return {
 			events: this.#events(channel, abort),
@@ -2300,10 +2304,10 @@ var Agent = class {
 	abort(reason) {
 		for (const abort of [...this.#runs]) abort.abort(reason);
 	}
-	async #pump(abort, outcome, timeout, channel, settled, think) {
+	async #pump(abort, outcome, timeout, channel, settled, think, schema, limit, budget) {
 		let failure;
 		try {
-			for await (const chunk of this.#run(abort, outcome, think)) channel.push(chunk);
+			for await (const chunk of this.#run(abort, outcome, think, schema, limit, budget)) channel.push(chunk);
 		} catch (error) {
 			failure = { error };
 		} finally {
@@ -2314,7 +2318,8 @@ var Agent = class {
 				channel.close();
 				const result = this.#result(outcome);
 				settled.resolve(result);
-				if (outcome.partial) this.#emitter.emit("abort", abort.signal.reason);
+				if (outcome.exhausted) this.#emitter.emit("exhaust", limit);
+				else if (outcome.partial) this.#emitter.emit("abort", abort.signal.reason);
 				this.#emitter.emit("finish", result);
 			} else {
 				this.#status = "error";
@@ -2331,51 +2336,83 @@ var Agent = class {
 			abort.abort();
 		}
 	}
-	async *#run(abort, outcome, think) {
+	async *#run(abort, outcome, think, schema, limit, budget) {
 		const messages = [...this.#context.build(this.#provider.format)];
 		const tools = this.#context.tools;
 		let content = "";
 		let thinking;
 		let usage;
+		let pending = false;
+		let broke = false;
 		const compaction = { futile: false };
 		const compacting = this.#window !== void 0 && this.#context.conversations.active?.summarizable === true;
 		if (compacting) {
 			this.#window?.clear();
 			if (!abort.signal.aborted) await this.#trim(messages, compaction, false);
 		}
-		for (let turn = 0; turn < this.#limit; turn += 1) {
+		for (let turn = 0; turn < limit; turn += 1) {
 			this.#emitter.emit("turn", turn);
 			if (turn > 0) try {
 				await this.#scheduler?.yield({ signal: abort.signal });
 			} catch (error) {
 				if (abort.signal.aborted) {
 					outcome.partial = true;
+					broke = true;
 					break;
 				}
 				throw error;
 			}
 			if (abort.signal.aborted) {
 				outcome.partial = true;
+				broke = true;
 				break;
 			}
 			const advertised = filterAllowList(this.#context.scope?.tools, tools.definitions(), (definition) => definition.name);
 			const definitions = advertised.length > 0 ? advertised : void 0;
+			let charged = 0;
+			let turnContent = "";
 			let result;
 			try {
-				result = yield* this.#provide(messages, abort.signal, definitions, think, (delta) => {
+				result = yield* this.#provide(messages, abort.signal, definitions, think, schema, (delta) => {
 					content += delta;
+					turnContent += delta;
+					const est = estimateTokens(turnContent);
+					if (est > charged) {
+						budget?.consume({
+							prompt: 0,
+							completion: est - charged,
+							total: est - charged
+						});
+						charged = est;
+					}
 				});
 			} catch (error) {
 				if (abort.signal.aborted) {
-					if (isProviderAbortError(error) && error.partial.thinking !== void 0) thinking = this.#thought(thinking, error.partial.thinking);
+					if (isProviderAbortError(error)) {
+						if (error.partial.thinking !== void 0) thinking = this.#thought(thinking, error.partial.thinking);
+						if (error.partial.usage !== void 0) {
+							const abortUsage = error.partial.usage;
+							budget?.consume({
+								prompt: abortUsage.prompt,
+								completion: Math.max(0, abortUsage.completion - charged),
+								total: Math.max(0, abortUsage.total - charged)
+							});
+							usage = this.#sum(usage, abortUsage);
+						}
+					}
 					outcome.partial = true;
+					broke = true;
 					break;
 				}
 				throw error;
 			}
 			if (result.thinking !== void 0 && result.thinking.length > 0) thinking = this.#thought(thinking, result.thinking);
 			if (result.usage !== void 0) {
-				this.#budget?.consume(result.usage);
+				budget?.consume({
+					prompt: result.usage.prompt,
+					completion: Math.max(0, result.usage.completion - charged),
+					total: Math.max(0, result.usage.total - charged)
+				});
 				usage = this.#sum(usage, result.usage);
 				this.#emitter.emit("usage", result.usage);
 				yield {
@@ -2408,6 +2445,7 @@ var Agent = class {
 					messages.push(toolMessage);
 				}
 				if (compacting) await this.#trim(messages, compaction, true);
+				pending = true;
 				continue;
 			}
 			messages.push(this.#context.messages.add({
@@ -2415,8 +2453,14 @@ var Agent = class {
 				content: result.content
 			}));
 			content = result.content;
+			pending = false;
+			broke = true;
 			break;
 		}
+		if (!broke && pending) {
+			outcome.partial = true;
+			outcome.exhausted = !abort.signal.aborted;
+		}
 		outcome.content = content;
 		outcome.thinking = thinking;
 		outcome.usage = usage;
@@ -2476,8 +2520,11 @@ var Agent = class {
 	#reason(error) {
 		return error instanceof Error ? error.message : String(error);
 	}
-	async *#provide(messages, signal, definitions, think, onDelta) {
-		const generator = this.#provider.stream(messages, signal, definitions, think === void 0 ? void 0 : { think });
+	async *#provide(messages, signal, definitions, think, schema, onDelta) {
+		const options = {};
+		if (think !== void 0) options.think = think;
+		if (schema !== void 0) options.schema = schema;
+		const generator = this.#provider.stream(messages, signal, definitions, Object.keys(options).length > 0 ? options : void 0);
 		let next = await generator.next();
 		while (!next.done) {
 			const delta = next.value;
@@ -2495,11 +2542,12 @@ var Agent = class {
 		}
 		return next.value;
 	}
-	#parents(timeout) {
+	#parents(timeout, budget, signal) {
 		const signals = [];
 		if (this.#signal !== void 0) signals.push(this.#signal);
+		if (signal !== void 0) signals.push(signal);
 		if (timeout !== void 0) signals.push(timeout.signal);
-		if (this.#budget !== void 0) signals.push(this.#budget.signal);
+		if (budget !== void 0) signals.push(budget.signal);
 		if (signals.length === 0) return void 0;
 		if (signals.length === 1) return signals[0];
 		return AbortSignal.any(signals);