Skip to content

Commit ec0ae01

Browse files
committed
saving screenshots in a file
1 parent c249c97 commit ec0ae01

File tree

2 files changed

+84
-16
lines changed

2 files changed

+84
-16
lines changed

lib/agent/AnthropicCUAClient.ts

+42-8
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import {
1313
} from "@/types/agent";
1414
import { AgentClient } from "./AgentClient";
1515
import { AgentScreenshotProviderError } from "@/types/stagehandErrors";
16+
import * as fs from "fs";
17+
import * as path from "path";
1618

1719
export type ResponseInputItem = AnthropicMessage | AnthropicToolResult;
1820

@@ -838,25 +840,57 @@ export class AnthropicCUAClient extends AgentClient {
838840
base64Image?: string;
839841
currentUrl?: string;
840842
}): Promise<string> {
843+
let imageData = "";
844+
841845
// Use provided options if available
842846
if (options?.base64Image) {
843-
return `data:image/png;base64,${options.base64Image}`;
847+
imageData = `data:image/png;base64,${options.base64Image}`;
844848
}
845-
846849
// Use the screenshot provider if available
847-
if (this.screenshotProvider) {
850+
else if (this.screenshotProvider) {
848851
try {
849852
const base64Image = await this.screenshotProvider();
850-
return `data:image/png;base64,${base64Image}`;
853+
imageData = `data:image/png;base64,${base64Image}`;
851854
} catch (error) {
852855
console.error("Error capturing screenshot:", error);
853856
throw error;
854857
}
858+
} else {
859+
throw new AgentScreenshotProviderError(
860+
"`screenshotProvider` has not been set. " +
861+
"Please call `setScreenshotProvider()` with a valid function that returns a base64-encoded image",
862+
);
863+
}
864+
865+
// Save the screenshot to file if we have valid image data
866+
if (imageData) {
867+
try {
868+
// Create screenshots directory if it doesn't exist
869+
const screenshotsDir = path.resolve("screenshots");
870+
if (!fs.existsSync(screenshotsDir)) {
871+
fs.mkdirSync(screenshotsDir, { recursive: true });
872+
}
873+
874+
// Generate filename with timestamp
875+
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
876+
const filename = path.join(
877+
screenshotsDir,
878+
`screenshot-${timestamp}.png`,
879+
);
880+
881+
// Extract base64 data without the data URL prefix
882+
const base64Data = imageData.replace(/^data:image\/png;base64,/, "");
883+
884+
// Write file
885+
fs.writeFileSync(filename, base64Data, "base64");
886+
console.log(`Screenshot saved to ${filename}`);
887+
} catch (saveError) {
888+
// Log error but don't affect the function's behavior
889+
console.error("Error saving screenshot to file:", saveError);
890+
// Intentionally not re-throwing the error to keep function working
891+
}
855892
}
856893

857-
throw new AgentScreenshotProviderError(
858-
"`screenshotProvider` has not been set. " +
859-
"Please call `setScreenshotProvider()` with a valid function that returns a base64-encoded image",
860-
);
894+
return imageData;
861895
}
862896
}

lib/agent/OpenAICUAClient.ts

+42-8
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import {
1212
} from "@/types/agent";
1313
import { AgentClient } from "./AgentClient";
1414
import { AgentScreenshotProviderError } from "@/types/stagehandErrors";
15+
import * as fs from "fs";
16+
import * as path from "path";
1517

1618
/**
1719
* Client for OpenAI's Computer Use Assistant API
@@ -558,25 +560,57 @@ export class OpenAICUAClient extends AgentClient {
558560
base64Image?: string;
559561
currentUrl?: string;
560562
}): Promise<string> {
563+
let imageData = "";
564+
561565
// Use provided options if available
562566
if (options?.base64Image) {
563-
return `data:image/png;base64,${options.base64Image}`;
567+
imageData = `data:image/png;base64,${options.base64Image}`;
564568
}
565-
566569
// Use the screenshot provider if available
567-
if (this.screenshotProvider) {
570+
else if (this.screenshotProvider) {
568571
try {
569572
const base64Image = await this.screenshotProvider();
570-
return `data:image/png;base64,${base64Image}`;
573+
imageData = `data:image/png;base64,${base64Image}`;
571574
} catch (error) {
572575
console.error("Error capturing screenshot:", error);
573576
throw error;
574577
}
578+
} else {
579+
throw new AgentScreenshotProviderError(
580+
"`screenshotProvider` has not been set. " +
581+
"Please call `setScreenshotProvider()` with a valid function that returns a base64-encoded image",
582+
);
575583
}
576584

577-
throw new AgentScreenshotProviderError(
578-
"`screenshotProvider` has not been set. " +
579-
"Please call `setScreenshotProvider()` with a valid function that returns a base64-encoded image",
580-
);
585+
// Save the screenshot to file if we have valid image data
586+
if (imageData) {
587+
try {
588+
// Create screenshots directory if it doesn't exist
589+
const screenshotsDir = path.resolve("screenshots");
590+
if (!fs.existsSync(screenshotsDir)) {
591+
fs.mkdirSync(screenshotsDir, { recursive: true });
592+
}
593+
594+
// Generate filename with timestamp
595+
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
596+
const filename = path.join(
597+
screenshotsDir,
598+
`screenshot-${timestamp}.png`,
599+
);
600+
601+
// Extract base64 data without the data URL prefix
602+
const base64Data = imageData.replace(/^data:image\/png;base64,/, "");
603+
604+
// Write file
605+
fs.writeFileSync(filename, base64Data, "base64");
606+
console.log(`Screenshot saved to ${filename}`);
607+
} catch (saveError) {
608+
// Log error but don't affect the function's behavior
609+
console.error("Error saving screenshot to file:", saveError);
610+
// Intentionally not re-throwing the error to keep function working
611+
}
612+
}
613+
614+
return imageData;
581615
}
582616
}

0 commit comments

Comments
 (0)