@tanstack/ai
Version:
Core TanStack AI library - Open source AI SDK
328 lines (327 loc) • 9.89 kB
JavaScript
import { aiEventClient } from "@tanstack/ai-event-client";
function shouldSkipInstrumentation(mw) {
return mw.name === "devtools";
}
function instrumentCtx(ctx) {
return {
requestId: ctx.requestId,
streamId: ctx.streamId,
clientId: ctx.conversationId,
timestamp: Date.now()
};
}
class MiddlewareRunner {
constructor(middlewares) {
this.middlewares = middlewares;
}
get hasMiddleware() {
return this.middlewares.length > 0;
}
/**
* Pipe config through all middleware onConfig hooks in order.
* Each middleware receives the merged config from previous middleware.
* Partial returns are shallow-merged with the current config.
*/
async runOnConfig(ctx, config) {
let current = config;
for (const mw of this.middlewares) {
if (mw.onConfig) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
const result = await mw.onConfig(ctx, current);
const hasTransform = result !== void 0 && result !== null;
if (hasTransform) {
current = { ...current, ...result };
}
if (!skip) {
const base = instrumentCtx(ctx);
aiEventClient.emit("middleware:hook:executed", {
...base,
middlewareName: mw.name || "unnamed",
hookName: "onConfig",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform
});
if (hasTransform) {
aiEventClient.emit("middleware:config:transformed", {
...base,
middlewareName: mw.name || "unnamed",
iteration: ctx.iteration,
changes: result
});
}
}
}
}
return current;
}
/**
* Call onStart on all middleware in order.
*/
async runOnStart(ctx) {
for (const mw of this.middlewares) {
if (mw.onStart) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onStart(ctx);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onStart",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Pipe a single chunk through all middleware onChunk hooks in order.
* Returns the resulting chunks (0..N) to yield to the consumer.
*
* - void: pass through unchanged
* - chunk: replace with this chunk
* - chunk[]: expand to multiple chunks
* - null: drop the chunk entirely
*/
async runOnChunk(ctx, chunk) {
let chunks = [chunk];
for (const mw of this.middlewares) {
if (!mw.onChunk) continue;
const skip = shouldSkipInstrumentation(mw);
const nextChunks = [];
for (const c of chunks) {
const result = await mw.onChunk(ctx, c);
if (result === null) {
if (!skip) {
aiEventClient.emit("middleware:chunk:transformed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
originalChunkType: c.type,
resultCount: 0,
wasDropped: true
});
}
continue;
} else if (result === void 0) {
nextChunks.push(c);
} else if (Array.isArray(result)) {
nextChunks.push(...result);
if (!skip) {
aiEventClient.emit("middleware:chunk:transformed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
originalChunkType: c.type,
resultCount: result.length,
wasDropped: false
});
}
} else {
nextChunks.push(result);
if (!skip) {
aiEventClient.emit("middleware:chunk:transformed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
originalChunkType: c.type,
resultCount: 1,
wasDropped: false
});
}
}
}
chunks = nextChunks;
}
return chunks;
}
/**
* Run onBeforeToolCall through middleware in order.
* Returns the first non-void decision, or undefined to continue normally.
*/
async runOnBeforeToolCall(ctx, hookCtx) {
for (const mw of this.middlewares) {
if (mw.onBeforeToolCall) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
const decision = await mw.onBeforeToolCall(ctx, hookCtx);
const hasTransform = decision !== void 0 && decision !== null;
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onBeforeToolCall",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform
});
}
if (hasTransform) {
return decision;
}
}
}
return void 0;
}
/**
* Run onAfterToolCall on all middleware in order.
*/
async runOnAfterToolCall(ctx, info) {
for (const mw of this.middlewares) {
if (mw.onAfterToolCall) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onAfterToolCall(ctx, info);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onAfterToolCall",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Run onUsage on all middleware in order.
*/
async runOnUsage(ctx, usage) {
for (const mw of this.middlewares) {
if (mw.onUsage) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onUsage(ctx, usage);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onUsage",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Run onFinish on all middleware in order.
*/
async runOnFinish(ctx, info) {
for (const mw of this.middlewares) {
if (mw.onFinish) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onFinish(ctx, info);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onFinish",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Run onAbort on all middleware in order.
*/
async runOnAbort(ctx, info) {
for (const mw of this.middlewares) {
if (mw.onAbort) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onAbort(ctx, info);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onAbort",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Run onError on all middleware in order.
*/
async runOnError(ctx, info) {
for (const mw of this.middlewares) {
if (mw.onError) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onError(ctx, info);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onError",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Run onIteration on all middleware in order.
* Called at the start of each agent loop iteration.
*/
async runOnIteration(ctx, info) {
for (const mw of this.middlewares) {
if (mw.onIteration) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onIteration(ctx, info);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onIteration",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
/**
* Run onToolPhaseComplete on all middleware in order.
* Called after all tool calls in an iteration have been processed.
*/
async runOnToolPhaseComplete(ctx, info) {
for (const mw of this.middlewares) {
if (mw.onToolPhaseComplete) {
const skip = shouldSkipInstrumentation(mw);
const start = Date.now();
await mw.onToolPhaseComplete(ctx, info);
if (!skip) {
aiEventClient.emit("middleware:hook:executed", {
...instrumentCtx(ctx),
middlewareName: mw.name || "unnamed",
hookName: "onToolPhaseComplete",
iteration: ctx.iteration,
duration: Date.now() - start,
hasTransform: false
});
}
}
}
}
}
export {
MiddlewareRunner
};
//# sourceMappingURL=compose.js.map