responding to PR comments

Author: Arjun Balaji

.github/workflows/terminal-bench.yaml

@@ -1,34 +1,32 @@
 # This is a terminal-bench workflow that is manually triggered
 # Template taken from https://github.com/actions/starter-workflows/blob/main/automation/manual.yml for reference 
 
-name: TB workflow
+name: Terminal-Bench
 
 # Controls when the action will run. Workflow runs when manually triggered using the UI
-# or API.
 on:
   workflow_dispatch:
-    # Inputs the workflow accepts.
     inputs:
       name:
-        # Friendly description to be shown in the UI instead of 'name'
-        description: 'Run terminal bench'
-        # Default value if no value is explicitly provided
+        description: 'Run terminal-bench workflow to test Q CLI in real terminal environments.' 
         default: 'all'
-        # Input has to be provided for the workflow to run
         required: true
-        # The data type of the input
         type: string
 
-# A workflow run is made up of one or more jobs that can run sequentially or in parallel
 jobs:
   run-benchmark:
+    # avoids disk storage issues
     runs-on: ubuntu-latest-8-cores
+    # makes these env vars available in main.py
     env:
-      FIGCHAT_GAMMA_ID: ${{ secrets.FIGCHAT_GAMMA_ID }}
+      CHAT_DOWNLOAD_ROLE_ARN: ${{ secrets.CHAT_DOWNLOAD_ROLE_ARN }}
+      CHAT_BUILD_BUCKET_NAME: ${{ secrets.CHAT_BUILD_BUCKET_NAME }}
     permissions:
       id-token: write
       contents: read
     steps:
+
+    # clear unnecessary storage to ensure docker containers have space
     - name: Cleanup and free disk space
       run: |
         sudo rm -rf /usr/share/dotnet
@@ -43,6 +41,7 @@ jobs:
     - name: Checkout repository
       uses: actions/checkout@v4
 
+    # Captures git hash of branch to query specific S3 bucket
     - name: Set git hash
       run: |
         if [ -n "$GITHUB_SHA" ]; then
@@ -64,18 +63,19 @@ jobs:
         python -m pip install --upgrade pip
         pip install terminal-bench
     
+    # OIDC enabled for github for ArjunPersonal
     - name: Configure AWS credentials
       uses: aws-actions/configure-aws-credentials@v4
       with:
-        role-to-assume: arn:aws:iam::${{ secrets.AWS_TB_ROLE }}:role/ArjunPersonal
+        role-to-assume: ${{ secrets.AWS_TB_ROLE }}
         aws-region: us-east-1
-       
 
     - name: Run terminal benchmark
       run: |
         cd terminal-bench-test
         tb run --agent-import-path main:AmazonQCLIAgent --dataset-name terminal-bench-core --dataset-version head
 
+    # uploads results if run fails as well to allow for easy log inspection
     - name: Upload results
       if: always()
       uses: actions/upload-artifact@v4

terminal-bench-test/main.py

@@ -14,22 +14,24 @@ class AmazonQCLIAgent(AbstractInstalledAgent):
     def name() -> str:
         return "Amazon Q CLI"
 
