flashLoan protocol fee

src/PartyPool.sol

@@ -412,9 +412,25 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         IERC20 token = IERC20(tokenAddr);
         require(amount <= token.balanceOf(address(this)));
         (uint256 fee, ) = _computeFee(amount, FLASH_FEE_PPM);
+
+        // Compute protocol share of flash fee
+        if (PROTOCOL_FEE_PPM > 0 && fee > 0) {
+            uint256 protoShare = (fee * PROTOCOL_FEE_PPM) / 1_000_000; // floor
+            if (protoShare > 0) {
+                uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
+                protocolFeesOwed[tokenIndex] += protoShare;
+            }
+        }
+
         require(token.transfer(address(receiver), amount));
         require(receiver.onFlashLoan(msg.sender, address(token), amount, fee, data) == FLASH_CALLBACK_SUCCESS);
         require(token.transferFrom(address(receiver), address(this), amount + fee));
+
+        // Update cached balance for the borrowed token
+        uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
+        uint256 balAfter = token.balanceOf(address(this));
+        _recordCachedBalance(tokenIndex, balAfter);
+
         return true;
     }