Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 15 additions & 0 deletions src/bedrock_agentcore/tools/code_interpreter_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,19 @@
from bedrock_agentcore._utils.endpoints import get_control_plane_endpoint, get_data_plane_endpoint
from bedrock_agentcore._utils.user_agent import build_user_agent_suffix

from .config import Certificate

DEFAULT_IDENTIFIER = "aws.codeinterpreter.v1"
DEFAULT_TIMEOUT = 900


def _to_dict(obj):
"""Convert an object to a dict, calling to_dict() if available."""
if hasattr(obj, "to_dict"):
return obj.to_dict()
return obj


class CodeInterpreter:
"""Client for interacting with the AWS Code Interpreter sandbox service.

Expand Down Expand Up @@ -134,6 +143,7 @@ def create_code_interpreter(
execution_role_arn: str,
network_configuration: Optional[Dict] = None,
description: Optional[str] = None,
certificates: Optional[List[Union[Certificate, Dict[str, Any]]]] = None,
tags: Optional[Dict[str, str]] = None,
client_token: Optional[str] = None,
) -> Dict:
Expand All @@ -155,6 +165,8 @@ def create_code_interpreter(
}
}
description (Optional[str]): Description of the interpreter (1-4096 chars)
certificates (Optional[List[Union[Certificate, Dict]]]): Root CA certificates
from Secrets Manager for the code interpreter to trust.
tags (Optional[Dict[str, str]]): Tags for the interpreter
client_token (Optional[str]): Idempotency token

Expand Down Expand Up @@ -193,6 +205,9 @@ def create_code_interpreter(
if description:
request_params["description"] = description

if certificates:
request_params["certificates"] = [_to_dict(c) for c in certificates]

if tags:
request_params["tags"] = tags

Expand Down
124 changes: 124 additions & 0 deletions tests/bedrock_agentcore/tools/test_code_interpreter_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1500,3 +1500,127 @@ def test_execute_code_auto_starts_session(self, mock_boto3, mock_get_data_endpoi
# Assert
client.data_plane_client.start_code_interpreter_session.assert_called_once()
client.data_plane_client.invoke_code_interpreter.assert_called_once()

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
def test_create_code_interpreter_with_certificates(
self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint
):
# Arrange
mock_session = MagicMock()
mock_control_client = MagicMock()
mock_data_client = MagicMock()
mock_session.client.side_effect = [mock_control_client, mock_data_client]
mock_boto3.Session.return_value = mock_session
client = CodeInterpreter("us-west-2")

mock_response = {
"codeInterpreterArn": "arn:aws:bedrock-agentcore:us-west-2:123456789012:code-interpreter/test-interp",
"codeInterpreterId": "test-interp-123",
"createdAt": datetime.datetime.now(),
"status": "CREATING",
}
client.control_plane_client.create_code_interpreter.return_value = mock_response

certificates = [
{
"location": {
"secretsManager": {"secretArn": "arn:aws:secretsmanager:us-west-2:123456789012:secret:my-ca"}
}
}
]

# Act
result = client.create_code_interpreter(
name="test_interpreter_with_certs",
execution_role_arn="arn:aws:iam::123456789012:role/InterpreterRole",
certificates=certificates,
description="Interpreter with custom root CA",
)

# Assert
client.control_plane_client.create_code_interpreter.assert_called_once_with(
name="test_interpreter_with_certs",
executionRoleArn="arn:aws:iam::123456789012:role/InterpreterRole",
networkConfiguration={"networkMode": "PUBLIC"},
description="Interpreter with custom root CA",
certificates=certificates,
)
assert result == mock_response

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
def test_create_code_interpreter_with_certificate_dataclass(
self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint
):
# Arrange
mock_session = MagicMock()
mock_control_client = MagicMock()
mock_data_client = MagicMock()
mock_session.client.side_effect = [mock_control_client, mock_data_client]
mock_boto3.Session.return_value = mock_session
client = CodeInterpreter("us-west-2")

mock_response = {
"codeInterpreterArn": "arn:aws:bedrock-agentcore:us-west-2:123456789012:code-interpreter/test-interp",
"codeInterpreterId": "test-interp-123",
"createdAt": datetime.datetime.now(),
"status": "CREATING",
}
client.control_plane_client.create_code_interpreter.return_value = mock_response

# Use a mock Certificate dataclass with to_dict method
mock_cert = MagicMock()
mock_cert.to_dict.return_value = {
"location": {"secretsManager": {"secretArn": "arn:aws:secretsmanager:us-west-2:123456789012:secret:my-ca"}}
}

# Act
result = client.create_code_interpreter(
name="test_interpreter_with_certs",
execution_role_arn="arn:aws:iam::123456789012:role/InterpreterRole",
certificates=[mock_cert],
)

# Assert
mock_cert.to_dict.assert_called_once()
client.control_plane_client.create_code_interpreter.assert_called_once_with(
name="test_interpreter_with_certs",
executionRoleArn="arn:aws:iam::123456789012:role/InterpreterRole",
networkConfiguration={"networkMode": "PUBLIC"},
certificates=[mock_cert.to_dict.return_value],
)
assert result == mock_response

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
def test_create_code_interpreter_without_certificates(
self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint
):
"""Verify that omitting certificates does not include it in the API call (backward compatibility)."""
# Arrange
mock_session = MagicMock()
mock_control_client = MagicMock()
mock_data_client = MagicMock()
mock_session.client.side_effect = [mock_control_client, mock_data_client]
mock_boto3.Session.return_value = mock_session
client = CodeInterpreter("us-west-2")

mock_response = {
"codeInterpreterId": "test-interp-123",
"status": "CREATING",
}
client.control_plane_client.create_code_interpreter.return_value = mock_response

# Act
client.create_code_interpreter(
name="test_interpreter_no_certs",
execution_role_arn="arn:aws:iam::123456789012:role/InterpreterRole",
)

# Assert — certificates key should NOT be in the call
call_kwargs = client.control_plane_client.create_code_interpreter.call_args[1]
assert "certificates" not in call_kwargs
96 changes: 95 additions & 1 deletion tests_integ/tools/test_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,13 @@
To run: pytest tests_integ/tools/test_code.py -v
"""

