From a1e0090fe428e55fbb0269a3fa1a1643db91da81 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 7 Apr 2026 05:58:02 +0100 Subject: [PATCH] perf(agents): trim fast tool test seams --- src/agents/openclaw-tools.agents.test.ts | 35 +++----- .../openclaw-tools.session-status.test.ts | 67 +++++++++++++- .../compact.hooks.harness.ts | 7 ++ .../pi-embedded-runner/compact.hooks.test.ts | 89 +++++++++---------- src/agents/pi-embedded-runner/compact.ts | 82 ++++++++++++++--- 5 files changed, 194 insertions(+), 86 deletions(-) diff --git a/src/agents/openclaw-tools.agents.test.ts b/src/agents/openclaw-tools.agents.test.ts index 4a26b62ec97..9285cc0db0a 100644 --- a/src/agents/openclaw-tools.agents.test.ts +++ b/src/agents/openclaw-tools.agents.test.ts @@ -1,5 +1,6 @@ -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { createPerSenderSessionConfig } from "./test-helpers/session-config.js"; +import { createAgentsListTool } from "./tools/agents-list-tool.js"; let configOverride: ReturnType<(typeof import("../config/config.js"))["loadConfig"]> = { session: createPerSenderSessionConfig(), @@ -14,10 +15,6 @@ vi.mock("../config/config.js", async () => { }; }); -import "./test-helpers/fast-core-tools.js"; - -let createOpenClawTools: typeof import("./openclaw-tools.js").createOpenClawTools; - describe("agents_list", () => { type AgentConfig = NonNullable["list"]>[number]; @@ -30,14 +27,10 @@ describe("agents_list", () => { }; } - function requireAgentsListTool() { - const tool = createOpenClawTools({ + function createTool() { + return createAgentsListTool({ agentSessionKey: "main", - }).find((candidate) => candidate.name === "agents_list"); - if (!tool) { - throw new Error("missing agents_list tool"); - } - return tool; + }); } function readAgentList(result: unknown) { @@ -45,17 +38,11 @@ describe("agents_list", () => { .details?.agents; } - beforeEach(async () => { - vi.resetModules(); + it("defaults to the requester agent only", async () => { configOverride = { session: createPerSenderSessionConfig(), }; - await import("./test-helpers/fast-core-tools.js"); - ({ createOpenClawTools } = await import("./openclaw-tools.js")); - }); - - it("defaults to the requester agent only", async () => { - const tool = requireAgentsListTool(); + const tool = createTool(); const result = await tool.execute("call1", {}); expect(result.details).toMatchObject({ requester: "main", @@ -80,7 +67,7 @@ describe("agents_list", () => { }, ]); - const tool = requireAgentsListTool(); + const tool = createTool(); const result = await tool.execute("call2", {}); const agents = readAgentList(result); expect(agents?.map((agent) => agent.id)).toEqual(["main", "research"]); @@ -108,7 +95,7 @@ describe("agents_list", () => { }, }; - const tool = requireAgentsListTool(); + const tool = createTool(); const result = await tool.execute("call2b", {}); const agents = readAgentList(result); expect(agents?.map((agent) => agent.id)).toEqual(["main", "research"]); @@ -132,7 +119,7 @@ describe("agents_list", () => { }, ]); - const tool = requireAgentsListTool(); + const tool = createTool(); const result = await tool.execute("call3", {}); expect(result.details).toMatchObject({ allowAny: true, @@ -151,7 +138,7 @@ describe("agents_list", () => { }, ]); - const tool = requireAgentsListTool(); + const tool = createTool(); const result = await tool.execute("call4", {}); const agents = readAgentList(result); expect(agents?.map((agent) => agent.id)).toEqual(["main", "research"]); diff --git a/src/agents/openclaw-tools.session-status.test.ts b/src/agents/openclaw-tools.session-status.test.ts index bf320fb5833..40f44b29013 100644 --- a/src/agents/openclaw-tools.session-status.test.ts +++ b/src/agents/openclaw-tools.session-status.test.ts @@ -8,8 +8,12 @@ const loadSessionStoreMock = vi.fn(); const updateSessionStoreMock = vi.fn(); const callGatewayMock = vi.fn(); const loadCombinedSessionStoreForGatewayMock = vi.fn(); -const buildStatusMessageMock = vi.hoisted(() => vi.fn(() => "OpenClaw\n🧠 Model: GPT-5.4")); -const resolveQueueSettingsMock = vi.hoisted(() => vi.fn(() => ({ mode: "interrupt" }))); +const buildStatusMessageMock = vi.hoisted(() => + vi.fn((_params?: unknown) => "OpenClaw\n🧠 Model: GPT-5.4"), +); +const resolveQueueSettingsMock = vi.hoisted(() => + vi.fn((_params?: unknown) => ({ mode: "interrupt" })), +); const listTasksForRelatedSessionKeyForOwnerMock = vi.hoisted(() => vi.fn( (_: { relatedSessionKey: string; callerOwnerKey: string }) => @@ -170,6 +174,64 @@ function createProviderUsageModuleMock() { }; } +function createCommandsStatusRuntimeModuleMock() { + return { + buildStatusText: async (params: { + sessionKey: string; + sessionEntry: SessionEntry; + statusChannel: string; + provider?: string; + model: string; + primaryModelLabelOverride?: string; + includeTranscriptUsage?: boolean; + taskLineOverride?: string; + resolveDefaultThinkingLevel?: () => Promise | unknown; + }) => { + resolveQueueSettingsMock({ + channel: params.statusChannel, + sessionEntry: params.sessionEntry, + }); + const parsed = params.sessionKey.startsWith("agent:") ? params.sessionKey.split(":") : null; + const agentId = parsed?.[1] || "main"; + const configuredAgent = Array.isArray( + (mockConfig as { agents?: { list?: Array> } }).agents?.list, + ) + ? (mockConfig as { agents?: { list?: Array> } }).agents?.list?.find( + (entry) => entry.id === agentId, + ) + : undefined; + const primary = + params.primaryModelLabelOverride ?? + [params.provider, params.model].filter(Boolean).join("/") ?? + params.model; + const customAuth = params.provider + ? resolveUsableCustomProviderApiKeyMock({ provider: params.provider }) + : null; + const envAuth = + !customAuth && params.provider ? resolveEnvApiKeyMock(params.provider, process.env) : null; + const modelAuth = customAuth + ? `api-key (${customAuth.source})` + : envAuth + ? "api-key (env)" + : undefined; + buildStatusMessageMock({ + agentId, + agent: { + model: { primary }, + thinkingDefault: + configuredAgent?.thinkingDefault ?? (await params.resolveDefaultThinkingLevel?.()), + }, + sessionEntry: params.sessionEntry, + modelAuth, + includeTranscriptUsage: params.includeTranscriptUsage, + }); + return ["OpenClaw", `🧠 Model: ${primary}`, params.taskLineOverride] + .filter(Boolean) + .join("\n"); + }, + }; +} + vi.mock("../config/sessions.js", createSessionsModuleMock); vi.mock("../gateway/call.js", createGatewayCallModuleMock); vi.mock("../gateway/session-utils.js", createGatewaySessionUtilsModuleMock); @@ -187,6 +249,7 @@ vi.mock("../plugins/providers.runtime.js", () => ({ vi.mock("../agents/auth-profiles.js", createAuthProfilesModuleMock); vi.mock("../agents/model-auth.js", createModelAuthModuleMock); vi.mock("../infra/provider-usage.js", createProviderUsageModuleMock); +vi.mock("../auto-reply/reply/commands-status.runtime.js", createCommandsStatusRuntimeModuleMock); vi.mock("../auto-reply/group-activation.js", () => ({ normalizeGroupActivation: (value: unknown) => value ?? "always", })); diff --git a/src/agents/pi-embedded-runner/compact.hooks.harness.ts b/src/agents/pi-embedded-runner/compact.hooks.harness.ts index 9c2abb57254..103b6e5f14a 100644 --- a/src/agents/pi-embedded-runner/compact.hooks.harness.ts +++ b/src/agents/pi-embedded-runner/compact.hooks.harness.ts @@ -87,6 +87,7 @@ export const createOpenClawCodingToolsMock = vi.fn(() => []); export const resolveEmbeddedAgentStreamFnMock: Mock< (params?: unknown) => MockEmbeddedAgentStreamFn > = vi.fn((_params?: unknown) => vi.fn()); +export const registerProviderStreamForModelMock: Mock<(params?: unknown) => unknown> = vi.fn(); export const applyExtraParamsToAgentMock = vi.fn(() => ({ effectiveExtraParams: {} })); export const resolveAgentTransportOverrideMock: Mock<(params?: unknown) => string | undefined> = vi.fn(() => undefined); @@ -133,6 +134,8 @@ export function resetCompactSessionStateMocks(): void { sessionAbortCompactionMock.mockReset(); resolveEmbeddedAgentStreamFnMock.mockReset(); resolveEmbeddedAgentStreamFnMock.mockImplementation((_params?: unknown) => vi.fn()); + registerProviderStreamForModelMock.mockReset(); + registerProviderStreamForModelMock.mockReturnValue(undefined); applyExtraParamsToAgentMock.mockReset(); applyExtraParamsToAgentMock.mockReturnValue({ effectiveExtraParams: {} }); resolveAgentTransportOverrideMock.mockReset(); @@ -201,6 +204,10 @@ export async function loadCompactHooksHarness(): Promise<{ ensureRuntimePluginsLoaded, })); + vi.doMock("../provider-stream.js", () => ({ + registerProviderStreamForModel: registerProviderStreamForModelMock, + })); + vi.doMock("../../hooks/internal-hooks.js", async () => { const actual = await vi.importActual( "../../hooks/internal-hooks.js", diff --git a/src/agents/pi-embedded-runner/compact.hooks.test.ts b/src/agents/pi-embedded-runner/compact.hooks.test.ts index 13f13fd3086..236578ab15d 100644 --- a/src/agents/pi-embedded-runner/compact.hooks.test.ts +++ b/src/agents/pi-embedded-runner/compact.hooks.test.ts @@ -1,7 +1,5 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; -import { getApiProvider, unregisterApiProviders } from "@mariozechner/pi-ai"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; -import { getCustomApiRegistrySourceId } from "../custom-api-registry.js"; import { applyExtraParamsToAgentMock, contextEngineCompactMock, @@ -10,7 +8,7 @@ import { getMemorySearchManagerMock, hookRunner, loadCompactHooksHarness, - resolveAgentTransportOverrideMock, + registerProviderStreamForModelMock, resolveContextEngineMock, resolveEmbeddedAgentStreamFnMock, resolveMemorySearchConfigMock, @@ -163,7 +161,6 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { details: { ok: true }, }); resetCompactSessionStateMocks(); - unregisterApiProviders(getCustomApiRegistrySourceId("ollama")); }); it("bootstraps runtime plugins with the resolved workspace", async () => { @@ -218,15 +215,29 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { applyExtraParamsToAgentMock.mockReturnValue({ effectiveExtraParams: { transport: "websocket" }, }); - resolveContextEngineMock.mockResolvedValue({ info: { ownsCompaction: false } } as never); - resolveAgentTransportOverrideMock.mockReturnValue("websocket"); + const session = { + agent: { + streamFn: vi.fn(), + }, + messages: [{ role: "user", content: "hello" }], + }; - await compactEmbeddedPiSessionDirect({ + compactTesting.prepareCompactionSessionAgent({ + session: session as never, + providerStreamFn: vi.fn(), + shouldUseWebSocketTransport: false, sessionId: "session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp/workspace", + signal: new AbortController().signal, + effectiveModel: { provider: "openai", id: "fake", api: "responses", input: [] } as never, + resolvedApiKey: undefined, + authStorage: { setRuntimeApiKey: vi.fn() }, + config: undefined, provider: "openai", - model: "gpt-5.4", + modelId: "gpt-5.4", + thinkLevel: "off", + sessionAgentId: "main", + effectiveWorkspace: "/tmp/workspace", + agentDir: "/tmp/workspace", }); expect(resolveEmbeddedAgentStreamFnMock).toHaveBeenCalledWith( @@ -238,15 +249,6 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { expect(applyExtraParamsToAgentMock).toHaveBeenCalledWith( expect.objectContaining({ streamFn: resolvedStreamFn, - transport: "sse", - state: expect.objectContaining({ - messages: expect.arrayContaining([ - expect.objectContaining({ - role: "user", - content: "hello", - }), - ]), - }), }), undefined, "openai", @@ -259,9 +261,8 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { provider: "openai", id: "fake", api: "responses", - contextWindow: 128_000, }), - "/tmp", + "/tmp/workspace", ); }); @@ -596,39 +597,35 @@ describe("compactEmbeddedPiSessionDirect hooks", () => { }); it("registers the Ollama api provider before compaction", async () => { - resolveContextEngineMock.mockResolvedValue({ info: { ownsCompaction: false } } as never); - resolveModelMock.mockReturnValue({ - model: { + const streamFn = vi.fn(); + registerProviderStreamForModelMock.mockReturnValue(streamFn); + + const result = compactTesting.resolveCompactionProviderStream({ + effectiveModel: { provider: "ollama", api: "ollama", id: "qwen3:8b", input: ["text"], baseUrl: "http://127.0.0.1:11434", headers: { Authorization: "Bearer ollama-cloud" }, - }, - error: null, - authStorage: { setRuntimeApiKey: vi.fn() }, - modelRegistry: {}, - } as never); - sessionCompactImpl.mockImplementation(async () => { - expect(getApiProvider("ollama" as Parameters[0])).toBeDefined(); - return { - summary: "summary", - firstKeptEntryId: "entry-1", - tokensBefore: 120, - details: { ok: true }, - }; + } as never, + config: undefined, + agentDir: "/tmp", + effectiveWorkspace: "/tmp", }); - const result = await compactEmbeddedPiSessionDirect({ - sessionId: "session-1", - sessionKey: "agent:main:session-1", - sessionFile: "/tmp/session.jsonl", - workspaceDir: "/tmp", - customInstructions: "focus on decisions", - }); - - expect(result.ok).toBe(true); + expect(result).toBe(streamFn); + expect(registerProviderStreamForModelMock).toHaveBeenCalledWith( + expect.objectContaining({ + model: expect.objectContaining({ + provider: "ollama", + api: "ollama", + id: "qwen3:8b", + }), + agentDir: "/tmp", + workspaceDir: "/tmp", + }), + ); }); it("aborts in-flight compaction when the caller abort signal fires", async () => { diff --git a/src/agents/pi-embedded-runner/compact.ts b/src/agents/pi-embedded-runner/compact.ts index b7be44f96f5..ec842f0882b 100644 --- a/src/agents/pi-embedded-runner/compact.ts +++ b/src/agents/pi-embedded-runner/compact.ts @@ -217,6 +217,63 @@ function createCompactionDiagId(): string { return `cmp-${Date.now().toString(36)}-${generateSecureToken(4)}`; } +function prepareCompactionSessionAgent(params: { + session: { agent: { streamFn?: unknown } }; + providerStreamFn: unknown; + shouldUseWebSocketTransport: boolean; + wsApiKey?: string; + sessionId: string; + signal: AbortSignal; + effectiveModel: ProviderRuntimeModel; + resolvedApiKey?: string; + authStorage: unknown; + config?: OpenClawConfig; + provider: string; + modelId: string; + thinkLevel: ThinkLevel; + sessionAgentId: string; + effectiveWorkspace: string; + agentDir: string; +}) { + params.session.agent.streamFn = resolveEmbeddedAgentStreamFn({ + currentStreamFn: resolveEmbeddedAgentBaseStreamFn({ session: params.session as never }), + providerStreamFn: params.providerStreamFn as never, + shouldUseWebSocketTransport: params.shouldUseWebSocketTransport, + wsApiKey: params.wsApiKey, + sessionId: params.sessionId, + signal: params.signal, + model: params.effectiveModel, + resolvedApiKey: params.resolvedApiKey, + authStorage: params.authStorage as never, + }); + return applyExtraParamsToAgent( + params.session.agent as never, + params.config, + params.provider, + params.modelId, + undefined, + params.thinkLevel, + params.sessionAgentId, + params.effectiveWorkspace, + params.effectiveModel, + params.agentDir, + ); +} + +function resolveCompactionProviderStream(params: { + effectiveModel: ProviderRuntimeModel; + config?: OpenClawConfig; + agentDir: string; + effectiveWorkspace: string; +}) { + return registerProviderStreamForModel({ + model: params.effectiveModel, + cfg: params.config, + agentDir: params.agentDir, + workspaceDir: params.effectiveWorkspace, + }); +} + function normalizeObservedTokenCount(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) && value > 0 ? Math.floor(value) @@ -779,11 +836,11 @@ export async function compactEmbeddedPiSessionDirect( sandboxEnabled: !!sandbox?.enabled, }); - const providerStreamFn = registerProviderStreamForModel({ - model: effectiveModel, - cfg: params.config, + const providerStreamFn = resolveCompactionProviderStream({ + effectiveModel, + config: params.config, agentDir, - workspaceDir: effectiveWorkspace, + effectiveWorkspace, }); const shouldUseWebSocketTransport = shouldUseOpenAIWebSocketTransport({ provider, @@ -824,29 +881,24 @@ export async function compactEmbeddedPiSessionDirect( applySystemPromptOverrideToSession(session, buildSystemPromptOverride(thinkLevel)()); // Compaction builds the same embedded system prompt, so it must flow // through the same transport/payload shaping stack as normal turns. - session.agent.streamFn = resolveEmbeddedAgentStreamFn({ - currentStreamFn: resolveEmbeddedAgentBaseStreamFn({ session }), + prepareCompactionSessionAgent({ + session, providerStreamFn, shouldUseWebSocketTransport, wsApiKey, sessionId: params.sessionId, signal: runAbortController.signal, - model: effectiveModel, + effectiveModel, resolvedApiKey: hasRuntimeAuthExchange ? undefined : apiKeyInfo?.apiKey, authStorage, - }); - applyExtraParamsToAgent( - session.agent, - params.config, + config: params.config, provider, modelId, - undefined, thinkLevel, sessionAgentId, effectiveWorkspace, - effectiveModel, agentDir, - ); + }); const prior = await sanitizeSessionHistory({ messages: session.messages, @@ -1392,6 +1444,8 @@ export const __testing = { estimateTokensAfterCompaction, buildBeforeCompactionHookMetrics, hardenManualCompactionBoundary, + resolveCompactionProviderStream, + prepareCompactionSessionAgent, runBeforeCompactionHooks, runAfterCompactionHooks, runPostCompactionSideEffects,