Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions mcp/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@
"neo4j-driver": "5.28.3",
"node-cache": "5.1.2",
"p-queue": "^6.6.2",
"pg": "8.13.1",
"redis": "^6.2.1",
"simple-git": "3.30.0",
"tsx": "4.21.0",
Expand All @@ -103,6 +104,7 @@
"@types/jsonwebtoken": "9.0.10",
"@types/node": "22.19.9",
"@types/node-cache": "4.2.5",
"@types/pg": "8.11.10",
"@types/xmldom": "0.1.34",
"clipboardy": "4.0.0",
"esbuild": "0.25.12",
Expand Down
30 changes: 27 additions & 3 deletions mcp/src/tools/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@ import { bearerToken } from "./utils.js";
import { Express } from "express";
import * as stakgraph from "./stakgraph/index.js";
import * as stagehand from "./stagehand/tools.js";
import { getMcpTools } from "./utils.js";
import * as verify from "./verify/index.js";
import { getMcpTools, use_stagehand } from "./utils.js";
import {
CallToolRequestSchema,
ListToolsRequestSchema,
Expand Down Expand Up @@ -36,7 +37,8 @@ export function graph_mcp_routes(app: Express) {
}

graphServer.setRequestHandler(ListToolsRequestSchema, async () => {
return { tools: getMcpTools() };
const tools = getMcpTools();
return { tools: use_stagehand() ? [...tools, ...verify.VERIFY_TOOLS] : tools };
});

graphServer.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
Expand All @@ -61,9 +63,31 @@ graphServer.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
case stakgraph.GetRulesFilesTool.name: {
return await stakgraph.getRulesFiles();
}
case verify.HttpRequestTool.name: {
const sid = extra.sessionId || "default-session-id";
return await verify.httpRequest(sid, args || {});
}
case verify.SampleTool.name: {
const sid = extra.sessionId || "default-session-id";
return await verify.sampleUrl(sid, args || {});
}
case verify.DbQueryTool.name: {
const sid = extra.sessionId || "default-session-id";
return await verify.dbQuery(sid, args || {});
}
case verify.RunCommandTool.name: {
const sid = extra.sessionId || "default-session-id";
return await verify.runCommand(sid, args || {});
}
case verify.SubmitVerdictTool.name: {
const sid = extra.sessionId || "default-session-id";
return await verify.submitVerdict(sid, args || {});
}
default:
if (name.startsWith("stagehand_")) {
return await stagehand.call(name, args || {}, extra.sessionId);
const sid = extra.sessionId || "default-session-id";
const result = await stagehand.call(name, args || {}, extra.sessionId);
return verify.tagEvidence(sid, name, result);
}
throw new Error(`Unknown tool: ${name}`);
}
Expand Down
137 changes: 111 additions & 26 deletions mcp/src/tools/stagehand/core.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { Stagehand, Page } from "@browserbasehq/stagehand";
import { getProvider } from "./providers.js";
import { Stagehand, Page, AISdkClient } from "@browserbasehq/stagehand";
import { resolveBrowserModel } from "./providers.js";
import { getModelDetails } from "../../aieo/src/provider.js";

