diff --git a/models.py b/models.py
index c4293bc..4dd80b1 100644
--- a/models.py
+++ b/models.py
@@ -153,3 +153,15 @@ class CreateManualPaymentRequest(BaseModel):
 
     amount: int
     description: str
+
+
+class GeneratePaymentInvoice(BaseModel):
+    """Generate payment invoice request"""
+
+    amount: int
+
+
+class RecordPayment(BaseModel):
+    """Record a payment"""
+
+    payment_hash: str
diff --git a/static/js/index.js b/static/js/index.js
index 0265353..0396725 100644
--- a/static/js/index.js
+++ b/static/js/index.js
@@ -34,6 +34,7 @@ window.app = Vue.createApp({
         amount: null,
         paymentRequest: null,
         paymentHash: null,
+        checkWalletKey: null,
         loading: false
       },
       settingsDialog: {
@@ -326,6 +327,7 @@ window.app = Vue.createApp({
         // Show the payment request in the dialog
         this.payDialog.paymentRequest = response.data.payment_request
         this.payDialog.paymentHash = response.data.payment_hash
+        this.payDialog.checkWalletKey = response.data.check_wallet_key
 
         this.$q.notify({
           type: 'positive',
@@ -334,21 +336,21 @@ window.app = Vue.createApp({
         })
 
         // Poll for payment completion
-        this.pollForPayment(response.data.payment_hash)
+        this.pollForPayment(response.data.payment_hash, response.data.check_wallet_key)
       } catch (error) {
         LNbits.utils.notifyApiError(error)
       } finally {
         this.payDialog.loading = false
       }
     },
-    async pollForPayment(paymentHash) {
+    async pollForPayment(paymentHash, checkWalletKey) {
       // Poll every 2 seconds for payment status
       const checkPayment = async () => {
         try {
           const response = await LNbits.api.request(
             'GET',
             `/api/v1/payments/${paymentHash}`,
-            this.g.user.wallets[0].inkey
+            checkWalletKey
           )
 
           if (response.data && response.data.paid) {
@@ -415,6 +417,7 @@ window.app = Vue.createApp({
       this.payDialog.amount = Math.abs(this.balance.balance)
       this.payDialog.paymentRequest = null
       this.payDialog.paymentHash = null
+      this.payDialog.checkWalletKey = null
       this.payDialog.show = true
     },
     async showReceivableDialog() {
diff --git a/views_api.py b/views_api.py
index 4991374..0316958 100644
--- a/views_api.py
+++ b/views_api.py
@@ -34,8 +34,10 @@ from .models import (
     CreateEntryLine,
     CreateJournalEntry,
     ExpenseEntry,
+    GeneratePaymentInvoice,
     JournalEntry,
     ReceivableEntry,
+    RecordPayment,
     RevenueEntry,
     UserBalance,
     UserWalletSettings,
@@ -446,13 +448,14 @@ async def api_get_all_balances(
 
 @castle_api_router.post("/api/v1/generate-payment-invoice")
 async def api_generate_payment_invoice(
-    amount: int,
+    data: GeneratePaymentInvoice,
     wallet: WalletTypeInfo = Depends(require_invoice_key),
 ) -> dict:
     """
     Generate an invoice on the Castle wallet for user to pay their balance.
     User can then pay this invoice to settle their debt.
     """
+    from lnbits.core.crud.wallets import get_wallet
     from lnbits.core.models import CreateInvoice
     from lnbits.core.services import create_payment_request
 
@@ -462,7 +465,7 @@ async def api_generate_payment_invoice(
     # Create invoice on castle wallet
     invoice_data = CreateInvoice(
         out=False,
-        amount=amount,
+        amount=data.amount,
         memo=f"Payment from user {wallet.wallet.user[:8]} to Castle",
         unit="sat",
         extra={"user_id": wallet.wallet.user, "type": "castle_payment"},
@@ -470,17 +473,25 @@ async def api_generate_payment_invoice(
 
     payment = await create_payment_request(castle_wallet_id, invoice_data)
 
+    # Get castle wallet to return its inkey for payment checking
+    castle_wallet = await get_wallet(castle_wallet_id)
+    if not castle_wallet:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND, detail="Castle wallet not found"
+        )
+
     return {
         "payment_hash": payment.payment_hash,
         "payment_request": payment.bolt11,
-        "amount": amount,
+        "amount": data.amount,
         "memo": invoice_data.memo,
+        "check_wallet_key": castle_wallet.inkey,  # Key to check payment status
     }
 
 
 @castle_api_router.post("/api/v1/record-payment")
 async def api_record_payment(
-    payment_hash: str,
+    data: RecordPayment,
     wallet: WalletTypeInfo = Depends(require_invoice_key),
 ) -> dict:
     """
@@ -490,7 +501,7 @@ async def api_record_payment(
     from lnbits.core.crud.payments import get_standalone_payment
 
     # Get the payment details
-    payment = await get_standalone_payment(payment_hash)
+    payment = await get_standalone_payment(data.payment_hash)
     if not payment:
         raise HTTPException(
             status_code=HTTPStatus.NOT_FOUND, detail="Payment not found"
@@ -518,7 +529,7 @@ async def api_record_payment(
     # This reduces what the user owes
     entry_data = CreateJournalEntry(
         description=f"Lightning payment from user {wallet.wallet.user[:8]}",
-        reference=payment_hash,
+        reference=data.payment_hash,
         lines=[
             CreateEntryLine(
                 account_id=lightning_account.id,