diff --git a/src/context.sw b/src/context.sw
index da209a21e40..d146e10440f 100644
--- a/src/context.sw
+++ b/src/context.sw
@@ -6,16 +6,17 @@ use ::contract_id::ContractId;
 /// Get the current contract's id when called in an internal context.
 /// **Note !** If called in an external context, this will **not** return a contract ID.
 // @dev If called externally, will actually return a pointer to the transaction ID.
-pub fn contract_id() -> b256 {
-    asm() {
+pub fn contract_id() -> ContractId {
+    ~ContractId::from(asm() {
         fp: b256
-    }
+    })
 }
 
 /// Get the current contracts balance of coin `asset_id`
 pub fn this_balance(asset_id: b256) -> u64 {
-    asm(balance, token: asset_id) {
-        bal balance token fp;
+    let this_id = contract_id();
+    asm(balance, token: asset_id, contract_id: this_id) {
+        bal balance token contract_id;
         balance: u64
     }
 }
@@ -36,11 +37,11 @@ pub fn msg_amount() -> u64 {
 }
 
 /// Get the asset_id of coins being sent.
-pub fn msg_asset_id() -> b256 {
-    asm(asset_id) {
+pub fn msg_asset_id() -> ContractId {
+    ~ContractId::from(asm(asset_id) {
         addi asset_id fp i32;
         asset_id: b256
-    }
+    })
 }
 
 /// Get the remaining gas in the context.
diff --git a/tests/test_projects/token_ops/src/main.sw b/tests/test_projects/token_ops/src/main.sw
index 5abd3d20291..48bad07462a 100644
--- a/tests/test_projects/token_ops/src/main.sw
+++ b/tests/test_projects/token_ops/src/main.sw
@@ -1,6 +1,6 @@
 contract;
 
----------use std::{address::Address,contract_id::ContractId,token::*,context::balance_of_contract};
+use std::{address::Address, context::balance_of_contract, contract_id::ContractId, token::*};
 
 /// Parameters for `force_transfer` function.
 pub struct ParamsForceTransfer {