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-loop.test.ts
+35
View File
@@ -9,6 +9,7 @@ import {
import { Type } from "typebox";
import { describe, expect, it } from "vitest";
import { agentLoop, agentLoopContinue } from "../src/agent-loop.ts";
import { setDefaultStreamFn } from "../src/index.ts";
import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.ts";
// Mock stream for testing - mimics MockAssistantStream
@@ -80,6 +81,40 @@ function identityConverter(messages: AgentMessage[]): Message[] {
return messages.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];
}
describe("default stream function compatibility", () => {
it("uses the configured default when a legacy caller omits streamFn", async () => {
let calls = 0;
setDefaultStreamFn(() => {
calls++;
const stream = new MockAssistantStream();
queueMicrotask(() => {
stream.push({
type: "done",
reason: "stop",
message: createAssistantMessage([{ type: "text", text: "fallback" }]),
});
});
return stream;
});
try {
const context: AgentContext = { systemPrompt: "", messages: [], tools: [] };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter };
const stream = Reflect.apply(agentLoop, undefined, [
[createUserMessage("Hello")],
context,
config,
undefined,
]) as ReturnType<typeof agentLoop>;
await stream.result();
expect(calls).toBe(1);
} finally {
setDefaultStreamFn(undefined);
}
});
});
describe("agentLoop with AgentMessage", () => {
it("should emit events with AgentMessage types", async () => {
const context: AgentContext = {
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(() => {
packages/agent/test/e2e.test.ts
+10 -10
View File
@@ -38,7 +38,7 @@ afterEach(() => {
async function basicPrompt(model: Model<string>) {
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant. Keep your responses concise.",
model,
@@ -61,7 +61,7 @@ async function basicPrompt(model: Model<string>) {
async function toolExecution(model: Model<string>) {
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant. Always use the calculator tool for math.",
model,
@@ -101,7 +101,7 @@ async function toolExecution(model: Model<string>) {
async function abortExecution(model: Model<string>) {
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant.",
model,
@@ -129,7 +129,7 @@ async function abortExecution(model: Model<string>) {
async function stateUpdates(model: Model<string>) {
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant.",
model,
@@ -162,7 +162,7 @@ async function stateUpdates(model: Model<string>) {
async function multiTurnConversation(model: Model<string>) {
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant.",
model,
@@ -244,7 +244,7 @@ describe("Agent integration with faux provider", () => {
faux.setResponses([fauxAssistantMessage([fauxThinking("step by step"), fauxText("4")])]);
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant.",
model: faux.getModel(),
@@ -269,7 +269,7 @@ describe("Agent.continue() with faux provider", () => {
it("throws when no messages in context", async () => {
const faux = createFauxRegistration();
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "Test",
model: faux.getModel(),
@@ -283,7 +283,7 @@ describe("Agent.continue() with faux provider", () => {
const faux = createFauxRegistration();
const model = faux.getModel();
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "Test",
model,
@@ -318,7 +318,7 @@ describe("Agent.continue() with faux provider", () => {
const faux = createFauxRegistration();
faux.setResponses([fauxAssistantMessage("HELLO WORLD")]);
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt: "You are a helpful assistant. Follow instructions exactly.",
model: faux.getModel(),
@@ -353,7 +353,7 @@ describe("Agent.continue() with faux provider", () => {
const model = faux.getModel();
faux.setResponses([fauxAssistantMessage("The answer is 8.")]);
const agent = new Agent({
streamFunction: streamSimple,
streamFn: streamSimple,
initialState: {
systemPrompt:
"You are a helpful assistant. After getting a calculation result, state the answer clearly.",