from bedrock_agentcore.tools.code_interpreter_client import code_session
import os
import time

import boto3

from bedrock_agentcore._utils.endpoints import get_control_plane_endpoint
from bedrock_agentcore.tools.code_interpreter_client import CodeInterpreter, code_session

# Test 1: Basic code execution with system interpreter
print("Test 1: Basic code execution using execute_code()")
Expand Down Expand Up @@ -223,6 +229,94 @@
print(event["result"]["content"])
print("✅ Test 10 passed\n")

# Test 11: Create code interpreter with custom root CA certificate
print("Test 11: Code interpreter with custom root CA certificate")

REGION = "us-west-2"
secret_arn = os.environ.get("TEST_ROOT_CA_SECRET_ARN")
execution_role_arn = os.environ.get("TEST_EXECUTION_ROLE_ARN")

if secret_arn and execution_role_arn:
cp_client = boto3.client(
"bedrock-agentcore-control",
region_name=REGION,
endpoint_url=get_control_plane_endpoint(REGION),
)

# Create interpreter with root CA
import time

response = cp_client.create_code_interpreter(
name=f"integ_test_rootca_{int(time.time())}",
executionRoleArn=execution_role_arn,
networkConfiguration={"networkMode": "PUBLIC"},
certificates=[{"location": {"secretsManager": {"secretArn": secret_arn}}}],
description="Integration test: code interpreter with custom root CA",
)
interpreter_id = response["codeInterpreterId"]
print(f" Created interpreter: {interpreter_id}")

# Wait for ready
ci_client = CodeInterpreter(REGION)
for _ in range(40):
info = ci_client.get_code_interpreter(interpreter_id)
if info["status"] == "READY":
break
elif info["status"] == "CREATE_FAILED":
raise RuntimeError(f"Interpreter creation failed: {info.get('failureReason')}")
time.sleep(3)

print(f" Interpreter status: {info['status']}")

# Start session and test SSL connection to untrusted-root.badssl.com
ci_client.start(identifier=interpreter_id)
print(f" Session started: {ci_client.session_id}")

result = ci_client.invoke(
"executeCode",
{
"code": (
"import urllib.request\n"
"response = urllib.request.urlopen("
'"https://untrusted-root.badssl.com")\n'
'print(f"Status: {response.status}")'
),
"language": "python",
},
)

success = False
for event in result.get("stream", []):
if "result" in event:
structured = event["result"].get("structuredContent", {})
stdout = structured.get("stdout", "")
if "200" in stdout:
success = True
print(f" SSL connection succeeded: {stdout.strip()}")

ci_client.stop()

# Cleanup - delete interpreter
sessions = ci_client.list_sessions(interpreter_id=interpreter_id, status="READY")
for s in sessions.get("items", []):
try:
ci_client.data_plane_client.stop_code_interpreter_session(
codeInterpreterIdentifier=interpreter_id, sessionId=s["sessionId"]
)
except Exception:
pass
time.sleep(5)
cp_client.delete_code_interpreter(codeInterpreterId=interpreter_id)
print(f" Cleaned up interpreter: {interpreter_id}")

assert success, "SSL connection to untrusted-root.badssl.com should succeed with custom root CA"
print("✅ Test 11 passed\n")
else:
print(" ⏭️ Skipped (set TEST_ROOT_CA_SECRET_ARN and TEST_EXECUTION_ROLE_ARN env vars to enable)")
print(" Example:")
print(" export TEST_ROOT_CA_SECRET_ARN=arn:aws:secretsmanager:us-west-2:123456789012:secret:badssl-root-ca")
print(" export TEST_EXECUTION_ROLE_ARN=arn:aws:iam::123456789012:role/AgentCoreRole")
print("✅ Test 11 skipped\n")

print("=" * 50)
print("All integration tests passed! ✅")
Expand Down
Loading