let STATE: {
[sessionId: string]: {
Expand Down Expand Up @@ -56,19 +57,17 @@ export async function getOrCreateStagehand(sessionIdMaybe?: string) {
return STATE[sessionId].stagehand;
}

let provider = getProvider();
console.log("initializing stagehand!", provider.model);
const { model: modelName, apiKey } = resolveBrowserModel();
console.log("initializing stagehand!", modelName);
const { model } = getModelDetails(modelName, apiKey);
const sh = new Stagehand({
env: "LOCAL",
domSettleTimeout: 60000,
localBrowserLaunchOptions: {
headless: true,
viewport: { width: 1024, height: 768 },
},
model: {
modelName: provider.model,
apiKey: process.env[provider.api_key_env_var_name],
},
llmClient: new AISdkClient({ model: model as any }),
});
await sh.init();

Expand All @@ -80,25 +79,9 @@ export async function getOrCreateStagehand(sessionIdMaybe?: string) {
networkEntries: [],
};

// Set up console log listener on the active page
// Set up console + network capture on the active page via CDP
const page = getActivePage(sh);
page.on("console", (msg) => {
const location = msg.location();
addConsoleLog(sessionId, {
timestamp: new Date().toISOString(),
type: msg.type(),
text: msg.text(),
location: {
url: location.url || "",
lineNumber: location.lineNumber || 0,
columnNumber: location.columnNumber || 0,
},
});
});

// Note: Network monitoring via request/response events is not available
// in the new Stagehand V3 API. The Page class only exposes "console" events.
// Network monitoring would need to be implemented via CDP directly if needed.
attachCapture(page as any, sessionId);

// Check if we need to evict old sessions (LRU)
if (Object.keys(STATE).length > MAX_SESSIONS) {
Expand All @@ -112,6 +95,108 @@ export async function getOrCreateStagehand(sessionIdMaybe?: string) {
return sh;
}

const NET_RESOURCE: Record<string, string> = { XHR: "xhr", Fetch: "fetch", Document: "document" };
const PENDING: { [sessionId: string]: Map<string, { method: string; url: string; type: string }> } = {};

// Attach CDP-based console + network capture to a page's main session.
// Populates the session's console logs (incl. uncaught exceptions) and network entries.
function attachCapture(page: any, sessionId: string): void {
try {
page.on("console", (msg: any) => {
const location = msg.location?.() || {};
addConsoleLog(sessionId, {
timestamp: new Date().toISOString(),
type: msg.type?.() || "log",
text: msg.text?.() || "",
location: {
url: location.url || "",
lineNumber: location.lineNumber || 0,
columnNumber: location.columnNumber || 0,
},
});
});

const session = page?.mainFrame?.()?.session;
if (!session) return;
PENDING[sessionId] = new Map();

const pushLog = (type: string, text: string, url = "") =>
addConsoleLog(sessionId, {
timestamp: new Date().toISOString(),
type,
text,
location: { url, lineNumber: 0, columnNumber: 0 },
});

session.send("Runtime.enable").catch(() => {});
session.on("Runtime.exceptionThrown", (p: any) => {
const d = p?.exceptionDetails;
pushLog("error", d?.exception?.description ?? d?.text ?? "uncaught exception", d?.url ?? "");
});

session.send("Log.enable").catch(() => {});
session.on("Log.entryAdded", (p: any) => {
const e = p?.entry;
if (e && (e.level === "error" || e.level === "warning")) {
pushLog(e.level, e.text ?? "", e.url ?? "");
}
});

session.send("Network.enable").catch(() => {});
session.on("Network.requestWillBeSent", (p: any) => {
const type = p?.type ?? "";
if (p?.redirectResponse && NET_RESOURCE[type]) {
addNetworkEntry(sessionId, {
id: `${p.requestId}-r`,
timestamp: new Date().toISOString(),
type: "response",
method: p?.request?.method ?? "GET",
url: p?.redirectResponse?.url ?? p?.request?.url ?? "",
status: p?.redirectResponse?.status ?? 0,
resourceType: NET_RESOURCE[type],
});
}
PENDING[sessionId]?.set(p.requestId, {
method: p?.request?.method ?? "GET",
url: p?.request?.url ?? "",
type,
});
});
session.on("Network.responseReceived", (p: any) => {
const type = p?.type ?? PENDING[sessionId]?.get(p?.requestId)?.type ?? "";
const req = PENDING[sessionId]?.get(p?.requestId);
PENDING[sessionId]?.delete(p?.requestId);
if (!NET_RESOURCE[type]) return;
addNetworkEntry(sessionId, {
id: p?.requestId ?? "",
timestamp: new Date().toISOString(),
type: "response",
method: req?.method ?? "GET",
url: p?.response?.url ?? req?.url ?? "",
status: p?.response?.status ?? 0,
resourceType: NET_RESOURCE[type],
});
});
session.on("Network.loadingFailed", (p: any) => {
const req = PENDING[sessionId]?.get(p?.requestId);
const type = req?.type ?? p?.type ?? "";
PENDING[sessionId]?.delete(p?.requestId);
if (!NET_RESOURCE[type]) return;
addNetworkEntry(sessionId, {
id: p?.requestId ?? "",
timestamp: new Date().toISOString(),
type: "response",
method: req?.method ?? "",
url: req?.url ?? "(unknown)",
status: 0,
resourceType: NET_RESOURCE[type],
});
});
} catch {
/* a CDP quirk must never fail session creation */
}
}

export function addConsoleLog(sessionId: string, log: ConsoleLog): void {
// Add to global logs (backward compatibility)
if (!STATE[sessionId]) {
Expand Down
26 changes: 24 additions & 2 deletions mcp/src/tools/stagehand/providers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,18 @@ export interface ProviderData {
export const PROVIDER_MODELS: Record<Provider, ProviderData> = {
anthropic: {
name: "anthropic",
model: "claude-3-7-sonnet-latest",
// Frontier model, matched to aieo (MODELS.anthropic.sonnet = "claude-sonnet-5").
// aieo-format ("provider/model") so it resolves through getModelDetails.
model: "anthropic/claude-sonnet-5",
// Computer-use / agent model stays a native Stagehand string (used only by
// the stagehand_agent CUA path, which does not go through aieo).
computer_use_model: "claude-sonnet-4-20250514",
api_key_env_var_name: "ANTHROPIC_API_KEY",
},
openai: {
name: "openai",
model: "gpt-4o",
// Frontier model, matched to aieo (MODELS.openai.gpt = "gpt-5").
model: "openai/gpt-5",
computer_use_model: "computer-use-preview",
api_key_env_var_name: "OPENAI_API_KEY",
},
Expand All @@ -29,3 +34,20 @@ export function getProvider(arg?: "anthropic" | "openai"): ProviderData {
}
return provider;
}

function apiKeyForModel(model: string): string {
if (model.startsWith("openrouter/")) return process.env.OPENROUTER_API_KEY || "";
if (model.startsWith("openai/")) return process.env.OPENAI_API_KEY || "";
return process.env.ANTHROPIC_API_KEY || "";
}

/**
* Single source of truth for the browser's LLM. Returns an aieo-format model
* string + its key, resolved once and always driven through getModelDetails —
* no separate "Stagehand built-in vs aieo" branch. The default comes from the
* configured provider; swap it in one place (or wire a registry) here.
*/
export function resolveBrowserModel(): { model: string; apiKey: string } {
const model = getProvider().model;
return { model, apiKey: apiKeyForModel(model) };
}
Loading
Loading