Skip to content

Commit 7817fcc

Browse files
tkattkattkattkat
authored
[feat]: add configurable timeout to agent tools (browserbase#1766)
1 parent 9df8fe6 commit 7817fcc

28 files changed

Lines changed: 631 additions & 296 deletions

‎.changeset/public-results-mate.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@browserbasehq/stagehand": patch
3+
---
4+
5+
Add configurable timeout to tools in agent

‎packages/core/lib/inference.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,11 @@ import { SupportedUnderstudyAction } from "./v3/types/private/handlers.js";
2222
export type { LLMParsedResponse, LLMUsage } from "./v3/llm/LLMClient.js";
2323

2424
function withLlmTimeout<T>(promise: Promise<T>, operation: string): Promise<T> {
25-
const timeoutMs = getEnvTimeoutMs("LLM_MAX_MS");
26-
if (!timeoutMs) return promise;
27-
return withTimeout(promise, timeoutMs, `LLM ${operation}`);
25+
return withTimeout(
26+
promise,
27+
getEnvTimeoutMs("LLM_MAX_MS"),
28+
`LLM ${operation}`,
29+
);
2830
}
2931

3032
export async function extract<T extends StagehandZodObject>({

‎packages/core/lib/v3/agent/tools/act.ts‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@ import { z } from "zod";
33
import type { V3 } from "../../v3.js";
44
import type { Action } from "../../types/public/methods.js";
55
import type { AgentModelConfig, Variables } from "../../types/public/agent.js";
6+
import { TimeoutError } from "../../types/public/sdkErrors.js";
67

78
export const actTool = (
89
v3: V3,
910
executionModel?: string | AgentModelConfig,
1011
variables?: Variables,
12+
toolTimeout?: number,
1113
) => {
1214
const hasVariables = variables && Object.keys(variables).length > 0;
1315
const actionDescription = hasVariables
@@ -34,8 +36,9 @@ export const actTool = (
3436
},
3537
});
3638
const options = executionModel
37-
? { model: executionModel, variables }
38-
: { variables };
39+
? { model: executionModel, variables, timeout: toolTimeout }
40+
: { variables, timeout: toolTimeout };
41+
3942
const result = await v3.act(action, options);
4043
const actions = (result.actions as Action[] | undefined) ?? [];
4144
v3.recordAgentReplayStep({
@@ -60,7 +63,22 @@ export const actTool = (
6063
}
6164
return response;
6265
} catch (error) {
63-
return { success: false, error: error?.message ?? String(error) };
66+
if (error instanceof TimeoutError) {
67+
const timeoutMessage = `TimeoutError while waiting for act() to complete (it may continue executing in the background)`;
68+
v3.logger({
69+
category: "agent",
70+
message: timeoutMessage,
71+
level: 0,
72+
});
73+
return {
74+
success: false,
75+
error: `${timeoutMessage} — try using a different description for the action`,
76+
};
77+
}
78+
return {
79+
success: false,
80+
error: error?.message ?? String(error),
81+
};
6482
}
6583
},
6684
});

‎packages/core/lib/v3/agent/tools/ariaTree.ts‎

Lines changed: 62 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,76 @@
11
import { tool } from "ai";
22
import { z } from "zod";
33
import type { V3 } from "../../v3.js";
4+
import { TimeoutError } from "../../types/public/sdkErrors.js";
45

