Skip to content

Commit e879b91

Browse files
cursoragentshrisukhani
cursoragent
and
shrisukhani
committed
Harden extract action fallback and token trimming
Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>
1 parent 629e5fc commit e879b91

2 files changed

Lines changed: 255 additions & 30 deletions

File tree

‎src/agent/actions/extract.test.ts‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import fs from "fs";
2+
import {
3+
ExtractActionDefinition,
4+
estimateTextTokenCount,
5+
trimMarkdownToTokenLimit,
6+
} from "@/agent/actions/extract";
7+
import type { ActionContext } from "@/types";
8+
import type { HyperAgentLLM } from "@/llm/types";
9+
10+
jest.mock("@/utils/html-to-markdown", () => ({
11+
parseMarkdown: jest.fn(),
12+
}));
13+
14+
jest.mock("@/cdp", () => ({
15+
getCDPClient: jest.fn(),
16+
}));
17+
18+
const { parseMarkdown } = jest.requireMock("@/utils/html-to-markdown") as {
19+
parseMarkdown: jest.Mock;
20+
};
21+
22+
const { getCDPClient } = jest.requireMock("@/cdp") as {
23+
getCDPClient: jest.Mock;
24+
};
25+
26+
function createMockLLM(invokeMock?: jest.Mock): HyperAgentLLM {
27+
return {
28+
invoke: invokeMock
29+
? (async (messages) => invokeMock(messages))
30+
: async () => ({
31+
role: "assistant",
32+
content: "extracted output",
33+
}),
34+
invokeStructured: async () => ({ rawText: "{}", parsed: null }),
35+
getProviderId: () => "mock",
36+
getModelId: () => "mock-model",
37+
getCapabilities: () => ({
38+
multimodal: true,
39+
toolCalling: true,
40+
jsonMode: true,
41+
}),
42+
};
43+
}
44+
45+
function createContext(
46+
llm?: HyperAgentLLM,
47+
overrides?: Partial<ActionContext>
48+
): ActionContext {
49+
return {
50+
page: {
51+
content: jest.fn().mockResolvedValue("<html>demo</html>"),
52+
} as unknown as ActionContext["page"],
53+
domState: {
54+
elements: new Map(),
55+
domState: "",
56+
xpathMap: {},
57+
backendNodeMap: {},
58+
},
59+
llm: llm ?? createMockLLM(),
60+
tokenLimit: 200,
61+
variables: [],
62+
invalidateDomCache: jest.fn(),
63+
...overrides,
64+
} as ActionContext;
65+
}
66+
67+
describe("extract action token helpers", () => {
68+
it("estimates token count as positive non-zero", () => {
69+
expect(estimateTextTokenCount("hello world")).toBeGreaterThan(0);
70+
});
71+
72+
it("trims markdown and appends truncation notice when over limit", () => {
73+
const markdown = "a".repeat(2000);
74+
const trimmed = trimMarkdownToTokenLimit(markdown, 20);
75+
76+
expect(trimmed).toContain("[Content truncated due to token limit]");
77+
expect(trimmed.length).toBeLessThan(markdown.length);
78+
});
79+
});
80+
81+
describe("ExtractActionDefinition.run", () => {
82+
beforeEach(() => {
83+
jest.clearAllMocks();
84+
parseMarkdown.mockResolvedValue("page markdown content");
85+
getCDPClient.mockResolvedValue({
86+
acquireSession: jest.fn().mockResolvedValue({
87+
send: jest.fn().mockResolvedValue({ data: "abc" }),
88+
}),
89+
});
90+
});
91+
92+
it("falls back to markdown-only extraction when screenshot capture fails", async () => {
93+
getCDPClient.mockRejectedValue(new Error("cdp unavailable"));
94+
const invoke = jest.fn().mockResolvedValue({
95+
role: "assistant",
96+
content: "fallback extraction",
97+
});
98+
const ctx = createContext(createMockLLM(invoke));
99+
100+
const result = await ExtractActionDefinition.run(ctx, {
101+
objective: "Extract price",
102+
});
103+
104+
expect(result.success).toBe(true);
105+
expect(invoke).toHaveBeenCalled();
106+
const messagesArg = invoke.mock.calls[0]?.[0];
107+
const contentParts = messagesArg?.[0]?.content as Array<{
108+
type: string;
109+
url?: string;
110+
}>;
111+
expect(contentParts).toHaveLength(1);
112+
expect(contentParts[0]?.type).toBe("text");
113+
});
114+
115+
it("does not fail when debug file writes throw", async () => {
116+
const writeSpy = jest.spyOn(fs, "writeFileSync").mockImplementation(() => {
117+
throw new Error("disk full");
118+
});
119+
const errorSpy = jest.spyOn(console, "error").mockImplementation(() => {});
120+
const ctx = createContext(undefined, { debugDir: "debug", debug: true });
121+
122+
try {
123+
const result = await ExtractActionDefinition.run(ctx, {
124+
objective: "Extract title",
125+
});
126+
expect(result.success).toBe(true);
127+
expect(errorSpy).toHaveBeenCalled();
128+
} finally {
129+
writeSpy.mockRestore();
130+
errorSpy.mockRestore();
131+
}
132+
});
133+
134+
it("returns failure when llm responds without text content", async () => {
135+
const emptyTextLlm = createMockLLM(
136+
jest.fn().mockResolvedValue({
137+
role: "assistant",
138+
content: [{ type: "tool_call", toolName: "noop", arguments: {} }],
139+
})
140+
);
141+
const ctx = createContext(emptyTextLlm);
142+
143+
const result = await ExtractActionDefinition.run(ctx, {
144+
objective: "Extract content",
145+
});
146+
147+
expect(result.success).toBe(false);
148+
expect(result.message).toContain("No content extracted");
149+
});
150+
});

