fix(vesting): handle claim and clawback race states

backend/src/services/scheduleExecutor.ts

@@ -6,6 +6,7 @@ import { scheduleService } from './scheduleService.js';
 import type { Schedule, ExecutionResult, PaymentRecipient } from '../types/schedule.js';
 import { Operation, Asset, Memo, Keypair } from '@stellar/stellar-sdk';
 import os from 'node:os';
+import type { PoolClient } from 'pg';
 
 export class ScheduleExecutor {
   private cronJob: ScheduledTask | null = null;
@@ -117,7 +118,7 @@ export class ScheduleExecutor {
 
           const executionResult = await this.executeSchedule(schedule);
 
-          await this.recordExecution(schedule.id, executionResult);
+          await this.recordExecution(schedule.id, executionResult, client);
 
           if (executionResult.success) {
             successCount++;
@@ -143,7 +144,7 @@ export class ScheduleExecutor {
                 message: error instanceof Error ? error.message : 'System error in executor',
                 details: error as any,
               },
-            });
+            }, client);
           } catch (recordError) {
             console.error(
               `[ScheduleExecutor] Failed to record execution error for schedule ID ${scheduleRow.id}:`,
@@ -152,7 +153,7 @@ export class ScheduleExecutor {
           }
         } finally {
           // Always release the claim so the row is available for the next cycle
-          await this.releaseClaim(scheduleRow.id);
+          await this.releaseClaim(scheduleRow.id, client);
         }
       }
 