-    def __init__(self, model_name: str | None = None, *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._model_name = model_name
-        self._start_url = 'https://amzn.awsapps.com/start'
-        self.region = 'us-east-1'
 
+    """
+    Makes necessary env vars available in docker containers
+    """
     @property
     def _env(self) -> dict[str, str]:
         # SIGv4 = 1 for AWS credentials
-        env = {}
-        env["AMAZON_Q_SIGV4"] = 1
-        env["FIGCHAT_GAMMA_ID"] = os.environ.get("FIGCHAT_GAMMA_ID", '')
-        env["AWS_ACCESS_KEY_ID"] = os.environ.get("AWS_ACCESS_KEY_ID", '')
-        env["AWS_SECRET_ACCESS_KEY"] = os.environ.get("AWS_SECRET_ACCESS_KEY", '')
-        env["AWS_SESSION_TOKEN"] = os.environ.get("AWS_SESSION_TOKEN", '')
-        env["GIT_HASH"] = os.environ.get("GIT_HASH", '')
+        env = {
+            "AMAZON_Q_SIGV4": 1,
+            "AWS_ACCESS_KEY_ID": os.environ.get("AWS_ACCESS_KEY_ID", ''),
+            "AWS_SECRET_ACCESS_KEY": os.environ.get("AWS_SECRET_ACCESS_KEY", ''),
+            "AWS_SESSION_TOKEN": os.environ.get("AWS_SESSION_TOKEN", ''),
+            "GIT_HASH": os.environ.get("GIT_HASH", ''),
+            "CHAT_DOWNLOAD_ROLE_ARN": os.environ.get("CHAT_DOWNLOAD_ROLE_ARN", ''),
+            "CHAT_BUILD_BUCKET_NAME": os.environ.get("CHAT_BUILD_BUCKET_NAME", '')
+        }
         return env
 
     @property
@@ -40,11 +42,10 @@ def _run_agent_commands(self, task_description: str) -> list[TerminalCommand]:
         escaped_description = shlex.quote(task_description)
 
         return [
-        # q chat with 30 min max timeout and also we wait on input. Using qchat cuz sigv4. 
-        # non-interactive for now --> check if needed or not
+        # q chat with 30 min max timeout and also we wait on input. Using qchat because of sigv4. 
             TerminalCommand(
                 command=f"qchat chat --no-interactive --trust-all-tools {escaped_description}",
-                max_timeout_sec=370, 
+                max_timeout_sec=1800, 
                 block=True,
             )
         ]

terminal-bench-test/setup_amazon_q.sh

@@ -27,20 +27,19 @@ chmod 600 ~/.aws/config
 
 # Assume role and capture temporary credentials --> needed for s3 bucket access for build
 echo "Assuming AWS s3 role"
-TEMP_CREDENTIALS=$(aws sts assume-role --role-arn arn:aws:iam::${FIGCHAT_GAMMA_ID}:role/FigIoChat-S3Access-Role-Gamma --role-session-name S3AccessSession 2>/dev/null || echo '{}')
+TEMP_CREDENTIALS=$(aws sts assume-role --role-arn ${CHAT_DOWNLOAD_ROLE_ARN} --role-session-name S3AccessSession 2>/dev/null || echo '{}')
 QCHAT_ACCESSKEY=$(echo $TEMP_CREDENTIALS | jq -r '.Credentials.AccessKeyId')
 Q_SECRET_ACCESS_KEY=$(echo $TEMP_CREDENTIALS | jq -r '.Credentials.SecretAccessKey')
 Q_SESSION_TOKEN=$(echo $TEMP_CREDENTIALS | jq -r '.Credentials.SessionToken')
 
 # Download specific build from S3 based on commit hash
 echo "Downloading Amazon Q CLI build from S3..."
-S3_BUCKET="fig-io-chat-build-output-${FIGCHAT_GAMMA_ID}-us-east-1"
 S3_PREFIX="main/${GIT_HASH}/x86_64-unknown-linux-musl"
 echo "Downloading qchat.zip from s3://.../${S3_PREFIX}/qchat.zip"
 
 # Try download, if hash is invalid we fail.
 AWS_ACCESS_KEY_ID="$QCHAT_ACCESSKEY" AWS_SECRET_ACCESS_KEY="$Q_SECRET_ACCESS_KEY" AWS_SESSION_TOKEN="$Q_SESSION_TOKEN" \
-  aws s3 cp s3://${S3_BUCKET}/${S3_PREFIX}/qchat.zip ./qchat.zip --region us-east-1
+  aws s3 cp s3://${CHAT_BUILD_BUCKET_NAME}/${S3_PREFIX}/qchat.zip ./qchat.zip --region us-east-1
 
 # Handle the zip file, copy the qchat executable to /usr/local/bin + symlink from old code
 echo "Extracting qchat.zip..."