Skip to content

Commit 9b823c6

Browse files
steipetesteipete
committed
refactor(agents): extract attempt context guards and trajectory
1 parent 7f73f6e commit 9b823c6

5 files changed

Lines changed: 532 additions & 211 deletions

File tree

‎src/agents/embedded-agent-runner/run/attempt-context-guards.test.ts‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
2+
import type { AgentMessage } from "../../runtime/index.js";
3+
import type { AgentSession } from "../../sessions/index.js";
4+
import type { MidTurnPrecheckRequest } from "./midturn-precheck.js";
5+
6+
const hoisted = vi.hoisted(() => ({
7+
installContextEngineLoopHook: vi.fn(),
8+
installToolResultContextGuard: vi.fn(),
9+
installHistoryImagePruneContextTransform: vi.fn(),
10+
invalidateComputerFrameIfMissing: vi.fn(),
11+
}));
12+
13+
vi.mock("../tool-result-context-guard.js", () => ({
14+
installContextEngineLoopHook: hoisted.installContextEngineLoopHook,
15+
installToolResultContextGuard: hoisted.installToolResultContextGuard,
16+
}));
17+
vi.mock("./history-image-prune.js", () => ({
18+
installHistoryImagePruneContextTransform: hoisted.installHistoryImagePruneContextTransform,
19+
}));
20+
vi.mock("../../tools/computer-tool.js", () => ({
21+
invalidateComputerFrameIfMissing: hoisted.invalidateComputerFrameIfMissing,
22+
}));
23+
24+
import { installEmbeddedAttemptContextGuards } from "./attempt-context-guards.js";
25+
26+
function createInput(overrides: Record<string, unknown> = {}) {
27+
const activeSession = {
28+
agent: { transformContext: undefined },
29+
} as unknown as AgentSession;
30+
const settingsManager = {
31+
getBlockImages: vi.fn(() => false),
32+
getCompactionReserveTokens: vi.fn(() => 64),
33+
} as unknown as AgentSession["settingsManager"];
34+
return {
35+
activeSession,
36+
agentDir: "/tmp/agent",
37+
attempt: {
38+
config: {
39+
agents: { defaults: { compaction: { midTurnPrecheck: { enabled: true } } } },
40+
},
41+
contextTokenBudget: 1_024,
42+
model: { api: "anthropic-messages", contextWindow: 2_048 },
43+
modelId: "model-1",
44+
provider: "provider-1",
45+
sessionFile: "/tmp/session.jsonl",
46+
},
47+
computerContextEpoch: { value: 3 },
48+
effectiveCwd: "/tmp/workspace",
49+
effectiveWorkspace: "/tmp/workspace",
50+
getPrePromptMessageCount: () => 4,
51+
getPromptCache: () => undefined,
52+
getPromptCacheRetention: () => "short" as const,
53+
getSystemPrompt: () => "system prompt",
54+
isOpenAIResponsesApi: false,
55+
repairToolUseResultPairing: false,
56+
sessionAgentId: "main",
57+
sessionManager: {},
58+
settingsManager,
59+
...overrides,
60+
};
61+
}
62+
63+
describe("installEmbeddedAttemptContextGuards", () => {
64+
beforeEach(() => {
65+
vi.clearAllMocks();
66+
hoisted.installContextEngineLoopHook.mockReturnValue(vi.fn());
67+
hoisted.installToolResultContextGuard.mockReturnValue(vi.fn());
68+
hoisted.installHistoryImagePruneContextTransform.mockReturnValue(vi.fn());
69+
});
70+
71+
it("tracks mid-turn requests and restores attempt-local transforms", async () => {
72+
const input = createInput();
73+
const originalTransform = input.activeSession.agent.transformContext;
74+
const guards = installEmbeddedAttemptContextGuards(input as never);
75+
const guardOptions = hoisted.installToolResultContextGuard.mock.calls[0]?.[0];
76+
const request: MidTurnPrecheckRequest = {
77+
route: "compact_then_truncate",
78+
estimatedPromptTokens: 1_200,
79+
promptBudgetBeforeReserve: 1_024,
80+
overflowTokens: 176,
81+
toolResultReducibleChars: 800,
82+
effectiveReserveTokens: 64,
83+
};
84+
guardOptions.midTurnPrecheck.onMidTurnPrecheck(request);
85+
86+
expect(guards.takePendingMidTurnPrecheckRequest()).toBe(request);
87+
expect(guards.takePendingMidTurnPrecheckRequest()).toBeNull();
88+
expect(guardOptions).toMatchObject({
89+
contextWindowTokens: 1_024,
90+
midTurnPrecheck: {
91+
enabled: true,
92+
contextTokenBudget: 1_024,
93+
toolResultMaxChars: expect.any(Number),
94+
},
95+
});
96+
97+
const messages: AgentMessage[] = [
98+
{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: 1 },
99+
];
100+
await input.activeSession.agent.transformContext?.(messages, new AbortController().signal);
101+
expect(hoisted.invalidateComputerFrameIfMissing).toHaveBeenCalledWith({
102+
contextEpoch: input.computerContextEpoch,
103+
imagesBlocked: false,
104+
messages,
105+
});
106+
107+
const removeToolResultGuard = hoisted.installToolResultContextGuard.mock.results[0]?.value;
108+
const removeHistoryGuard =
109+
hoisted.installHistoryImagePruneContextTransform.mock.results[0]?.value;
110+
guards.remove();
111+
expect(input.activeSession.agent.transformContext).toBe(originalTransform);
112+
expect(removeHistoryGuard).toHaveBeenCalledOnce();
113+
expect(removeToolResultGuard).toHaveBeenCalledOnce();
114+
});
115+
116+
it("composes context-engine and tool-result cleanup while exposing checkpoints", () => {
117+
const activeContextEngine = {
118+
info: { id: "test-engine", ownsCompaction: true },
119+
};
120+
const guards = installEmbeddedAttemptContextGuards(
121+
createInput({
122+
activeContextEngine,
123+
repairToolUseResultPairing: true,
124+
}) as never,
125+
);
126+
const loopOptions = hoisted.installContextEngineLoopHook.mock.calls[0]?.[0];
127+
loopOptions.onAfterTurnCheckpoint(17);
128+
129+
expect(guards.getAfterTurnCheckpoint()).toBe(17);
130+
expect(loopOptions).toMatchObject({
131+
contextEngine: activeContextEngine,
132+
modelId: "model-1",
133+
repairAssembledMessages: expect.any(Function),
134+
});
135+
136+
const removeLoopHook = hoisted.installContextEngineLoopHook.mock.results[0]?.value;
137+
const removeToolResultGuard = hoisted.installToolResultContextGuard.mock.results[0]?.value;
138+
guards.remove();
139+
expect(removeToolResultGuard).toHaveBeenCalledOnce();
140+
expect(removeLoopHook).toHaveBeenCalledOnce();
141+
});
142+
});