@@ -172,12 +173,21 @@ export class ScheduleExecutor {
   /**
    * Release the row-level claim after execution (success or failure).
    */
-  private async releaseClaim(scheduleId: number): Promise<void> {
+  private async releaseClaim(scheduleId: number, client: PoolClient): Promise<void> {
     try {
-      await pool.query(
-        'UPDATE schedules SET locked_by = NULL, locked_at = NULL WHERE id = $1',
-        [scheduleId]
+      const result = await client.query(
+        `UPDATE schedules
+         SET locked_by = NULL, locked_at = NULL
+         WHERE id = $1 AND locked_by = $2
+         RETURNING id`,
+        [scheduleId, this.podId]
       );
+
+      if (result.rowCount !== 1) {
+        console.warn(
+          `[ScheduleExecutor] Claim for schedule ID ${scheduleId} was not released because it is no longer owned by ${this.podId}`
+        );
+      }
     } catch (error) {
       console.error(`[ScheduleExecutor] Failed to release claim for schedule ID ${scheduleId}:`, error);
     }
@@ -298,15 +308,18 @@ export class ScheduleExecutor {
    * @param scheduleId - The schedule ID
    * @param result - The execution result
    */
-  async recordExecution(scheduleId: number, result: ExecutionResult): Promise<void> {
-    const client = await pool.connect();
+  async recordExecution(
+    scheduleId: number,
+    result: ExecutionResult,
+    client?: PoolClient,
+  ): Promise<void> {
+    const ownsClient = !client;
+    const dbClient = client ?? await pool.connect();
+
     try {
-      await client.query('BEGIN');
+      await dbClient.query('BEGIN');
 
-      // Determine execution status
       const status = result.success ? 'success' : 'failed';
-
-      // Insert into execution_history
       const insertQuery = `
         INSERT INTO execution_history (
           schedule_id,
@@ -323,31 +336,29 @@ export class ScheduleExecutor {
 
       const insertValues = [
         scheduleId,
-        new Date(), // executed_at
+        new Date(),
         status,
         result.transactionHash || null,
         result.success ? JSON.stringify({ hash: result.transactionHash }) : null,
         result.error?.message || null,
         result.error?.details ? JSON.stringify(result.error.details) : null,
       ];
 
-      await client.query(insertQuery, insertValues);
+      await dbClient.query(insertQuery, insertValues);
 
-      // Update schedule state using ScheduleService
-      await scheduleService.updateAfterExecution(scheduleId, result);
-
-      // Clear the lock now that execution is recorded
-      await client.query(
-        'UPDATE schedules SET locked_by = NULL, locked_at = NULL WHERE id = $1',
-        [scheduleId]
-      );
+      // Keep execution history and schedule-state mutation on the same transaction
+      // and physical connection that owns the schedule claim.
+      await scheduleService.updateAfterExecution(scheduleId, result, dbClient);
 
-      await client.query('COMMIT');
+      // Commit execution state before making the schedule claim available again.
+      await dbClient.query('COMMIT');
     } catch (error) {
-      await client.query('ROLLBACK');
+      await dbClient.query('ROLLBACK');
       throw error;
     } finally {
-      client.release();
+      if (ownsClient) {
+        dbClient.release();
+      }
     }
   }
 }

backend/src/services/scheduleService.ts

@@ -1,4 +1,5 @@
 import { DateTime } from 'luxon';
+import type { PoolClient } from 'pg';
 import { default as pool } from '../config/database.js';
 import type {
   Schedule,
@@ -361,10 +362,15 @@ export class ScheduleService {
   async updateAfterExecution(
     scheduleId: number,
     executionResult: ExecutionResult,
+    client?: PoolClient,
   ): Promise<void> {
-    const client = await pool.connect();
+    const ownsClient = !client;
+    const dbClient = client ?? await pool.connect();
+
     try {
-      await client.query('BEGIN');
+      if (ownsClient) {
+        await dbClient.query('BEGIN');
+      }
 
       // Query the schedule to get its frequency and configuration
       const selectQuery = `
@@ -379,7 +385,7 @@ export class ScheduleService {
         WHERE id = $1
       `;
 
-      const selectResult = await client.query(selectQuery, [scheduleId]);
+      const selectResult = await dbClient.query(selectQuery, [scheduleId]);
 
       if (selectResult.rows.length === 0) {
         throw new Error(`Schedule with ID ${scheduleId} not found`);
@@ -393,27 +399,20 @@ export class ScheduleService {
       let nextRunTimestamp: Date | null = null;
 
       if (!executionResult.success) {
-        // If execution failed, set status to 'failed'
         newStatus = 'failed';
+      } else if (schedule.frequency === 'once') {
+        newStatus = 'completed';
       } else {
-        // Execution succeeded
-        if (schedule.frequency === 'once') {
-          // For one-time schedules, set status to 'completed'
-          newStatus = 'completed';
-        } else {
-          // For recurring schedules, calculate new next_run_timestamp and keep status 'active'
-          newStatus = 'active';
-          nextRunTimestamp = this.calculateNextRun(
-            schedule.frequency,
-            schedule.timeOfDay,
-            new Date(schedule.startDate),
-            schedule.timezone,
-            executionTime, // Use execution time as lastRun
-          );
-        }
+        newStatus = 'active';
+        nextRunTimestamp = this.calculateNextRun(
+          schedule.frequency,
+          schedule.timeOfDay,
+          new Date(schedule.startDate),
+          schedule.timezone,
+          executionTime,
+        );
       }
 
-      // Update the schedule in the database
       const updateQuery = `
         UPDATE schedules
         SET 
@@ -424,19 +423,25 @@ export class ScheduleService {
         WHERE id = $4
       `;
 
-      await client.query(updateQuery, [
+      await dbClient.query(updateQuery, [
         executionTime,
         newStatus,
         nextRunTimestamp,
         scheduleId,
       ]);
 
-      await client.query('COMMIT');
+      if (ownsClient) {
+        await dbClient.query('COMMIT');
+      }
     } catch (error) {
-      await client.query('ROLLBACK');
+      if (ownsClient) {
+        await dbClient.query('ROLLBACK');
+      }
       throw error;
     } finally {
-      client.release();
+      if (ownsClient) {
+        dbClient.release();
+      }
     }
   }
 }

contracts/vesting_escrow/src/lib.rs

@@ -17,8 +17,10 @@ pub enum ContractError {
     Unauthorized          = 3,
     UpgradeAlreadyPending = 4,
     NoPendingUpgrade      = 5,
-    TimelockNotExpired    = 6,
-    TimestampOverflow     = 7,
+    TimelockNotExpired                 = 6,
+    TimestampOverflow                  = 7,
+    ClaimUnavailableAfterClawback      = 8,
+    ClawbackUnavailableAfterFullClaim  = 9,
 }
 
 impl From<CommonError> for ContractError {
@@ -150,54 +152,63 @@ impl VestingContract {
         client.transfer(&funder, &e.current_contract_address(), &amount);
     }
 
-    pub fn claim(e: Env) {
+    pub fn claim(e: Env) -> Result<(), ContractError> {
         let mut config: VestingConfig = e.storage().instance().get(&DataKey::Config).expect("Not initialized");
-        
+
         config.beneficiary.require_auth();
-        
+
         let vested = Self::calc_vested(&e, &config);
         let claimable = vested - config.claimed_amount;
 
         if claimable <= 0 {
-            // Nothing to claim, just return
-            return;
+            // Once clawback has capped the schedule, an exhausted grant is a
+            // terminal state rather than a silent no-op. This is what a claim
+            // observes when a full clawback wins the race.
+            if !config.is_active {
+                return Err(ContractError::ClaimUnavailableAfterClawback);
+            }
+            return Ok(());
         }
 
-        // Update state
         config.claimed_amount += claimable;
         e.storage().instance().set(&DataKey::Config, &config);
 
-        // Transfer tokens
         let client = token::Client::new(&e, &config.token);
         client.transfer(&e.current_contract_address(), &config.beneficiary, &claimable);
+        Ok(())
     }
-    
-    pub fn clawback(e: Env) {
+
+    pub fn clawback(e: Env) -> Result<(), ContractError> {
         let mut config: VestingConfig = e.storage().instance().get(&DataKey::Config).expect("Not initialized");
-        
+
         config.clawback_admin.require_auth();
-        
+
         if !config.is_active {
             panic!("Already revoked/inactive");
         }
 
-        // Calculate what has vested so far
+        // If a full claim wins the race there is nothing left for the admin to
+        // recover. Surface that terminal state explicitly instead of silently
+        // deactivating an already-settled grant.
+        if config.claimed_amount >= config.total_amount {
+            return Err(ContractError::ClawbackUnavailableAfterFullClaim);
+        }
+
         let vested = Self::calc_vested(&e, &config);
-
-        // The unvested amount is the total scheduled minus what has vested
         let unvested = config.total_amount - vested;
-        
-        // Update config to stop future vesting
-        // We set total_amount to vested, so effectively the grant is capped at what was vested at this moment
+
+        // Freeze vesting at the amount earned when clawback lands. Anything
+        // already vested but not yet claimed remains beneficiary-owned.
         config.total_amount = vested;
         config.is_active = false;
         e.storage().instance().set(&DataKey::Config, &config);
 
         if unvested > 0 {
-            // Return unvested tokens to admin
             let client = token::Client::new(&e, &config.token);
             client.transfer(&e.current_contract_address(), &config.clawback_admin, &unvested);
         }
+
+        Ok(())
     }
 
     pub fn get_vested_amount(e: Env) -> i128 {

contracts/vesting_escrow/src/test.rs

@@ -179,6 +179,71 @@ fn setup_upgradeable() -> UpgradeFixture {
     }
 }
 
+#[test]
+fn test_claim_after_full_clawback_returns_clear_error() {
+    let f = setup_upgradeable();
+
+    // Before the cliff, clawback is full because nothing has vested.
+    f.client.clawback();
+
+    assert_eq!(
+        f.client.try_claim(),
+        Err(Ok(ContractError::ClaimUnavailableAfterClawback))
+    );
+
+    let config = f.client.get_config();
+    assert_eq!(config.total_amount, 0);
+    assert_eq!(config.claimed_amount, 0);
+    assert!(!config.is_active);
+}
+
+#[test]
+fn test_clawback_after_full_claim_returns_clear_error() {
+    let f = setup_upgradeable();
+    let token_client = token::Client::new(&f.e, &f.token);
+
+    f.e.ledger().set_timestamp(f.start_time + 1_000);
+    f.client.claim();
+    assert_eq!(token_client.balance(&f.beneficiary), 10_000);
+
+    assert_eq!(
+        f.client.try_clawback(),
+        Err(Ok(ContractError::ClawbackUnavailableAfterFullClaim))
+    );
+
+    let config = f.client.get_config();
+    assert_eq!(config.claimed_amount, 10_000);
+    assert_eq!(config.total_amount, 10_000);
+    assert!(config.is_active);
+}
+
+#[test]
+fn test_partial_claim_before_clawback_preserves_remaining_vested_claim() {
+    let f = setup_upgradeable();
+    let token_client = token::Client::new(&f.e, &f.token);
+
+    f.e.ledger().set_timestamp(f.start_time + 200);
+    f.client.claim();
+    assert_eq!(token_client.balance(&f.beneficiary), 2_000);
+
+    f.e.ledger().set_timestamp(f.start_time + 500);
+    f.client.clawback();
+
+    let frozen = f.client.get_config();
+    assert_eq!(frozen.claimed_amount, 2_000);
+    assert_eq!(frozen.total_amount, 5_000);
+    assert!(!frozen.is_active);
+    assert_eq!(token_client.balance(&f.clawback_admin), 5_000);
+
+    // The 3_000 that vested before clawback remains claimable.
+    f.client.claim();
+    assert_eq!(token_client.balance(&f.beneficiary), 5_000);
+    assert_eq!(
+        f.client.try_claim(),
+        Err(Ok(ContractError::ClaimUnavailableAfterClawback))
+    );
+}
+
 fn new_wasm_hash(e: &Env) -> BytesN<32> {
     e.deployer()
         .upload_contract_wasm(Bytes::from_slice(e, upgraded_vesting::WASM))