diff --git a/README.md b/README.md
index 8200720..ce47b1a 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@
 
 <br/>
 
-Enables quick switching between AWS profiles. Supports both hard-coded keys and role assumption.
+Enables quick switching between AWS profiles. Supports both hard-coded keys and role assumption. Works with the [aws](https://github.com/oh-my-fish/plugin-aws) plugin by [@sagebind](https://github.com/sagebind).
 
 
 ## Prerequisites
diff --git a/functions/asp.fish b/functions/asp.fish
index e112008..afe42d1 100644
--- a/functions/asp.fish
+++ b/functions/asp.fish
@@ -1,23 +1,35 @@
 function asp -d 'Switches AWS profile' -a 'aws_profile'
   if test -n "$aws_profile"
     if fgrep -q "[profile $aws_profile]" ~/.aws/config
-      set -l region (awk "/\[profile\ $aws_profile\]/,/^\$/ { if (\$1 == \"region\") { print \$3 }}" ~/.aws/config)
+      set -l region \
+        (awk "/\[profile\ $aws_profile\]/,/^\$/ { if (\$1 == \"region\") { print \$3 }}" \
+          ~/.aws/config)
 
       if test -z "$region"
         echo "No region listed for $aws_profile profile"
         return 1
       end
 
-      set -l access_key (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"aws_access_key_id\") { print \$3 }}" ~/.aws/credentials)
-      set -l secret_key (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"aws_secret_access_key\") { print \$3 }}" ~/.aws/credentials)
+      set -l access_key \
+        (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"aws_access_key_id\") { print \$3 }}" \
+          ~/.aws/credentials)
+      set -l secret_key \
+        (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"aws_secret_access_key\") { print \$3 }}" \
+          ~/.aws/credentials)
       set -l session_token ""
 
       if test -z "$access_key" -o -z "$secret_key"
-        set -l role_arn (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"role_arn\") { print \$3 }}" ~/.aws/credentials)
-        set -l source_profile (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"source_profile\") { print \$3 }}" ~/.aws/credentials)
+        set -l role_arn \
+          (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"role_arn\") { print \$3 }}" \
+            ~/.aws/credentials)
+        set -l source_profile \
+          (awk "/\[$aws_profile\]/,/^\$/ { if (\$1 == \"source_profile\") { print \$3 }}" \
+            ~/.aws/credentials)
 
         if test -n "$role_arn" -a -n "$source_profile"
-          set -l json (aws sts assume-role --profile "$source_profile" --role-arn "$role_arn" --role-session-name "$aws_profile" --output json)
+          set -l json \
+            (aws sts assume-role --profile "$source_profile" --role-arn "$role_arn" \
+              --role-session-name "$aws_profile" --output json)
 
           set access_key (echo $json | jq -r '.Credentials.AccessKeyId')
           set secret_key (echo $json | jq -r '.Credentials.SecretAccessKey')
@@ -34,6 +46,7 @@ function asp -d 'Switches AWS profile' -a 'aws_profile'
       set -gx AWS_SECURITY_TOKEN "$AWS_SESSION_TOKEN"
       set -gx AWS_DEFAULT_REGION "$region"
       set -gx AWS_DEFAULT_PROFILE "$aws_profile"
+      set -gx aws_profile "$aws_profile"
     else
       echo "No $aws_profile profile found in ~/.aws/config"
       return 1