Skip to content

Commit 7395436

Browse files
author
bgagent
committed
fix(browser): address PR review feedback — security, error handling, simplification
- Remove Gateway from AgentBrowser construct (unused; agent calls Lambda directly via boto3, avoids extra Cognito user pool) - Add URL allowlist in browser-tool Lambda: HTTPS-only, github.com only, reject RFC1918/link-local/IMDS to prevent SSRF - Use discriminated union type for BrowserToolResponse - Sanitize taskId in S3 key path to [a-zA-Z0-9_-] - Add persistent WebSocket error handler, wrap JSON.parse in try/catch, replace non-null assertions with descriptive error checks - Log session cleanup failures instead of silently swallowing - Increase presigned URL expiry to 7 days (bounded by credential lifetime) - Distinguish config errors (ERROR) from transient errors (WARN) in browser.py; check FunctionError on Lambda invoke response - Check gh pr edit return code; deduplicate ## Screenshots section on retry - Log exception details in pipeline.py screenshot catch block - Add 9 tests for capture_pr_screenshots and _append_screenshots_to_pr
1 parent 7d1c947 commit 7395436

10 files changed

Lines changed: 294 additions & 121 deletions

File tree

‎agent/src/browser.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,18 +31,33 @@ def capture_screenshot(url: str, task_id: str = "") -> str | None:
3131
return None
3232
try:
3333
client = _get_lambda_client()
34+
except ValueError as e:
35+
print(f"[browser] [ERROR] Configuration error: {e}", flush=True)
36+
return None
37+
try:
3438
payload = json.dumps({"action": "screenshot", "url": url, "taskId": task_id})
3539
response = client.invoke(
3640
FunctionName=function_name,
3741
InvocationType="RequestResponse",
3842
Payload=payload,
3943
)
44+
if "FunctionError" in response:
45+
error_payload = json.loads(response["Payload"].read())
46+
print(
47+
f"[browser] [ERROR] Lambda function crashed: "
48+
f"{response['FunctionError']} — {error_payload}",
49+
flush=True,
50+
)
51+
return None
4052
result = json.loads(response["Payload"].read())
4153
if result.get("status") == "success":
4254
print(f"[browser] Screenshot captured: {result.get('screenshotS3Key')}", flush=True)
4355
return result.get("presignedUrl")
4456
print(f"[browser] Screenshot failed: {result.get('error', 'unknown')}", flush=True)
4557
return None
4658
except Exception as e:
47-
print(f"[browser] [WARN] capture_screenshot failed: {type(e).__name__}: {e}", flush=True)
59+
print(
60+
f"[browser] [WARN] capture_screenshot failed (transient): {type(e).__name__}: {e}",
61+
flush=True,
62+
)
4863
return None