‎src/agent/actions/extract.ts‎

Lines changed: 105 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { ActionContext, ActionOutput, AgentActionDefinition } from "@/types";
33
import { parseMarkdown } from "@/utils/html-to-markdown";
44
import fs from "fs";
55
import { getCDPClient } from "@/cdp";
6+
import type { HyperAgentContentPart } from "@/llm/types";
67

78
export const ExtractAction = z
89
.object({
@@ -14,6 +15,67 @@ export const ExtractAction = z
1415

1516
export type ExtractActionType = z.infer<typeof ExtractAction>;
1617

18+
const EXTRACT_TRUNCATION_NOTICE = "\n[Content truncated due to token limit]";
19+
20+
export function estimateTextTokenCount(text: string): number {
21+
const wordCount = text.match(/[A-Za-z0-9_]+/g)?.length ?? 0;
22+
const cjkCount =
23+
text.match(/[\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]/g)
24+
?.length ?? 0;
25+
const symbolCount = text.match(/[^\sA-Za-z0-9_]/g)?.length ?? 0;
26+
const characterEstimate = Math.ceil(text.length / 3.8);
27+
const lexicalEstimate = Math.ceil(
28+
wordCount * 1.1 + cjkCount + symbolCount * 0.3
29+
);
30+
return Math.max(1, characterEstimate, lexicalEstimate);
31+
}
32+
33+
export function trimMarkdownToTokenLimit(
34+
markdown: string,
35+
tokenLimit: number
36+
): string {
37+
if (estimateTextTokenCount(markdown) <= tokenLimit) {
38+
return markdown;
39+
}
40+
41+
const suffixTokens = estimateTextTokenCount(EXTRACT_TRUNCATION_NOTICE);
42+
if (tokenLimit <= suffixTokens) {
43+
return EXTRACT_TRUNCATION_NOTICE;
44+
}
45+
46+
const targetPrefixTokens = tokenLimit - suffixTokens;
47+
let low = 0;
48+
let high = markdown.length;
49+
let best = 0;
50+
51+
while (low <= high) {
52+
const mid = Math.floor((low + high) / 2);
53+
const prefix = markdown.slice(0, mid);
54+
if (estimateTextTokenCount(prefix) <= targetPrefixTokens) {
55+
best = mid;
56+
low = mid + 1;
57+
} else {
58+
high = mid - 1;
59+
}
60+
}
61+
62+
return markdown.slice(0, best) + EXTRACT_TRUNCATION_NOTICE;
63+
}
64+
65+
function writeDebugFileSafe(
66+
filePath: string,
67+
content: Buffer | string,
68+
debug?: boolean
69+
): void {
70+
try {
71+
fs.writeFileSync(filePath, content);
72+
} catch (error) {
73+
if (debug) {
74+
console.error(`[extract] Failed to write debug file "${filePath}":`, error);
75+
}
76+
}
77+
}
78+
1779
export const ExtractActionDefinition: AgentActionDefinition = {
1880
type: "extract" as const,
1981
actionParams: ExtractAction,
@@ -26,50 +88,63 @@ export const ExtractActionDefinition: AgentActionDefinition = {
2688
const markdown = await parseMarkdown(content);
2789
const objective = action.objective;
2890

29-
// Take a screenshot of the page
30-
const cdpClient = await getCDPClient(ctx.page);
31-
const cdpSession = await cdpClient.acquireSession("screenshot");
32-
const screenshot = await cdpSession.send<{ data: string }>(
33-
"Page.captureScreenshot"
34-
);
91+
// Try to take a screenshot of the page; continue with text-only extraction if unavailable
92+
let screenshotData: string | null = null;
93+
try {
94+
const cdpClient = await getCDPClient(ctx.page);
95+
const cdpSession = await cdpClient.acquireSession("screenshot");
96+
const screenshot = await cdpSession.send<{ data: string }>(
97+
"Page.captureScreenshot"
98+
);
99+
screenshotData = screenshot.data;
100+
} catch (error) {
101+
if (ctx.debug) {
102+
console.warn(
103+
"[extract] Screenshot capture unavailable, falling back to markdown-only extraction:",
104+
error
105+
);
106+
}
107+
}
35108

36109
// Save screenshot to debug dir if exists
37-
if (ctx.debugDir) {
38-
fs.writeFileSync(
110+
if (ctx.debugDir && screenshotData) {
111+
writeDebugFileSafe(
39112
`${ctx.debugDir}/extract-screenshot.png`,
40-
Buffer.from(screenshot.data, "base64")
113+
Buffer.from(screenshotData, "base64"),
114+
ctx.debug
41115
);
42116
}
43117

44-
// Trim markdown to stay within token limit
45-
// TODO: this is a hack, we should use a better token counting method
46-
const avgTokensPerChar = 0.75; // Conservative estimate of tokens per character
47-
const maxChars = Math.floor(ctx.tokenLimit / avgTokensPerChar);
48-
const trimmedMarkdown =
49-
markdown.length > maxChars
50-
? markdown.slice(0, maxChars) + "\n[Content truncated due to length]"
51-
: markdown;
118+
const trimmedMarkdown = trimMarkdownToTokenLimit(markdown, ctx.tokenLimit);
52119
if (ctx.debugDir) {
53-
fs.writeFileSync(
120+
writeDebugFileSafe(
54121
`${ctx.debugDir}/extract-markdown-content.md`,
55-
trimmedMarkdown
122+
trimmedMarkdown,
123+
ctx.debug
56124
);
57125
}
58126

127+
const textPrompt = screenshotData
128+
? `Extract the following information from the page according to this objective: "${objective}"\n\nPage content:\n${trimmedMarkdown}\nHere is a screenshot of the page:\n`
129+
: `Extract the following information from the page according to this objective: "${objective}"\n\nPage content:\n${trimmedMarkdown}\nNo screenshot was available. Use the page content to extract the answer.`;
130+
const contentParts: HyperAgentContentPart[] = [
131+
{
132+
type: "text",
133+
text: textPrompt,
134+
},
135+
];
136+
if (screenshotData) {
137+
contentParts.push({
138+
type: "image",
139+
url: `data:image/png;base64,${screenshotData}`,
140+
mimeType: "image/png",
141+
});
142+
}
143+
59144
const response = await ctx.llm.invoke([
60145
{
61146
role: "user",
62-
content: [
63-
{
64-
type: "text",
65-
text: `Extract the following information from the page according to this objective: "${objective}"\n\nPage content:\n${trimmedMarkdown}\nHere is a screenshot of the page:\n`,
66-
},
67-
{
68-
type: "image",
69-
url: `data:image/png;base64,${screenshot.data}`,
70-
mimeType: "image/png",
71-
},
72-
],
147+
content: contentParts,
73148
},
74149
]);
75150
// Handle both string and HyperAgentContentPart[] responses

0 commit comments

Comments
 (0)