feat: Make ai usage counter polymorphic

Author: kofimokome

database/migrations/2026_06_16_201941_create_ai_usage_counters_table.php

@@ -12,11 +12,10 @@ public function up(): void
     {
         Schema::create('ai_usage_counters', function (Blueprint $table) {
             $table->id();
-            $table->unsignedBigInteger('user_id');
+            $table->morphs('owner');
             $table->timestamp('period_start');
             $table->integer('tokens_used');
 
-            $table->foreign('user_id')->references('id')->on('users');
             $table->timestamps();
         });
     }

src/Models/AiUsageCounter.php

@@ -4,6 +4,7 @@
 
 use Illuminate\Database\Eloquent\Model;
 use Illuminate\Database\Eloquent\Relations\BelongsTo;
+use Illuminate\Database\Eloquent\Relations\MorphTo;
 
 class AiUsageCounter extends Model
 {
@@ -20,7 +21,8 @@ class AiUsageCounter extends Model
      * @var array<int, string>
      */
     protected $fillable = [
-        'user_id',
+        'owner_id',
+        'owner_type',
         'period_start',
         'tokens_used',
     ];
@@ -37,8 +39,8 @@ class AiUsageCounter extends Model
     /**
      * Get the user that owns this usage counter.
      */
-    public function user(): BelongsTo
+    public function owner(): MorphTo
     {
-        return $this->belongsTo(\App\Models\User::class);
+        return $this->morphTo();
     }
 }

src/Support/CloudEntitlements.php

@@ -62,7 +62,8 @@ public function remaining(?Model $owner, string $meter): int|float
         }
 
         $periodStart = Carbon::now()->startOfMonth();
-        $used = AiUsageCounter::where('user_id', $owner->id)
+        $used = AiUsageCounter::where('owner_id', $owner->id)
+            ->where('owner_type', get_class($owner))
             ->where('period_start', $periodStart)
             ->value('tokens_used') ?? 0;
 
@@ -87,29 +88,19 @@ public function consume(?Model $owner, string $meter, int $amount): void
             return;
         }
 
-        // Extract message ID from debug backtrace to ensure idempotency
-        $messageId = $this->getCurrentChatMessageId();
-        if ($messageId === null) {
-            return;
-        }
-
-        $cacheKey = "cloud_ai_consumed_message_{$messageId}";
-        if (Cache::has($cacheKey)) {
-            return;
-        }
-        Cache::put($cacheKey, true, now()->addDays(30));
-
         $periodStart = Carbon::now()->startOfMonth();
 
-        $counter = AiUsageCounter::where('user_id', $owner->id)
+        $counter = AiUsageCounter::where('owner_id', $owner->id)
+            ->where('owner_type', get_class($owner))
             ->where('period_start', $periodStart)
             ->first();
 
         if ($counter) {
             $counter->increment('tokens_used', $amount);
         } else {
             AiUsageCounter::create([
-                'user_id' => $owner->id,
+                'owner_id' => $owner->id,
+                'owner_type' => get_class($owner),
                 'period_start' => $periodStart,
                 'tokens_used' => $amount,
             ]);
@@ -135,17 +126,4 @@ private function getPlanCode(Model $owner): string
 
         return 'free';
     }
-
-    /**
-     * Find the chat message ID from the backtrace to key the idempotent increment.
-     */
-    private function getCurrentChatMessageId(): ?int
-    {
-        foreach (debug_backtrace(DEBUG_BACKTRACE_PROVIDE_OBJECT) as $frame) {
-            if (isset($frame['object']) && $frame['object'] instanceof \App\Jobs\ProcessChatMessageJob) {
-                return $frame['object']->assistantMessage->id;
-            }
-        }
-        return null;
-    }
 }