AI Agent with controlled access to AWS services in Strands Agents

Strands Agents is a framework for building AI agents that comes with a long list of predefined tools. The tool use_aws allows the Agent to interact with AWS Services. In this post, I will explain how to restrict the use of this tool with an IAM role that is distinct from the IAM role of the agent itself.

Github repo

You can find the code in this post and a complete working example in this github repo. I am using AWS CDK to define the cloud resources in python. It should be relatively straightforward to translate the code to CloudFormation or Terraform if this is what you use.

The problem

When building agents in Strands Agents, we

a) first test our code locally
b) and then we deploy to Amazon Bedrock AgentCore.

In scenario a), we run the code with an IAM profile that might have the complete IAM permissions of our user (the developer). This might even provide administrator access to the AWS account.

In scenario b), we create an AgentCore Runtime and pass an IAM role. This role requires a set of permissions that allows AgentCore to host the agent for us, e.g. downloading the ECR image (ecr:BatchGetImage), invoking the Bedrock model (bedrock:InvokeModel) or putting logs in CloudWatch (logs:PutLogEvents).

In both scenarios, the IAM permissions needed to run our code is very different from the IAM permissions that we want to give to our Agent itself (permissions for calling AWS services). In this blog post, we assume that the agent will only help us understand our AWS Account, and for that reason it only needs read-only access.

Our goal is to run the predefined use_aws tool with a specific IAM role that has specific permissions (read-only in our example).

The IAM role

First define the IAM role. In AWS CDK syntax, this is

tool_use_aws_role = iam.Role(
    self,
    "tool-use-aws-role",
    role_name=f"{prefix}-tool-use-aws",
    assumed_by=iam.AccountRootPrincipal(),
    managed_policies=[iam.ManagedPolicy.from_aws_managed_policy_name("ReadOnlyAccess")],
)

We attach the managed AWS policy named ReadOnlyAccess. We also allow the whole account to assume the role, but you could restrict this role so that it can only be assumed by specific users and the role we will pass to AgentCore Runtime.

Passing the role

In Strands Agents, to use the use_aws tool we first have to import it

from strands_tools import use_aws

and then pass it to the agent

agent = Agent(
    model=bedrock_model,
    callback_handler=None,
    tools=[use_aws],
)

As we are not using a constructor when using the tool, we cannot pass any parameters. Notice that we just import the whole module.

We can have a closer look at the specification of the tool. This can be achieved in the IDE by looking at the source code of the following dictionary (or see the source code on github).

from strands_tools.use_aws import TOOL_SPEC

Among the specification of the tool, we find a parameter called profile_name. Not that this is a parameter for the Agent to use. The agent can choose or be instructed to use a specific profile. This is documentation mainly for the agent.

"profile_name": {
    "type": "string",
    "description": (
        "Optional: AWS profile name to use from ~/.aws/credentials. "
        "Defaults to default profile if not specified."
    ),
},

What we could do is to always use a predefined profile strands-playground-too-use-aws that will assume the role we defined earlier.

Defining a profile locally

Locally, in the development environment, we can easily define a profile by adding the following to ~/.aws/config

[profile strands-playground-tool-use-aws]
region=eu-west-1
role_arn = arn:aws:iam::<ACCOUNT ID>:role/strands-playground-tool-use-aws
source_profile = <YOUR MAIN PROFILE>

The role_arn will contain your own account id (I have hidden mine). The source profile will be a set of credentials that gives you access to the specific account and is allowed to assume the role.

Defining a profile on AgentCore

There are several options for deploying to Amazon Bedrock AgentCore Runtime. In order to define a profile we will have to go with the option of a Custom Agent with Dockerfile. In the Dockerfile, we need to do the following

ARG AWS_ACCOUNT_ID
ARG AWS_REGION_NAME
ARG AGENT_TOOL_USE_AWS_PROFILE_NAME="strands-playground-tool-use-aws"
ARG AGENT_TOOL_USE_AWS_ROLE_NAME="strands-playground-tool-use-aws"

RUN mkdir /root/.aws/
RUN echo "[profile ${AGENT_TOOL_USE_AWS_PROFILE_NAME}]"                                 >> /root/.aws/config
RUN echo "region=${AWS_REGION_NAME}"                                                    >> /root/.aws/config
RUN echo "role_arn=arn:aws:iam::${AWS_ACCOUNT_ID}:role/${AGENT_TOOL_USE_AWS_ROLE_NAME}" >> /root/.aws/config
# https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/security-credentials-management.html
RUN echo "credential_source=Ec2InstanceMetadata"                                        >> /root/.aws/config

Some observations:

• The complete Dockefile can be found here

• I am using a python container and the home directory is /root/. If you use a different base image, then make sure the profile is defined in <HOME DIRECTORY>/.aws/config.

• According to the documentation, the agent runs in a MicroVM that obtains credentials similarly to how EC2 instances do. Therefore we set credential_source=Ec2InstanceMetadata

• The Dockerfile arguments are passed in when we build the docker image. In CDK this is done like this:

•         docker_image = ecr_assets.DockerImageAsset(
              self,
              "agent-docker-image",
              directory=TOP_DIRECTORY,
              file="docker/images/agentcore/Dockerfile",
              build_args={
                  "AWS_ACCOUNT_ID": AWS_ACCOUNT_ID,
                  "AWS_REGION_NAME": AWS_REGION_NAME,
                  "AGENT_TOOL_USE_AWS_PROFILE_NAME": AGENT_TOOL_USE_AWS_PROFILE_NAME,
              },
          )
  