5-
export const ariaTreeTool = (v3: V3) =>
6+
export const ariaTreeTool = (v3: V3, toolTimeout?: number) =>
67
tool({
78
description:
89
"gets the accessibility (ARIA) hybrid tree text for the current page. use this to understand structure and content.",
910
inputSchema: z.object({}),
1011
execute: async () => {
11-
v3.logger({
12-
category: "agent",
13-
message: `Agent calling tool: ariaTree`,
14-
level: 1,
15-
});
16-
const page = await v3.context.awaitActivePage();
17-
const { pageText } = (await v3.extract()) as { pageText: string };
18-
const pageUrl = page.url();
12+
try {
13+
v3.logger({
14+
category: "agent",
15+
message: `Agent calling tool: ariaTree`,
16+
level: 1,
17+
});
18+
const page = await v3.context.awaitActivePage();
19+
const extractOptions = toolTimeout
20+
? { timeout: toolTimeout }
21+
: undefined;
22+
const { pageText } = (await v3.extract(extractOptions)) as {
23+
pageText: string;
24+
};
25+
const pageUrl = page.url();
1926

20-
let content = pageText;
21-
const MAX_TOKENS = 70000; // rough cap, assume ~4 chars per token for conservative truncation
22-
const estimatedTokens = Math.ceil(content.length / 4);
23-
if (estimatedTokens > MAX_TOKENS) {
24-
const maxChars = MAX_TOKENS * 4;
25-
content =
26-
content.substring(0, maxChars) +
27-
"\n\n[CONTENT TRUNCATED: Exceeded 70,000 token limit]";
27+
let content = pageText;
28+
const MAX_TOKENS = 70000; // rough cap, assume ~4 chars per token for conservative truncation
29+
const estimatedTokens = Math.ceil(content.length / 4);
30+
if (estimatedTokens > MAX_TOKENS) {
31+
const maxChars = MAX_TOKENS * 4;
32+
content =
33+
content.substring(0, maxChars) +
34+
"\n\n[CONTENT TRUNCATED: Exceeded 70,000 token limit]";
35+
}
36+
37+
return { success: true, content, pageUrl };
38+
} catch (error) {
39+
if (error instanceof TimeoutError) {
40+
const timeoutMessage = `TimeoutError: ariaTree() timed out — the page may be too large`;
41+
v3.logger({
42+
category: "agent",
43+
message: timeoutMessage,
44+
level: 0,
45+
});
46+
return {
47+
content: "",
48+
error: timeoutMessage,
49+
success: false,
50+
pageUrl: "",
51+
};
52+
}
53+
return {
54+
content: "",
55+
error: error?.message ?? String(error),
56+
success: false,
57+
pageUrl: "",
58+
};
59+
}
60+
},
61+
toModelOutput: (result) => {
62+
if (result.success === false || result.error !== undefined) {
63+
return {
64+
type: "content",
65+
value: [{ type: "text", text: JSON.stringify(result) }],
66+
};
2867
}
2968

30-
return { content, pageUrl };
69+
return {
70+
type: "content",
71+
value: [
72+
{ type: "text", text: `Accessibility Tree:\n${result.content}` },
73+
],
74+
};
3175
},
32-
toModelOutput: (result) => ({
33-
type: "content",
34-
value: [{ type: "text", text: `Accessibility Tree:\n${result.content}` }],
35-
}),
3676
});

‎packages/core/lib/v3/agent/tools/click.ts‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -87,37 +87,30 @@ export const clickTool = (v3: V3, provider?: string) =>
8787
}
8888
},
8989
toModelOutput: (result) => {
90-
if (result.success) {
91-
const content: ModelOutputContentItem[] = [
92-
{
93-
type: "text",
94-
text: JSON.stringify({
95-
success: result.success,
96-
describe: result.describe,
97-
coordinates: result.coordinates,
98-
}),
99-
},
100-
];
101-
if (result.screenshotBase64) {
102-
content.push({
103-
type: "media",
104-
mediaType: "image/png",
105-
data: result.screenshotBase64,
106-
});
107-
}
108-
return { type: "content", value: content };
90+
if (result.success === false || result.error !== undefined) {
91+
return {
92+
type: "content",
93+
value: [{ type: "text", text: JSON.stringify(result) }],
94+
};
10995
}
110-
return {
111-
type: "content",
112-
value: [
113-
{
114-
type: "text",
115-
text: JSON.stringify({
116-
success: result.success,
117-
error: result.error,
118-
}),
119-
},
120-
],
121-
};
96+
97+
const content: ModelOutputContentItem[] = [
98+
{
99+
type: "text",
100+
text: JSON.stringify({
101+
success: result.success,
102+
describe: result.describe,
103+
coordinates: result.coordinates,
104+
}),
105+
},
106+
];
107+
if (result.screenshotBase64) {
108+
content.push({
109+
type: "media",
110+
mediaType: "image/png",
111+
data: result.screenshotBase64,
112+
});
113+
}
114+
return { type: "content", value: content };
122115
},
123116
});