‎src/agents/embedded-agent-runner/run/attempt-context-guards.ts‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
/** Installs attempt-local context engine, tool-result, image, and frame guards. */
2+
import { OPENCLAW_EMBEDDED_CONTEXT_ENGINE_HOST } from "../../../context-engine/host-compat.js";
3+
import { buildContextEngineRuntimeSettings } from "../../../context-engine/runtime-settings.js";
4+
import type { ContextEngine } from "../../../context-engine/types.js";
5+
import { isHeartbeatLifecycleRunKind } from "../../bootstrap-mode.js";
6+
import { DEFAULT_CONTEXT_TOKENS } from "../../defaults.js";
7+
import type { guardSessionManager } from "../../session-tool-result-guard-wrapper.js";
8+
import type { AgentSession } from "../../sessions/index.js";
9+
import { invalidateComputerFrameIfMissing } from "../../tools/computer-tool.js";
10+
import { readLastCacheTtlTimestamp } from "../cache-ttl.js";
11+
import {
12+
installContextEngineLoopHook,
13+
installToolResultContextGuard,
14+
} from "../tool-result-context-guard.js";
15+
import { resolveLiveToolResultMaxChars } from "../tool-result-truncation.js";
16+
import { repairAttemptToolUseResultPairing } from "./attempt-transcript-helpers.js";
17+
import { buildLoopPromptCacheInfo } from "./attempt.context-engine-helpers.js";
18+
import { buildAfterTurnRuntimeContext } from "./attempt.prompt-helpers.js";
19+
import { installHistoryImagePruneContextTransform } from "./history-image-prune.js";
20+
import type { MidTurnPrecheckRequest } from "./midturn-precheck.js";
21+
import type { EmbeddedRunAttemptParams, EmbeddedRunAttemptResult } from "./types.js";
22+
23+
type PromptCacheRetention = Parameters<typeof buildLoopPromptCacheInfo>[0]["retention"];
24+
25+
export function installEmbeddedAttemptContextGuards(input: {
26+
activeContextEngine?: ContextEngine;
27+
activeSession: AgentSession;
28+
agentDir: string;
29+
attempt: EmbeddedRunAttemptParams;
30+
computerContextEpoch: { value: number };
31+
effectiveCwd: string;
32+
effectiveWorkspace: string;
33+
getPrePromptMessageCount: () => number;
34+
getPromptCache: () => EmbeddedRunAttemptResult["promptCache"];
35+
getPromptCacheRetention: () => PromptCacheRetention;
36+
getSystemPrompt: () => string;
37+
isOpenAIResponsesApi: boolean;
38+
repairToolUseResultPairing: boolean;
39+
sessionAgentId: string;
40+
sessionManager: ReturnType<typeof guardSessionManager>;
41+
settingsManager: AgentSession["settingsManager"];
42+
}): {
43+
getAfterTurnCheckpoint: () => number | null;
44+
remove: () => void;
45+
takePendingMidTurnPrecheckRequest: () => MidTurnPrecheckRequest | null;
46+
} {
47+
const { activeContextEngine, activeSession, attempt, settingsManager } = input;
48+
const contextTokenBudget = Math.max(
49+
1,
50+
Math.floor(
51+
attempt.contextTokenBudget ??
52+
attempt.model.contextWindow ??
53+
attempt.model.maxTokens ??
54+
DEFAULT_CONTEXT_TOKENS,
55+
),
56+
);
57+
const toolResultMaxChars = resolveLiveToolResultMaxChars({
58+
contextWindowTokens: contextTokenBudget,
59+
cfg: attempt.config,
60+
agentId: input.sessionAgentId,
61+
});
62+
let pendingMidTurnPrecheckRequest: MidTurnPrecheckRequest | null = null;
63+
let afterTurnCheckpoint: number | null = null;
64+
const midTurnPrecheckOptions =
65+
attempt.config?.agents?.defaults?.compaction?.midTurnPrecheck?.enabled === true
66+
? {
67+
midTurnPrecheck: {
68+
enabled: true,
69+
contextTokenBudget,
70+
reserveTokens: () => settingsManager.getCompactionReserveTokens(),
71+
toolResultMaxChars,
72+
getSystemPrompt: input.getSystemPrompt,
73+
getPrePromptMessageCount: input.getPrePromptMessageCount,
74+
onMidTurnPrecheck: (request: MidTurnPrecheckRequest) => {
75+
pendingMidTurnPrecheckRequest = request;
76+
},
77+
},
78+
}
79+
: {};
80+
81+
let removeLoopGuard: () => void;
82+
if (activeContextEngine?.info.ownsCompaction === true) {
83+
const selectedContextEngineId = activeContextEngine.info.id;
84+
const runtimeSettings = buildContextEngineRuntimeSettings({
85+
contextEngineHost: OPENCLAW_EMBEDDED_CONTEXT_ENGINE_HOST,
86+
provider: attempt.provider,
87+
requestedModel: attempt.requestedModelId,
88+
resolvedModel: attempt.modelId,
89+
selectedContextEngineId,
90+
contextEngineSelectionSource: selectedContextEngineId === "legacy" ? "default" : "configured",
91+
promptTokenBudget: attempt.contextTokenBudget,
92+
fallbackReason: attempt.fallbackReason,
93+
degradedReason: attempt.degradedReason,
94+
});
95+
const removeContextEngineLoopHook = installContextEngineLoopHook({
96+
agent: activeSession.agent,
97+
contextEngine: activeContextEngine,
98+
sessionId: attempt.sessionId,
99+
sessionKey: attempt.sessionKey,
100+
sessionTarget: attempt.sessionTarget,
101+
sessionFile: attempt.sessionFile,
102+
tokenBudget: attempt.contextTokenBudget,
103+
modelId: attempt.modelId,
104+
...(input.repairToolUseResultPairing
105+
? {
106+
repairAssembledMessages: (messages) =>
107+
repairAttemptToolUseResultPairing(messages, input.isOpenAIResponsesApi),
108+
}
109+
: {}),
110+
getPrePromptMessageCount: input.getPrePromptMessageCount,
111+
onAfterTurnCheckpoint: (messageCount) => {
112+
afterTurnCheckpoint = messageCount;
113+
},
114+
getRuntimeContext: ({ messages, prePromptMessageCount }) =>
115+
buildAfterTurnRuntimeContext({
116+
attempt,
117+
workspaceDir: input.effectiveWorkspace,
118+
cwd: input.effectiveCwd,
119+
agentDir: input.agentDir,
120+
tokenBudget: attempt.contextTokenBudget,
121+
promptCache:
122+
input.getPromptCache() ??
123+
buildLoopPromptCacheInfo({
124+
messagesSnapshot: messages,
125+
prePromptMessageCount,
126+
retention: input.getPromptCacheRetention(),
127+
fallbackLastCacheTouchAt: readLastCacheTtlTimestamp(input.sessionManager, {
128+
provider: attempt.provider,
129+
modelId: attempt.modelId,
130+
}),
131+
}),
132+
}),
133+
runtimeSettings,
134+
isHeartbeat: isHeartbeatLifecycleRunKind(attempt.bootstrapContextRunKind),
135+
});
136+
const removeToolResultGuard = installToolResultContextGuard({
137+
agent: activeSession.agent,
138+
contextWindowTokens: contextTokenBudget,
139+
...midTurnPrecheckOptions,
140+
});
141+
removeLoopGuard = () => {
142+
removeToolResultGuard();
143+
removeContextEngineLoopHook();
144+
};
145+
} else {
146+
removeLoopGuard = installToolResultContextGuard({
147+
agent: activeSession.agent,
148+
contextWindowTokens: contextTokenBudget,
149+
...midTurnPrecheckOptions,
150+
});
151+
}
152+
153+
const removeHistoryImagePruneContextTransform = installHistoryImagePruneContextTransform(
154+
activeSession.agent,
155+
);
156+
const previousComputerFrameTransform = activeSession.agent.transformContext;
157+
activeSession.agent.transformContext = async (messages, signal) => {
158+
const transformed = previousComputerFrameTransform
159+
? await previousComputerFrameTransform.call(activeSession.agent, messages, signal)
160+
: messages;
161+
const modelContext = Array.isArray(transformed) ? transformed : messages;
162+
invalidateComputerFrameIfMissing({
163+
contextEpoch: input.computerContextEpoch,
164+
messages: modelContext,
165+
imagesBlocked: settingsManager.getBlockImages(),
166+
});
167+
return modelContext;
168+
};
169+
170+
return {
171+
getAfterTurnCheckpoint: () => afterTurnCheckpoint,
172+
remove: () => {
173+
activeSession.agent.transformContext = previousComputerFrameTransform;
174+
removeHistoryImagePruneContextTransform();
175+
removeLoopGuard();
176+
},
177+
takePendingMidTurnPrecheckRequest: () => {
178+
const request = pendingMidTurnPrecheckRequest;
179+
pendingMidTurnPrecheckRequest = null;
180+
return request;
181+
},
182+
};
183+
}

0 commit comments

Comments
 (0)