ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): restore streamFn extension compatibility

Browse Source 
Keep streamFn required for typed callers while preserving the legacy runtime fallback for extensions that omit it.\n\nfixes #6915
This commit is contained in:
Mario Zechner
2026-07-21 18:34:56 +02:00
parent b142504125
commit b9e5c5d941
26 changed files with 210 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/agent/test/agent.test.ts
+48 -20
View File
@@ -1,7 +1,14 @@
import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@earendil-works/pi-ai/compat";
import { Type } from "typebox";
import { describe, expect, it } from "vitest";
import { Agent, type AgentEvent, type AgentTool, type AgentToolUpdateCallback, type StreamFn } from "../src/index.ts";
import {
Agent,
type AgentEvent,
type AgentTool,
type AgentToolUpdateCallback,
type StreamFn,
setDefaultStreamFn,
} from "../src/index.ts";
// Mock stream that mimics AssistantMessageEventStream
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
@@ -75,8 +82,29 @@ function createDeferred(): {
}
describe("Agent", () => {
it("uses the configured default when a legacy caller omits streamFn", async () => {
let calls = 0;
setDefaultStreamFn(() => {
calls++;
const stream = new MockAssistantStream();
queueMicrotask(() => {
const message = createAssistantMessage("fallback");
stream.push({ type: "done", reason: "stop", message });
});
return stream;
});
try {
const agent = Reflect.construct(Agent, [{}]) as Agent;
await agent.prompt("Hello");
expect(calls).toBe(1);
} finally {
setDefaultStreamFn(undefined);
}
});
it("should create an agent instance with default state", () => {
const agent = new Agent({ streamFunction: unusedStreamFunction });
const agent = new Agent({ streamFn: unusedStreamFunction });
expect(agent.state).toBeDefined();
expect(agent.state.systemPrompt).toBe("");
@@ -93,7 +121,7 @@ describe("Agent", () => {
it("should create an agent instance with custom initial state", () => {
const customModel = getModel("openai", "gpt-4o-mini");
const agent = new Agent({
streamFunction: unusedStreamFunction,
streamFn: unusedStreamFunction,
initialState: {
systemPrompt: "You are a helpful assistant.",
model: customModel,
@@ -107,7 +135,7 @@ describe("Agent", () => {
});
it("should subscribe to events", () => {
const agent = new Agent({ streamFunction: unusedStreamFunction });
const agent = new Agent({ streamFn: unusedStreamFunction });
let eventCount = 0;
const unsubscribe = agent.subscribe((_event) => {
@@ -130,7 +158,7 @@ describe("Agent", () => {
it("emits full lifecycle events for thrown run failures", async () => {
const agent = new Agent({
streamFunction: () => {
streamFn: () => {
throw new Error("provider exploded");
},
});
@@ -162,7 +190,7 @@ describe("Agent", () => {
it("should await async subscribers before prompt resolves", async () => {
const barrier = createDeferred();
const agent = new Agent({
streamFunction: () => {
streamFn: () => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
@@ -200,7 +228,7 @@ describe("Agent", () => {
it("waitForIdle should wait for async subscribers", async () => {
const barrier = createDeferred();
const agent = new Agent({
streamFunction: () => {
streamFn: () => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
@@ -235,7 +263,7 @@ describe("Agent", () => {
it("should pass the active abort signal to subscribers", async () => {
let receivedSignal: AbortSignal | undefined;
const agent = new Agent({
streamFunction: (_model, _context, options) => {
streamFn: (_model, _context, options) => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({ type: "start", partial: createAssistantMessage("") });
@@ -298,7 +326,7 @@ describe("Agent", () => {
};
const agent = new Agent({
initialState: { tools: [tool] },
streamFunction: () => {
streamFn: () => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({
@@ -373,7 +401,7 @@ describe("Agent", () => {
};
const agent = new Agent({
initialState: { tools: [settledTool, slowTool] },
streamFunction: () => {
streamFn: () => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({
@@ -412,7 +440,7 @@ describe("Agent", () => {
});
it("should update state with mutators", () => {
const agent = new Agent({ streamFunction: unusedStreamFunction });
const agent = new Agent({ streamFn: unusedStreamFunction });
// Test setSystemPrompt
agent.state.systemPrompt = "Custom prompt";
@@ -451,7 +479,7 @@ describe("Agent", () => {
});
it("should support steering message queue", async () => {
const agent = new Agent({ streamFunction: unusedStreamFunction });
const agent = new Agent({ streamFn: unusedStreamFunction });
const message = { role: "user" as const, content: "Steering message", timestamp: Date.now() };
agent.steer(message);
@@ -461,7 +489,7 @@ describe("Agent", () => {
});
it("should support follow-up message queue", async () => {
const agent = new Agent({ streamFunction: unusedStreamFunction });
const agent = new Agent({ streamFn: unusedStreamFunction });
const message = { role: "user" as const, content: "Follow-up message", timestamp: Date.now() };
agent.followUp(message);
@@ -471,7 +499,7 @@ describe("Agent", () => {
});
it("should handle abort controller", () => {
const agent = new Agent({ streamFunction: unusedStreamFunction });
const agent = new Agent({ streamFn: unusedStreamFunction });
// Should not throw even if nothing is running
expect(() => agent.abort()).not.toThrow();
@@ -481,7 +509,7 @@ describe("Agent", () => {
let abortSignal: AbortSignal | undefined;
const agent = new Agent({
// Use a stream function that responds to abort
streamFunction: (_model, _context, options) => {
streamFn: (_model, _context, options) => {
abortSignal = options?.signal;
const stream = new MockAssistantStream();
queueMicrotask(() => {
@@ -520,7 +548,7 @@ describe("Agent", () => {
it("should throw when continue() called while streaming", async () => {
let abortSignal: AbortSignal | undefined;
const agent = new Agent({
streamFunction: (_model, _context, options) => {
streamFn: (_model, _context, options) => {
abortSignal = options?.signal;
const stream = new MockAssistantStream();
queueMicrotask(() => {
@@ -555,7 +583,7 @@ describe("Agent", () => {
it("continue() should process queued follow-up messages after an assistant turn", async () => {
const agent = new Agent({
streamFunction: () => {
streamFn: () => {
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") });
@@ -594,7 +622,7 @@ describe("Agent", () => {
it("continue() should keep one-at-a-time steering semantics from assistant tail", async () => {
let responseCount = 0;
const agent = new Agent({
streamFunction: () => {
streamFn: () => {
const stream = new MockAssistantStream();
responseCount++;
queueMicrotask(() => {
@@ -652,7 +680,7 @@ describe("Agent", () => {
sawAbortSignal = signal instanceof AbortSignal;
return undefined;
},
streamFunction: () => {
streamFn: () => {
requestCount++;
const stream = new MockAssistantStream();
queueMicrotask(() => {
@@ -680,7 +708,7 @@ describe("Agent", () => {
let receivedSessionId: string | undefined;
const agent = new Agent({
sessionId: "session-abc",
streamFunction: (_model, _context, options) => {
streamFn: (_model, _context, options) => {
receivedSessionId = options?.sessionId;
const stream = new MockAssistantStream();
queueMicrotask(() => {