wehub-resource-sync
122 lines
4.3 KiB
TypeScript
122 lines
4.3 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, it } from "bun:test";
|
|
import * as path from "node:path";
|
|
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
|
|
import type { AssistantMessage, Context } from "@oh-my-pi/pi-ai";
|
|
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
|
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
|
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
|
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
|
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
|
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
|
import { type CustomMessage, convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
|
|
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
|
import { TempDir } from "@oh-my-pi/pi-utils";
|
|
import { type } from "arktype";
|
|
|
|
const zeroUsage = {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
} satisfies AssistantMessage["usage"];
|
|
|
|
describe("AgentSession tool-call loop guard", () => {
|
|
let tempDir: TempDir;
|
|
let authStorage: AuthStorage;
|
|
let session: AgentSession | undefined;
|
|
|
|
beforeEach(async () => {
|
|
tempDir = TempDir.createSync("@pi-tool-call-loop-guard-");
|
|
authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db"));
|
|
authStorage.setRuntimeApiKey("openai", "openai-test-key");
|
|
});
|
|
|
|
afterEach(async () => {
|
|
await session?.dispose();
|
|
authStorage.close();
|
|
tempDir.removeSync();
|
|
});
|
|
|
|
it("injects a hidden redirect before the next model call", async () => {
|
|
const model = createMockModel({ provider: "openai", id: "gpt-test" }).model;
|
|
const modelRegistry = new ModelRegistry(authStorage);
|
|
const contexts: Context[] = [];
|
|
const bashTool: AgentTool = {
|
|
name: "bash",
|
|
label: "Bash",
|
|
description: "Mock bash tool",
|
|
parameters: type({ "command?": "string" }),
|
|
execute: async () => ({ content: [{ type: "text" as const, text: "1263 passed, 4 skipped" }] }),
|
|
};
|
|
let callCount = 0;
|
|
const agent = new Agent({
|
|
getApiKey: () => "test-key",
|
|
initialState: { model, systemPrompt: ["Test"], tools: [bashTool], messages: [] },
|
|
convertToLlm,
|
|
streamFn: (_model, context) => {
|
|
contexts.push(context);
|
|
const toolCallTurn = callCount < 5;
|
|
const toolCallId = `tc-${callCount}`;
|
|
callCount++;
|
|
const message: AssistantMessage = toolCallTurn
|
|
? {
|
|
role: "assistant",
|
|
content: [{ type: "toolCall", id: toolCallId, name: "bash", arguments: { command: "pytest -q" } }],
|
|
api: model.api,
|
|
provider: model.provider,
|
|
model: model.id,
|
|
usage: zeroUsage,
|
|
stopReason: "toolUse",
|
|
timestamp: Date.now(),
|
|
}
|
|
: {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "Stopped repeating." }],
|
|
api: model.api,
|
|
provider: model.provider,
|
|
model: model.id,
|
|
usage: zeroUsage,
|
|
stopReason: "stop",
|
|
timestamp: Date.now(),
|
|
};
|
|
const stream = new AssistantMessageEventStream();
|
|
queueMicrotask(() => {
|
|
stream.push({ type: "start", partial: message });
|
|
stream.push({ type: "done", reason: toolCallTurn ? "toolUse" : "stop", message });
|
|
});
|
|
return stream;
|
|
},
|
|
});
|
|
const settings = Settings.isolated({
|
|
"compaction.enabled": false,
|
|
"todo.enabled": false,
|
|
"model.toolCallLoopGuard.enabled": true,
|
|
"model.toolCallLoopGuard.threshold": 5,
|
|
"model.toolCallLoopGuard.exemptTools": ["job", "irc"],
|
|
});
|
|
settings.setModelRole("default", `${model.provider}/${model.id}`);
|
|
session = new AgentSession({
|
|
agent,
|
|
sessionManager: SessionManager.inMemory(tempDir.path()),
|
|
settings,
|
|
modelRegistry,
|
|
toolRegistry: new Map([[bashTool.name, bashTool]]),
|
|
});
|
|
|
|
await session.prompt("run checks");
|
|
await session.waitForIdle();
|
|
|
|
expect(contexts).toHaveLength(6);
|
|
expect(JSON.stringify(contexts[5]!.messages)).toContain("tool_call_loop_detected");
|
|
expect(JSON.stringify(contexts[5]!.messages)).toContain("1263 passed, 4 skipped");
|
|
const redirects = session.agent.state.messages.filter(
|
|
(message): message is CustomMessage =>
|
|
message.role === "custom" && message.customType === "tool-call-loop-redirect",
|
|
);
|
|
expect(redirects).toHaveLength(1);
|
|
expect(redirects[0]!.display).toBe(false);
|
|
});
|
|
});
|