working on rate limit / premium impl

Author: sawka

.roo/rules/rules.md

@@ -39,6 +39,7 @@ It has a TypeScript/React frontend and a Go backend. They talk together over `ws
     - NEVER use cursor-help (it looks terrible)
     - useAtom() and useAtomValue() are react HOOKS, so they must be called at the component level not inline in JSX
     - If you use React.memo(), make sure to add a displayName for the component
+- In general, when writing functions, we prefer _early returns_ rather than putting the majority of a function inside of an if block.
 
 ### Styling
 

aiprompts/wps-events.md

@@ -0,0 +1,293 @@
+# WPS Events Guide
+
+## Overview
+
+WPS (Wave PubSub) is Wave Terminal's publish-subscribe event system that enables different parts of the application to communicate asynchronously. The system uses a broker pattern to route events from publishers to subscribers based on event types and scopes.
+
+## Key Files
+
+- [`pkg/wps/wpstypes.go`](../pkg/wps/wpstypes.go) - Event type constants and data structures
+- [`pkg/wps/wps.go`](../pkg/wps/wps.go) - Broker implementation and core logic
+- [`pkg/wcore/wcore.go`](../pkg/wcore/wcore.go) - Example usage patterns
+
+## Event Structure
+
+Events in WPS have the following structure:
+
+```go
+type WaveEvent struct {
+    Event   string   `json:"event"`      // Event type constant
+    Scopes  []string `json:"scopes,omitempty"` // Optional scopes for targeted delivery
+    Sender  string   `json:"sender,omitempty"` // Optional sender identifier
+    Persist int      `json:"persist,omitempty"` // Number of events to persist in history
+    Data    any      `json:"data,omitempty"`    // Event payload
+}
+```
+
+## Adding a New Event Type
+
+### Step 1: Define the Event Constant
+
+Add your event type constant to [`pkg/wps/wpstypes.go`](../pkg/wps/wpstypes.go:8-19):
+
+```go
+const (
+    Event_BlockClose       = "blockclose"
+    Event_ConnChange       = "connchange"
+    // ... other events ...
+    Event_YourNewEvent     = "your:newevent"  // Use colon notation for namespacing
+)
+```
+
+**Naming Convention:**
+- Use descriptive PascalCase for the constant name with `Event_` prefix
+- Use lowercase with colons for the string value (e.g., "namespace:eventname")
+- Group related events with the same namespace prefix
+
+### Step 2: Define Event Data Structure (Optional)
+
+If your event carries structured data, define a type for it:
+
+```go
+type YourEventData struct {
+    Field1 string `json:"field1"`
+    Field2 int    `json:"field2"`
+}
+```
+
+### Step 3: Expose Type to Frontend (If Needed)
+
+If your event data type isn't already exposed via an RPC call, you need to add it to [`pkg/tsgen/tsgen.go`](../pkg/tsgen/tsgen.go:29-56) so TypeScript types are generated:
+
+```go
+// add extra types to generate here
+var ExtraTypes = []any{
+    waveobj.ORef{},
+    // ... other types ...
+    uctypes.RateLimitInfo{},  // Example: already added
+    YourEventData{},          // Add your new type here
+}
+```
+
+Then run code generation:
+
+```bash
+task generate
+```
+
+This will update [`frontend/types/gotypes.d.ts`](../frontend/types/gotypes.d.ts) with TypeScript definitions for your type, ensuring type safety in the frontend when handling these events.
+
+## Publishing Events
+
+### Basic Publishing
+
+To publish an event, use the global broker:
+
+```go
+import "github.com/wavetermdev/waveterm/pkg/wps"
+
+wps.Broker.Publish(wps.WaveEvent{
+    Event: wps.Event_YourNewEvent,
+    Data:  yourData,
+})
+```
+
+### Publishing with Scopes
+
+Scopes allow targeted event delivery. Subscribers can filter events by scope:
+
+```go
+wps.Broker.Publish(wps.WaveEvent{
+    Event:  wps.Event_WaveObjUpdate,
+    Scopes: []string{oref.String()},  // Target specific object
+    Data:   updateData,
+})
+```
+
+### Publishing in a Goroutine
+
+To avoid blocking the caller, publish events asynchronously:
+
+```go
+go func() {
+    wps.Broker.Publish(wps.WaveEvent{
+        Event: wps.Event_YourNewEvent,
+        Data:  data,
+    })
+}()
+```
+
+**When to use goroutines:**
+- When publishing from performance-critical code paths
+- When the event is informational and doesn't need immediate delivery
+- When publishing from code that holds locks (to prevent deadlocks)
+
+### Event Persistence
+
+Events can be persisted in memory for late subscribers:
+
+```go
+wps.Broker.Publish(wps.WaveEvent{
+    Event:   wps.Event_YourNewEvent,
+    Persist: 100,  // Keep last 100 events
+    Data:    data,
+})
+```
+
+## Complete Example: Rate Limit Updates
+
+This example shows how rate limit information is published when AI chat responses include rate limit headers.
+
+### 1. Define the Event Type
+
+In [`pkg/wps/wpstypes.go`](../pkg/wps/wpstypes.go:19):
+
+```go
+const (
+    // ... other events ...
+    Event_WaveAIRateLimit  = "waveai:ratelimit"
+)
+```
+
+### 2. Publish the Event
+
+In [`pkg/aiusechat/usechat.go`](../pkg/aiusechat/usechat.go:94-108):
+
+```go
+import "github.com/wavetermdev/waveterm/pkg/wps"
+
+func updateRateLimit(info *uctypes.RateLimitInfo) {
+    if info == nil {
+        return
+    }
+    rateLimitLock.Lock()
+    defer rateLimitLock.Unlock()
+    globalRateLimitInfo = info
+    
+    // Publish event in goroutine to avoid blocking
+    go func() {
+        wps.Broker.Publish(wps.WaveEvent{
+            Event: wps.Event_WaveAIRateLimit,
+            Data:  info,  // RateLimitInfo struct
+        })
+    }()
+}
+```
+
+### 3. Subscribe to the Event (Frontend)
+
+In the frontend, subscribe to events via WebSocket:
+
+```typescript
+// Subscribe to rate limit updates
+const subscription = {
+    event: "waveai:ratelimit",
+    allscopes: true  // Receive all rate limit events
+};
+```
+
+## Subscribing to Events
+
+### From Go Code
+
+```go
+// Subscribe to all events of a type
+wps.Broker.Subscribe(routeId, wps.SubscriptionRequest{
+    Event:     wps.Event_YourNewEvent,
+    AllScopes: true,
+})
+
+// Subscribe to specific scopes
+wps.Broker.Subscribe(routeId, wps.SubscriptionRequest{
+    Event:  wps.Event_WaveObjUpdate,
+    Scopes: []string{"workspace:123"},
+})
+
+// Unsubscribe
+wps.Broker.Unsubscribe(routeId, wps.Event_YourNewEvent)
+```
+
+### Scope Matching
+
+Scopes support wildcard matching:
+- `*` matches a single scope segment
+- `**` matches multiple scope segments
+
+```go
+// Subscribe to all workspace events
+wps.Broker.Subscribe(routeId, wps.SubscriptionRequest{
+    Event:  wps.Event_WaveObjUpdate,
+    Scopes: []string{"workspace:*"},
+})
+```
+
+## Best Practices
+
+1. **Use Namespaces**: Prefix event names with a namespace (e.g., `waveai:`, `workspace:`, `block:`)
+
+2. **Don't Block**: Use goroutines when publishing from performance-critical code or while holding locks
+
+3. **Type-Safe Data**: Define struct types for event data rather than using maps
+
+4. **Scope Wisely**: Use scopes to limit event delivery and reduce unnecessary processing
+
+5. **Document Events**: Add comments explaining when events are fired and what data they carry
+
+6. **Consider Persistence**: Use `Persist` for events that late subscribers might need (like status updates)
+
+## Common Event Patterns
+
+### Status Updates
+
+```go
+wps.Broker.Publish(wps.WaveEvent{
+    Event:   wps.Event_ControllerStatus,
+    Scopes:  []string{blockId},
+    Persist: 1,  // Keep only latest status
+    Data:    statusData,
+})
+```
+
+### Object Updates
+
+```go
+wps.Broker.Publish(wps.WaveEvent{
+    Event:  wps.Event_WaveObjUpdate,
+    Scopes: []string{oref.String()},
+    Data: waveobj.WaveObjUpdate{
+        UpdateType: waveobj.UpdateType_Update,
+        OType:      obj.GetOType(),
+        OID:        waveobj.GetOID(obj),
+        Obj:        obj,
+    },
+})
+```
+
+### Batch Updates
+
+```go
+// Helper function for multiple updates
+func (b *BrokerType) SendUpdateEvents(updates waveobj.UpdatesRtnType) {
+    for _, update := range updates {
+        b.Publish(WaveEvent{
+            Event:  Event_WaveObjUpdate,
+            Scopes: []string{waveobj.MakeORef(update.OType, update.OID).String()},
+            Data:   update,
+        })
+    }
+}
+```
+
+## Debugging
+
+To debug event flow:
+
+1. Check broker subscription map: `wps.Broker.SubMap`
+2. View persisted events: `wps.Broker.ReadEventHistory(eventType, scope, maxItems)`
+3. Add logging in publish/subscribe methods
+4. Monitor WebSocket traffic in browser dev tools
+
+## Related Documentation
+
+- [Configuration System](config-system.md) - Uses WPS events for config updates
+- [Wave AI Architecture](waveai-architecture.md) - AI-related events

frontend/app/aipanel/aipanel.tsx

@@ -104,6 +104,10 @@ const AIPanelComponentInner = memo(({ className, onClose }: AIPanelProps) => {
         loadMessages();
     }, [model, setMessages]);
 
+    useEffect(() => {
+        model.ensureRateLimitSet();
+    }, [model]);
+
     const handleSubmit = async (e: React.FormEvent) => {
         e.preventDefault();
         if (!input.trim() || status !== "ready" || isLoadingChat) return;

frontend/app/aipanel/waveai-model.tsx

@@ -172,4 +172,19 @@ export class WaveAIModel {
             return [];
         }
     }
+
+    async ensureRateLimitSet() {
+        const currentInfo = globalStore.get(atoms.waveAIRateLimitInfoAtom);
+        if (currentInfo != null) {
+            return;
+        }
+        try {
+            const rateLimitInfo = await RpcApi.GetWaveAIRateLimitCommand(TabRpcClient);
+            if (rateLimitInfo != null) {
+                globalStore.set(atoms.waveAIRateLimitInfoAtom, rateLimitInfo);
+            }
+        } catch (error) {
+            console.error("Failed to fetch rate limit info:", error);
+        }
+    }
 }

frontend/app/store/global.ts

@@ -149,6 +149,7 @@ function initGlobalAtoms(initOpts: GlobalInitOptions) {
     const notificationPopoverModeAtom = atom<boolean>(false);
     const reinitVersion = atom(0);
     const waveAIFocusedAtom = atom(false);
+    const rateLimitInfoAtom = atom(null) as PrimitiveAtom<RateLimitInfo>;
     atoms = {
         // initialized in wave.ts (will not be null inside of application)
         clientId: clientIdAtom,
@@ -173,6 +174,7 @@ function initGlobalAtoms(initOpts: GlobalInitOptions) {
         reinitVersion,
         isTermMultiInput: atom(false),
         waveAIFocusedAtom,
+        waveAIRateLimitInfoAtom: rateLimitInfoAtom,
     };
 }
 
@@ -213,6 +215,13 @@ function initGlobalWaveEventSubs(initOpts: WaveInitOpts) {
                     fileSubject.next(fileData);
                 }
             },
+        },
+        {
+            eventType: "waveai:ratelimit",
+            handler: (event) => {
+                const rateLimitInfo: RateLimitInfo = event.data;
+                globalStore.set(atoms.waveAIRateLimitInfoAtom, rateLimitInfo);
+            },
         }
     );
 }

frontend/app/store/wshclientapi.ts

@@ -287,6 +287,11 @@ class RpcApiType {
         return client.wshRpcCall("getwaveaichat", data, opts);
     }
 
+    // command "getwaveairatelimit" [call]
+    GetWaveAIRateLimitCommand(client: WshClient, opts?: RpcOpts): Promise<RateLimitInfo> {
+        return client.wshRpcCall("getwaveairatelimit", null, opts);
+    }
+
     // command "message" [call]
     MessageCommand(client: WshClient, data: CommandMessageData, opts?: RpcOpts): Promise<void> {
         return client.wshRpcCall("message", data, opts);

frontend/types/custom.d.ts

@@ -29,6 +29,7 @@ declare global {
         reinitVersion: jotai.PrimitiveAtom<number>;
         isTermMultiInput: jotai.PrimitiveAtom<boolean>;
         waveAIFocusedAtom: jotai.PrimitiveAtom<boolean>;
+        waveAIRateLimitInfoAtom: jotai.PrimitiveAtom<RateLimitInfo>;
     };
 
     type WritableWaveObjectAtom<T extends WaveObj> = jotai.WritableAtom<T, [value: T], void>;

frontend/types/gotypes.d.ts

@@ -673,6 +673,14 @@ declare global {
         y: number;
     };
 
+    // uctypes.RateLimitInfo
+    type RateLimitInfo = {
+        remaining: number;
+        premiumremaining: number;
+        expirationepoch: number;
+        unknown?: boolean;
+    };
+
     // wshrpc.RemoteInfo
     type RemoteInfo = {
         clientarch: string;