‎agent/src/pipeline.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -337,8 +337,11 @@ def run_task(
337337
screenshot_urls = capture_pr_screenshots(pr_url, config.task_id)
338338
if screenshot_urls:
339339
_append_screenshots_to_pr(config, setup, screenshot_urls)
340-
except Exception:
341-
log("WARN", "Screenshot capture failed (non-fatal)")
340+
except Exception as exc:
341+
log(
342+
"WARN",
343+
f"Screenshot capture failed (non-fatal): {type(exc).__name__}: {exc}",
344+
)
342345

343346
post_span.set_attribute("build.passed", build_passed)
344347
post_span.set_attribute("lint.passed", lint_passed)

‎agent/src/post_hooks.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -376,16 +376,33 @@ def _append_screenshots_to_pr(
376376
images_md = "\n".join(
377377
f"![Screenshot {i + 1}]({url})" for i, url in enumerate(screenshot_urls)
378378
)
379-
updated_body = f"{current_body}\n\n## Screenshots\n\n{images_md}"
379+
screenshots_section = f"## Screenshots\n\n{images_md}"
380+
381+
if re.search(r"## Screenshots", current_body):
382+
updated_body = re.sub(
383+
r"## Screenshots\n.*?(?=\n## |\Z)",
384+
screenshots_section,
385+
current_body,
386+
flags=re.DOTALL,
387+
)
388+
else:
389+
updated_body = f"{current_body}\n\n{screenshots_section}"
380390

381-
subprocess.run(
391+
edit_result = subprocess.run(
382392
["gh", "pr", "edit", setup.branch, "--repo", config.repo_url, "--body", updated_body],
383393
cwd=setup.repo_dir,
384394
capture_output=True,
385395
text=True,
386396
timeout=30,
387397
)
388-
log("POST", f"Appended {len(screenshot_urls)} screenshot(s) to PR body")
398+
if edit_result.returncode == 0:
399+
log("POST", f"Appended {len(screenshot_urls)} screenshot(s) to PR body")
400+
else:
401+
log(
402+
"WARN",
403+
f"gh pr edit failed (rc={edit_result.returncode}): "
404+
f"{edit_result.stderr.strip()[:200]}",
405+
)
389406
except Exception as e:
390407
log("WARN", f"Failed to append screenshots to PR: {type(e).__name__}: {e}")
391408

‎agent/tests/test_post_hooks.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
"""Unit tests for post_hooks screenshot functions."""
2+
3+
from unittest.mock import MagicMock, patch
4+
5+
from post_hooks import _append_screenshots_to_pr, capture_pr_screenshots
6+
7+
8+
class TestCapturePrScreenshots:
9+
def test_returns_urls_on_success(self):
10+
with patch("browser.capture_screenshot", return_value="https://s3/img.png"):
11+
result = capture_pr_screenshots("https://github.com/owner/repo/pull/1", "task-1")
12+
assert result == ["https://s3/img.png"]
13+
14+
def test_returns_empty_list_when_pr_url_empty(self):
15+
result = capture_pr_screenshots("", "task-1")
16+
assert result == []
17+
18+
def test_returns_empty_list_when_pr_url_not_github(self):
19+
result = capture_pr_screenshots("https://gitlab.com/owner/repo/pull/1", "task-1")
20+
assert result == []
21+
22+
def test_returns_empty_list_on_exception(self):
23+
with patch("browser.capture_screenshot", side_effect=RuntimeError("boom")):
24+
result = capture_pr_screenshots("https://github.com/owner/repo/pull/1", "task-1")
25+
assert result == []
26+
27+
28+
class TestAppendScreenshotsToPr:
29+
def _make_mocks(self):
30+
config = MagicMock()
31+
config.repo_url = "https://github.com/owner/repo"
32+
setup = MagicMock()
33+
setup.branch = "bgagent/task-1"
34+
setup.repo_dir = "/tmp/repo"
35+
return config, setup
36+
37+
def test_appends_screenshots_section(self):
38+
config, setup = self._make_mocks()
39+
view_result = MagicMock(returncode=0, stdout="## Summary\n\nSome PR body")
40+
edit_result = MagicMock(returncode=0, stderr="")
41+
with patch("post_hooks.subprocess.run", side_effect=[view_result, edit_result]) as mock_run:
42+
_append_screenshots_to_pr(config, setup, ["https://s3/img1.png"])
43+
edit_call = mock_run.call_args_list[1]
44+
body_arg = edit_call[0][0][edit_call[0][0].index("--body") + 1]
45+
assert "## Screenshots" in body_arg
46+
assert "![Screenshot 1](https://s3/img1.png)" in body_arg
47+
48+
def test_replaces_existing_screenshots_section(self):
49+
config, setup = self._make_mocks()
50+
existing_body = "## Summary\n\nBody\n\n## Screenshots\n\n![Screenshot 1](https://old.png)"
51+
view_result = MagicMock(returncode=0, stdout=existing_body)
52+
edit_result = MagicMock(returncode=0, stderr="")
53+
with patch("post_hooks.subprocess.run", side_effect=[view_result, edit_result]) as mock_run:
54+
_append_screenshots_to_pr(config, setup, ["https://s3/new.png"])
55+
edit_call = mock_run.call_args_list[1]
56+
body_arg = edit_call[0][0][edit_call[0][0].index("--body") + 1]
57+
assert "![Screenshot 1](https://s3/new.png)" in body_arg
58+
assert "https://old.png" not in body_arg
59+
# Should only have one ## Screenshots section
60+
assert body_arg.count("## Screenshots") == 1
61+
62+
def test_handles_gh_pr_view_failure(self):
63+
config, setup = self._make_mocks()
64+
view_result = MagicMock(returncode=1, stdout="", stderr="not found")
65+
with patch("post_hooks.subprocess.run", return_value=view_result):
66+
# Should not raise
67+
_append_screenshots_to_pr(config, setup, ["https://s3/img.png"])
68+
69+
def test_handles_gh_pr_edit_failure(self):
70+
config, setup = self._make_mocks()
71+
view_result = MagicMock(returncode=0, stdout="## Summary\n\nBody")
72+
edit_result = MagicMock(returncode=1, stderr="permission denied")
73+
with patch("post_hooks.subprocess.run", side_effect=[view_result, edit_result]):
74+
# Should not raise
75+
_append_screenshots_to_pr(config, setup, ["https://s3/img.png"])
76+
77+
def test_does_nothing_when_urls_empty(self):
78+
config, setup = self._make_mocks()
79+
with patch("post_hooks.subprocess.run") as mock_run:
80+
_append_screenshots_to_pr(config, setup, [])
81+
mock_run.assert_not_called()

‎cdk/src/constructs/agent-browser.ts‎

Lines changed: 0 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ export interface AgentBrowserProps {
3434

3535
export class AgentBrowser extends Construct {
3636
public readonly browser: agentcore.BrowserCustom;
37-
public readonly gateway: agentcore.Gateway;
3837
public readonly browserToolFn: lambda.NodejsFunction;
3938
public readonly screenshotBucket: s3.Bucket;
4039

@@ -108,79 +107,6 @@ export class AgentBrowser extends Construct {
108107
},
109108
], true);
110109

111-
// --- Gateway with Lambda target ---
112-
this.gateway = new agentcore.Gateway(this, 'Gateway');
113-
114-
const toolSchema = agentcore.ToolSchema.fromInline([
115-
{
116-
name: 'screenshot',
117-
description: 'Capture a screenshot of a web page at the given URL',
118-
inputSchema: {
119-
type: agentcore.SchemaDefinitionType.OBJECT,
120-
properties: {
121-
action: {
122-
type: agentcore.SchemaDefinitionType.STRING,
123-
description: 'The action to perform. Must be "screenshot".',
124-
},
125-
url: {
126-
type: agentcore.SchemaDefinitionType.STRING,
127-
description: 'The URL of the web page to capture',
128-
},
129-
taskId: {
130-
type: agentcore.SchemaDefinitionType.STRING,
131-
description: 'Optional task ID for organizing screenshots',
132-
},
133-
},
134-
required: ['action', 'url'],
135-
},
136-
outputSchema: {
137-
type: agentcore.SchemaDefinitionType.OBJECT,
138-
properties: {
139-
status: {
140-
type: agentcore.SchemaDefinitionType.STRING,
141-
description: 'Result status: success or error',
142-
},
143-
screenshotS3Key: {
144-
type: agentcore.SchemaDefinitionType.STRING,
145-
description: 'S3 key of the stored screenshot',
146-
},
147-
presignedUrl: {
148-
type: agentcore.SchemaDefinitionType.STRING,
149-
description: 'Presigned URL to download the screenshot',
150-
},
151-
error: {
152-
type: agentcore.SchemaDefinitionType.STRING,
153-
description: 'Error message if the action failed',
154-
},
155-
},
156-
},
157-
},
158-
]);
159-
160-
this.gateway.addLambdaTarget('BrowserToolTarget', {
161-
lambdaFunction: this.browserToolFn,
162-
toolSchema,
163-
});
164-
165-
NagSuppressions.addResourceSuppressions(this.gateway, [
166-
{
167-
id: 'AwsSolutions-IAM5',
168-
reason: 'Gateway execution role requires wildcard permissions — generated by CDK L2 construct',
169-
},
170-
{
171-
id: 'AwsSolutions-COG1',
172-
reason: 'Gateway default Cognito user pool uses M2M client credentials flow — password policy not applicable',
173-
},
174-
{
175-
id: 'AwsSolutions-COG2',
176-
reason: 'Gateway default Cognito user pool uses M2M client credentials flow — MFA not applicable',
177-
},
178-
{
179-
id: 'AwsSolutions-COG3',
180-
reason: 'Gateway default Cognito user pool uses M2M client credentials flow — advanced security not required',
181-
},
182-
], true);
183-
184110
NagSuppressions.addResourceSuppressions(this.browser, [
185111
{
186112
id: 'AwsSolutions-IAM5',

0 commit comments

Comments
 (0)