alexion/dev
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore(vendor): import pi 1.3.0

Browse Source 
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree.

git-subtree-dir: third_party/pi

git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
This commit is contained in:
alexion
2026-10-15 17:58:48 -04:00
parent 719370c8f7
commit 6081e3e325
1743 changed files with 205460 additions and 131705 deletions
Download Patch File Download Diff File Expand all files Collapse all files

167
third_party/pi/packages/coding-agent/test/sdk-stream-options.test.ts vendored
View File

@@ -1,4 +1,4 @@
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
@@ -6,14 +6,16 @@ import {
type AssistantMessage,
createAssistantMessageEventStream,
type Model,
normalizeContext,
type SimpleStreamOptions,
} from "@earendil-works/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts";
import type { ExtensionFactory } from "../src/core/extensions/types.ts";
import { DefaultResourceLoader } from "../src/core/resource-loader.ts";
import { createAgentSession } from "../src/core/sdk.ts";
import { SessionManager } from "../src/core/session-manager.ts";
import { type Settings, SettingsManager } from "../src/core/settings-manager.ts";
import { createModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts";
describe("createAgentSession stream options", () => {
@@ -51,9 +53,8 @@ describe("createAgentSession stream options", () => {
};
}
function createDoneStream(api: Api) {
const stream = createAssistantMessageEventStream();
const message: AssistantMessage = {
function createDoneMessage(api: Api, promptTokens = 0): AssistantMessage {
return {
role: "assistant",
content: [{ type: "text", text: "ok" }],
api,
@@ -62,15 +63,19 @@ describe("createAgentSession stream options", () => {
usage: {
input: 0,
output: 0,
cacheRead: 0,
cacheRead: promptTokens,
cacheWrite: 0,
totalTokens: 0,
totalTokens: promptTokens,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: Date.now(),
};
stream.end(message);
}
function createDoneStream(api: Api, promptTokens = 0) {
const stream = createAssistantMessageEventStream();
stream.end(createDoneMessage(api, promptTokens));
return stream;
}
@@ -78,15 +83,18 @@ describe("createAgentSession stream options", () => {
api: Api,
settings: Partial<Settings>,
requestOptions: SimpleStreamOptions = {},
extensionSource?: string,
extensionFactory?: ExtensionFactory,
providerEvent?: unknown,
): Promise<SimpleStreamOptions | undefined> {
const model = createModel(api);
const settingsManager = SettingsManager.inMemory(settings);
if (extensionSource) {
const extensionsDir = join(agentDir, "extensions");
mkdirSync(extensionsDir, { recursive: true });
writeFileSync(join(extensionsDir, "headers.ts"), extensionSource);
}
const resourceLoader = new DefaultResourceLoader({
cwd,
agentDir,
settingsManager,
extensionFactories: extensionFactory ? [extensionFactory] : [],
});
await resourceLoader.reload();
const authStorage = AuthStorage.create(join(agentDir, "auth.json"));
await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "test-api-key" }));
@@ -96,9 +104,16 @@ describe("createAgentSession stream options", () => {
modelRegistry.registerProvider(model.provider, {
api,
headers: { "x-provider": "provider" },
streamSimple: (_model, _context, providerOptions) => {
streamSimple: (requestModel, _context, providerOptions) => {
capturedOptions = providerOptions;
return createDoneStream(api);
if (providerEvent === undefined) return createDoneStream(api);
const stream = createAssistantMessageEventStream();
void (async () => {
await providerOptions?.onProviderStreamEvent?.(providerEvent, requestModel);
stream.end(createDoneMessage(api));
})();
return stream;
},
});
@@ -111,11 +126,20 @@ describe("createAgentSession stream options", () => {
modelRuntime,
settingsManager,
sessionManager,
resourceLoader,
});
try {
const stream = await session.agent.streamFunction(model, { messages: [] }, requestOptions);
await stream.result();
if (providerEvent === undefined) {
const stream = await session.agent.streamFunction(
model,
normalizeContext({ messages: [] }),
requestOptions,
);
await stream.result();
} else {
await session.prompt("test");
}
return capturedOptions;
} finally {
session.dispose();
@@ -123,6 +147,80 @@ describe("createAgentSession stream options", () => {
}
}
async function createCacheWarmingSession(populate?: (manager: SessionManager, model: Model<Api>) => void) {
const model: Model<Api> = {
...createModel("anthropic-messages"),
cost: { input: 10, output: 50, cacheRead: 0.25, cacheWrite: 12.5 },
promptCache: { short: 300 },
};
const authStorage = AuthStorage.create(join(agentDir, "auth.json"));
await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "test-api-key" }));
const modelRegistry = await createModelRegistry(authStorage, join(agentDir, "models.json"));
let providerCalls = 0;
modelRegistry.registerProvider(model.provider, {
api: model.api,
streamSimple: () => {
providerCalls++;
return createDoneStream(model.api, 100_000);
},
});
const sessionManager = SessionManager.inMemory(cwd);
populate?.(sessionManager, model);
const { session } = await createAgentSession({
cwd,
agentDir,
model,
modelRuntime: getModelRuntime(modelRegistry),
settingsManager: SettingsManager.inMemory({ cacheWarming: "idle" }),
sessionManager,
});
return {
session,
providerCalls: () => providerCalls,
dispose: () => {
session.dispose();
modelRegistry.unregisterProvider(model.provider);
},
};
}
it("schedules cache warming after a completed session request", async () => {
const fixture = await createCacheWarmingSession();
try {
await fixture.session.prompt("test");
expect(fixture.session.cacheWarmingStatus?.nextWarmAt).toBeGreaterThan(Date.now());
// Equivalent shallow copies remain current, but removing the request prefix does not.
fixture.session.agent.state.messages = [...fixture.session.agent.state.messages];
fixture.session.agent.state.model = { ...fixture.session.agent.state.model };
expect(fixture.session.cacheWarmingStatus?.nextWarmAt).toBeGreaterThan(Date.now());
fixture.session.agent.state.messages = fixture.session.agent.state.messages.slice(1);
expect(fixture.session.cacheWarmingStatus?.reason).toBe("conversation context changed");
} finally {
fixture.dispose();
}
});
it("waits for the next request instead of restoring cache warming", async () => {
const fixture = await createCacheWarmingSession((manager, model) => {
manager.appendModelChange(model.provider, model.id);
manager.appendThinkingLevelChange("off");
manager.appendMessage({ role: "user", content: "test", timestamp: Date.now() - 60_000 });
const assistant = { ...createDoneMessage(model.api, 100_000), timestamp: Date.now() - 59_000 };
manager.appendMessage(assistant);
manager.appendUsage("cache_warm", model.provider, model.id, assistant.usage);
});
try {
expect(fixture.providerCalls()).toBe(0);
expect(fixture.session.cacheWarmingStatus).toEqual({
state: "inactive",
reason: "waiting for first request",
});
} finally {
fixture.dispose();
}
});
it("forwards httpIdleTimeoutMs as timeoutMs for OpenAI Codex", async () => {
const options = await captureStreamOptions("openai-codex-responses", { httpIdleTimeoutMs: 1234 });
@@ -170,12 +268,41 @@ describe("createAgentSession stream options", () => {
expect(options?.maxRetryDelayMs).toBe(3000);
});
// Regression test for #9784.
it("forwards provider stream events to extensions", async () => {
const providerEvent = { openrouter_metadata: { strategy: "direct" } };
const extensionEvents: unknown[] = [];
const options = await captureStreamOptions(
"openai-completions",
{},
{},
(pi) => {
pi.on("provider_stream_event", (event) => {
extensionEvents.push(event);
});
},
providerEvent,
);
expect(options?.onProviderStreamEvent).toEqual(expect.any(Function));
expect(extensionEvents).toEqual([
{
data: providerEvent,
type: "provider_stream_event",
provider: "capture-provider",
api: "openai-completions",
model: "capture-model",
},
]);
});
it("runs before_provider_headers on assembled headers without forwarding the transform", async () => {
const options = await captureStreamOptions(
"openai-completions",
{},
{ headers: { "x-explicit": "explicit" } },
`export default function (pi) {
(pi) => {
pi.on("before_provider_headers", (event) => {
event.headers["x-hook"] = [
event.headers["x-provider"],
@@ -183,7 +310,7 @@ describe("createAgentSession stream options", () => {
event.headers["x-explicit"],
].join(":");
});
}`,
},
);
expect(options?.headers).toMatchObject({