Modifying the tool behaviour

What we have achieved up to now is

• Create the IAM role

• Define a profile in both scenarios (locally and on AgentCore Runtime) that assumes the IAM role

What we are still missing is how to modify the behaviour of our tool so that it always uses the give profile.

Option 1: Tool Interception

The first approach (and recommended by the framework) is to use Hooks. This is a very clean feature provided by the framework that allows to extend agent functionality. Specifically, we need to do tool interception that is described here.

We define an interceptor like this

import os
from typing import Any

from strands.hooks import BeforeToolCallEvent, HookProvider, HookRegistry
from strands_tools.use_aws import TOOL_SPEC

TOOL_NAME = TOOL_SPEC["name"]

# The environment variable must be set. I let it fail if not.
PROFILE_NAME = os.environ["AGENT_TOOL_USE_AWS_PROFILE_NAME"]


class UseAwsInterceptor(HookProvider):
    def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
        registry.add_callback(BeforeToolCallEvent, self.intercept_tool)

    def intercept_tool(self, event: BeforeToolCallEvent) -> None:
        if event.tool_use["name"] == TOOL_NAME:
            event.tool_use["input"]["profile_name"] = PROFILE_NAME

Every time the tool use_aws is used, we intercept the call and modify the input just before the tool is called. We modify the profile_name parameter and we pass the name of the profile that we defined in the previous section. Note that I am injecting the profile name as an environment variable, this environment variable is set as part of my local environment and also on the AgentCore Runtime (be reminded that you can check the github repo for a complete working example).

Once the interceptor is defined, as a last step we need to add it to the hooks of the agent.

agent = Agent(
    model=bedrock_model,
    tools=[use_aws],
    hooks=[UseAwsInterceptor()],
)

The main disadvantage of this option (option 1), is that the Agent will still see the profile_name parameter in the tool specification. This means the agent might attempt to pass a profile name or even converse with the user about this parameter. The model will not be aware of the fact that the parameter takes no effect and is secretly replaced.

Option 2: Clone and modify the tool

As a second option, we will try to clone the tool and modify it with surgical precision. We want to achieve our goal with the minimal changes and without copy pasting functionality. If the use_aws tool evolves, we want our future code to have the latest and greatest without any changes.

Let’s go back to basics, and understand how tools are defined in Strands Agents. The use_aws tool is defined as a module based tool. This means that

• We create a module that contains a TOOL_SPEC dictionary

• The name of the tool is defined in TOOL_SPEC[“name”]

• A tool function is defined with the same name as the tool and a specific signature

With this knowledge, we can hack together our own tool called controlled_use_aws

import copy
import os
from typing import Any

from strands.types.tools import ToolResult, ToolUse
from strands_tools.use_aws import TOOL_SPEC as ORIGINAL_TOOL_SPEC
from strands_tools.use_aws import use_aws as original_use_aws

# Make a deep copy of the original TOOL_SPEC.
# Strands Agents expects this to be called TOOL_SPEC within the module
TOOL_SPEC = copy.deepcopy(ORIGINAL_TOOL_SPEC)

# Set a new name for our own implementation
TOOL_SPEC["name"] = "controlled_use_aws"

# Remove the profile_name parameter from the spec
del TOOL_SPEC["inputSchema"]["json"]["properties"]["profile_name"]  # type: ignore


# The environment variable must be set. I let it fail if not.
PROFILE_NAME = os.environ["AGENT_TOOL_USE_AWS_PROFILE_NAME"]


def controlled_use_aws(tool: ToolUse, **kwargs: Any) -> ToolResult:
    tool["input"]["profile_name"] = PROFILE_NAME

    return original_use_aws(tool, **kwargs)

That’s all we need. Some observations:

• I am making a deep copy of the TOOL_SPEC from the original tool, this is because I am about to mutate it and I don’t want to mutate the dictionary in my library.

• I am changing the tool name in the new spec to controlled_use_aws.

• I delete the property profile_name from the new specification. This way our agent will not know that the tool can accept a custom profile name.

• The tool function has to be named as the tool, i.e. controlled_use_aws

• The tool function is 2 lines of code: we add the profile_name to the input and we pass it to the original tool function of the tool use_aws

Now that we have defined our own tool, we need to add it to the agent. Remember that you need to import the module as the tool. This means that if the above code is in src/playground/agents/tools/controlled_use_aws.py then you need to import it like this

from playground.agents.tools import controlled_use_aws

and then add it to the agent’s list of tools

agent = Agent(
    model=bedrock_model,
    tools=[controlled_use_aws],
)

The advantage of option 2 is that the profile_name is not part of the specification any more. The obvious disadvantage is that it is a bit of a hack, but this is due to the lack of configurability of the underlying tool.

Conclusion

The use_aws tool in Strands Agents can be restricted to run under a dedicated IAM role. This allows us to to control the permissions of the agent when interacting with AWS services on our behalf. We presented two implementation options, one using Strands Agents hooks and one writing our own custom tool reusing the logic in the original tool. In both cases, we showed how it needs to be configured and run in a) the local development environment and b) in a deployed AgentCore Runtime.