fix: [TRST-L-9] Cancel agreement if over-allocated

Author: matiasedgeandnode

packages/subgraph-service/contracts/SubgraphService.sol

@@ -574,10 +574,10 @@ contract SubgraphService is
      * @notice Internal function to handle closing an allocation
      * @dev This function is called when an allocation is closed, either by the indexer or by a third party
      * @param _allocationId The id of the allocation being closed
-     * @param _stale Whether the allocation is stale or not
+     * @param _forceClosed Whether the allocation was force closed
      */
-    function _onCloseAllocation(address _allocationId, bool _stale) internal {
-        IndexingAgreement._getStorageManager().onCloseAllocation(_allocationId, _stale);
+    function _onCloseAllocation(address _allocationId, bool _forceClosed) internal {
+        IndexingAgreement._getStorageManager().onCloseAllocation(_allocationId, _forceClosed);
     }
 
     /**
@@ -738,7 +738,20 @@ contract SubgraphService is
             _allocations.get(allocationId).indexer == _indexer,
             SubgraphServiceAllocationNotAuthorized(_indexer, allocationId)
         );
-        return _presentPOI(allocationId, poi_, poiMetadata_, _delegationRatio, paymentsDestination[_indexer]);
+
+        (uint256 paymentCollected, bool allocationForceClosed) = _presentPOI(
+            allocationId,
+            poi_,
+            poiMetadata_,
+            _delegationRatio,
+            paymentsDestination[_indexer]
+        );
+
+        if (allocationForceClosed) {
+            _onCloseAllocation(allocationId, true);
+        }
+
+        return paymentCollected;
     }
 
     /**

packages/subgraph-service/contracts/libraries/AllocationHandler.sol

@@ -281,14 +281,15 @@ library AllocationHandler {
      * @param allocationProvisionTracker The mapping of indexers to their locked tokens
      * @param _subgraphAllocatedTokens The mapping of subgraph deployment ids to their allocated tokens
      * @param params The parameters for the POI presentation
-     * @return The amount of tokens collected
+     * @return tokensCollected The amount of tokens collected
+     * @return allocationForceClosed True if the allocation was automatically closed due to over-allocation, false otherwise
      */
     function presentPOI(
         mapping(address allocationId => Allocation.State allocation) storage _allocations,
         mapping(address indexer => uint256 tokens) storage allocationProvisionTracker,
         mapping(bytes32 subgraphDeploymentId => uint256 tokens) storage _subgraphAllocatedTokens,
         PresentParams memory params
-    ) external returns (uint256) {
+    ) external returns (uint256, bool) {
         Allocation.State memory allocation = _allocations.get(params._allocationId);
         require(allocation.isOpen(), AllocationHandler.AllocationHandlerAllocationClosed(params._allocationId));
 
@@ -358,6 +359,7 @@ library AllocationHandler {
         );
 
         // Check if the indexer is over-allocated and force close the allocation if necessary
+        bool allocationForceClosed;
         if (
             _isOverAllocated(
                 allocationProvisionTracker,
@@ -366,6 +368,7 @@ library AllocationHandler {
                 params._delegationRatio
             )
         ) {
+            allocationForceClosed = true;
             _closeAllocation(
                 _allocations,
                 allocationProvisionTracker,
@@ -376,7 +379,7 @@ library AllocationHandler {
             );
         }
 
-        return tokensRewards;
+        return (tokensRewards, allocationForceClosed);
     }
 
     /**

packages/subgraph-service/contracts/libraries/IndexingAgreement.sol

@@ -450,10 +450,10 @@ library IndexingAgreement {
      *
      * @param self The indexing agreement storage manager
      * @param _allocationId The allocation ID
-     * @param stale Whether the allocation is stale or not
+     * @param forceClosed Whether the allocation was force closed
      *
      */
-    function onCloseAllocation(StorageManager storage self, address _allocationId, bool stale) external {
+    function onCloseAllocation(StorageManager storage self, address _allocationId, bool forceClosed) external {
         bytes16 agreementId = self.allocationToActiveAgreementId[_allocationId];
         if (agreementId == bytes16(0)) {
             return;
@@ -469,7 +469,7 @@ library IndexingAgreement {
             agreementId,
             wrapper.agreement,
             wrapper.collectorAgreement,
-            stale
+            forceClosed
                 ? IRecurringCollector.CancelAgreementBy.ThirdParty
                 : IRecurringCollector.CancelAgreementBy.ServiceProvider
         );

packages/subgraph-service/contracts/utilities/AllocationManager.sol

@@ -131,15 +131,16 @@ abstract contract AllocationManager is EIP712Upgradeable, GraphDirectory, Alloca
      * @param _poiMetadata The metadata associated with the POI. The data and encoding format is for off-chain components to define, this function will only emit the value in an event as-is.
      * @param _delegationRatio The delegation ratio to consider when locking tokens
      * @param _paymentsDestination The address where indexing rewards should be sent
-     * @return The amount of tokens collected
+     * @return tokensCollected The amount of tokens collected
+     * @return allocationForceClosed True if the allocation was automatically closed due to over-allocation, false otherwise
      */
     function _presentPOI(
         address _allocationId,
         bytes32 _poi,
         bytes memory _poiMetadata,
         uint32 _delegationRatio,
         address _paymentsDestination
-    ) internal returns (uint256) {
+    ) internal returns (uint256, bool) {
         return
             AllocationHandler.presentPOI(
                 _allocations,

packages/subgraph-service/test/unit/shared/HorizonStakingShared.t.sol

@@ -38,6 +38,12 @@ abstract contract HorizonStakingSharedTest is SubgraphBaseTest {
         staking.addToProvision(_indexer, address(subgraphService), _tokens);
     }
 
+    function _removeFromProvision(address _indexer, uint256 _tokens) internal {
+        staking.thaw(_indexer, address(subgraphService), _tokens);
+        skip(staking.getProvision(_indexer, address(subgraphService)).thawingPeriod + 1);
+        staking.deprovision(_indexer, address(subgraphService), 0);
+    }
+
     function _delegate(address _indexer, address _verifier, uint256 _tokens, uint256 _minSharesOut) internal {
         staking.delegate(_indexer, _verifier, _tokens, _minSharesOut);
     }

packages/subgraph-service/test/unit/subgraphService/indexing-agreement/integration.t.sol

@@ -101,6 +101,35 @@ contract SubgraphServiceIndexingAgreementIntegrationTest is SubgraphServiceIndex
         _sharedAssert(beforeCollect, afterCollect, expectedTokens, tokensCollected);
     }
 
+    function test_SubgraphService_CollectIndexingRewards_CancelsAgreementWhenOverAllocated_Integration(
+        Seed memory seed
+    ) public {
+        // Setup context and indexer with active agreement
+        Context storage ctx = _newCtx(seed);
+        IndexerState memory indexerState = _withIndexer(ctx);
+        (, bytes16 agreementId) = _withAcceptedIndexingAgreement(ctx, indexerState);
+
+        // Reduce indexer's provision to force over-allocation after collecting rewards
+        uint256 extraTokens = indexerState.tokens - minimumProvisionTokens;
+        vm.assume(extraTokens > 0);
+        _removeTokensFromProvision(indexerState, extraTokens);
+
+        // Verify indexer will be over-allocated after presenting POI
+        assertTrue(subgraphService.isOverAllocated(indexerState.addr));
+
+        // Collect indexing rewards - this should trigger allocation closure and agreement cancellation
+        bytes memory collectData = abi.encode(indexerState.allocationId, bytes32("poi"), bytes("metadata"));
+        resetPrank(indexerState.addr);
+        subgraphService.collect(indexerState.addr, IGraphPayments.PaymentTypes.IndexingRewards, collectData);
+
+        // Verify the indexing agreement was properly cancelled
+        IndexingAgreement.AgreementWrapper memory agreement = subgraphService.getIndexingAgreement(agreementId);
+        assertEq(
+            uint8(agreement.collectorAgreement.state),
+            uint8(IRecurringCollector.AgreementState.CanceledByServiceProvider)
+        );
+    }
+
     /* solhint-enable graph/func-name-mixedcase */
 
     function _sharedSetup(
@@ -195,10 +224,17 @@ contract SubgraphServiceIndexingAgreementIntegrationTest is SubgraphServiceIndex
         );
     }
 
-    function _addTokensToProvision(IndexerState memory _indexerState, uint256 _tokensToAddToProvision) private {
-        deal({ token: address(token), to: _indexerState.addr, give: _tokensToAddToProvision });
+    function _addTokensToProvision(IndexerState memory _indexerState, uint256 _tokens) private {
+        deal({ token: address(token), to: _indexerState.addr, give: _tokens });
+        vm.startPrank(_indexerState.addr);
+        _addToProvision(_indexerState.addr, _tokens);
+        vm.stopPrank();
+    }
+
+    function _removeTokensFromProvision(IndexerState memory _indexerState, uint256 _tokens) private {
+        deal({ token: address(token), to: _indexerState.addr, give: _tokens });
         vm.startPrank(_indexerState.addr);
-        _addToProvision(_indexerState.addr, _tokensToAddToProvision);
+        _removeFromProvision(_indexerState.addr, _tokens);
         vm.stopPrank();
     }