‎packages/core/lib/v3/agent/tools/dragAndDrop.ts‎

Lines changed: 23 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -102,36 +102,29 @@ export const dragAndDropTool = (v3: V3, provider?: string) =>
102102
}
103103
},
104104
toModelOutput: (result) => {
105-
if (result.success) {
106-
const content: ModelOutputContentItem[] = [
107-
{
108-
type: "text",
109-
text: JSON.stringify({
110-
success: result.success,
111-
describe: result.describe,
112-
}),
113-
},
114-
];
115-
if (result.screenshotBase64) {
116-
content.push({
117-
type: "media",
118-
mediaType: "image/png",
119-
data: result.screenshotBase64,
120-
});
121-
}
122-
return { type: "content", value: content };
105+
if (result.success === false || result.error !== undefined) {
106+
return {
107+
type: "content",
108+
value: [{ type: "text", text: JSON.stringify(result) }],
109+
};
123110
}
124-
return {
125-
type: "content",
126-
value: [
127-
{
128-
type: "text",
129-
text: JSON.stringify({
130-
success: result.success,
131-
error: result.error,
132-
}),
133-
},
134-
],
135-
};
111+
112+
const content: ModelOutputContentItem[] = [
113+
{
114+
type: "text",
115+
text: JSON.stringify({
116+
success: result.success,
117+
describe: result.describe,
118+
}),
119+
},
120+
];
121+
if (result.screenshotBase64) {
122+
content.push({
123+
type: "media",
124+
mediaType: "image/png",
125+
data: result.screenshotBase64,
126+
});
127+
}
128+
return { type: "content", value: content };
136129
},
137130
});

‎packages/core/lib/v3/agent/tools/extract.ts‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { tool } from "ai";
22
import { z, ZodTypeAny } from "zod";
33
import type { V3 } from "../../v3.js";
44
import type { AgentModelConfig } from "../../types/public/agent.js";
5+
import { TimeoutError } from "../../types/public/sdkErrors.js";
56

67
interface JsonSchema {
78
type?: string;
@@ -48,6 +49,7 @@ function jsonSchemaToZod(schema: JsonSchema): ZodTypeAny {
4849
export const extractTool = (
4950
v3: V3,
5051
executionModel?: string | AgentModelConfig,
52+
toolTimeout?: number,
5153
) =>
5254
tool({
5355
description: `Extract structured data from the current page based on a provided schema.
@@ -94,11 +96,23 @@ export const extractTool = (
9496
: undefined;
9597
const result = await v3.extract(instruction, parsedSchema, {
9698
...(executionModel ? { model: executionModel } : {}),
99+
timeout: toolTimeout,
97100
});
98101
return { success: true, result };
99102
} catch (error) {
100-
const err = error as Error;
101-
return { success: false, error: err?.message ?? String(error) };
103+
if (error instanceof TimeoutError) {
104+
const timeoutMessage = `TimeoutError: extract() timed out — try using a smaller or simpler schema`;
105+
v3.logger({
106+
category: "agent",
107+
message: timeoutMessage,
108+
level: 0,
109+
});
110+
return {
111+
success: false,
112+
error: timeoutMessage,
113+
};
114+
}
115+
return { success: false, error: error?.message ?? String(error) };
102116
}
103117
},
104118
});

0 commit comments

Comments
 (0)