feat: create and export UsageLimitedChatService (#1182)

* feat: create and export UsageLimitedChatService for when user exceeds usage limit

* tweak: change comment on usage-limited-chat to better explain action

* fixed whitespace

* tweak: remove test-app from gitignore

* tweak: remove extra spacing in AIChatService usage-limited-chat comment

* tweak: fixed whitespace

* tweak: remove tabbed whitespace in AIChatService

* tweak: remove .qodo from gitignore

* tweak: remove extra enter in AIChatService

* Revert "tweak: remove .qodo from gitignore"

tweak: correct mistake on removing extra enter

This reverts commit e066e294fa.

* tweak: add space at fallback and remove whitespace tabs

* tweak: remove whitespace on brackets

* tweak: remove .qodo and test-app from gitignore

* tweak: add accidentally deleted enter back in

---------

Co-authored-by: Rishabh Shinde <rishabhsshinde27@gail.com>

.gitignore

@@ -29,4 +29,4 @@ dist/
 
 # Local Netlify folder
 .netlify
-src/emulator/release/
+src/emulator/release/

src/backend/.gitignore

@@ -148,4 +148,4 @@ creds*
 thumbnail-service
 
 # init sql generated from ./run.sh
-init.sql
+init.sql

src/backend/src/modules/puterai/AIChatService.js

@@ -366,7 +366,7 @@ class AIChatService extends BaseService {
                 if ( ! event.allow ) {
                     test_mode = true;
                 }
-                
+
                 if ( parameters.messages ) {
                     parameters.messages =
                         Messages.normalize_messages(parameters.messages);
@@ -401,11 +401,23 @@ class AIChatService extends BaseService {
                 let model_used = this.get_model_from_request(parameters, {
                     intended_service
                 });
-                await this.check_usage_({
+
+                // Updated: Check usage and get a boolean result instead of throwing error
+                const usageAllowed = await this.check_usage_({
                     actor: Context.get('actor'),
                     service: service_used,
                     model: model_used,
                 });
+
+                // Handle usage limits reached case
+                if ( !usageAllowed ) {
+                    // The check_usage_ method has already updated the intended_service to 'usage-limited-chat'
+                    service_used = 'usage-limited-chat';
+                    model_used = 'usage-limited';
+                    // Update intended_service to match service_used
+                    intended_service = service_used;
+                }
+
                 try {
                     ret = await svc_driver.call_new_({
                         actor: Context.get('actor'),
@@ -423,7 +435,7 @@ class AIChatService extends BaseService {
                     tried.push(model);
 
                     error = e;
-                    
+
                     // Distinguishing between user errors and service errors
                     // is very messy because of different conventions between
                     // services. This is a best-effort attempt to catch user
@@ -482,45 +494,73 @@ class AIChatService extends BaseService {
                             fallback_model_name
                         });
 
-                        await this.check_usage_({
+                        // Check usage for fallback model too (with updated method)
+                        const fallbackUsageAllowed = await this.check_usage_({
                             actor: Context.get('actor'),
                             service: fallback_service_name,
                             model: fallback_model_name,
                         });
-                        try {
+                        
+                        // If usage not allowed for fallback, use usage-limited-chat instead
+                        if (!fallbackUsageAllowed) {
+                            // The check_usage_ method has already updated intended_service
+                            service_used = 'usage-limited-chat';
+                            model_used = 'usage-limited';
+                            // Clear the error to exit the fallback loop
+                            error = null;
+                            
+                            // Call the usage-limited service
                             ret = await svc_driver.call_new_({
                                 actor: Context.get('actor'),
-                                service_name: fallback_service_name,
+                                service_name: 'usage-limited-chat',
                                 skip_usage: true,
                                 iface: 'puter-chat-completion',
                                 method: 'complete',
-                                args: {
-                                    ...parameters,
+                                args: parameters,
+                            });
+                        } else {
+                            // Normal fallback flow continues
+                            try {
+                                ret = await svc_driver.call_new_({
+                                    actor: Context.get('actor'),
+                                    service_name: fallback_service_name,
+                                    skip_usage: true,
+                                    iface: 'puter-chat-completion',
+                                    method: 'complete',
+                                    args: {
+                                        ...parameters,
+                                        model: fallback_model_name,
+                                    },
+                                });
+                                error = null;
+                                service_used = fallback_service_name;
+                                model_used = fallback_model_name;
+                                response_metadata.fallback = {
+                                    service: fallback_service_name,
                                     model: fallback_model_name,
-                                },
-                            });
-                            error = null;
-                            service_used = fallback_service_name;
-                            model_used = fallback_model_name;
-                            response_metadata.fallback = {
-                                service: fallback_service_name,
-                                model: fallback_model_name,
-                                tried: tried,
-                            };
-                        } catch (e) {
-                            error = e;
-                            tried.push(fallback_model_name);
-                            this.log.error('error calling fallback', {
-                                intended_service,
-                                model,
-                                error: e,
-                            });
+                                    tried: tried,
+                                };
+                            } catch (e) {
+                                error = e;
+                                tried.push(fallback_model_name);
+                                this.log.error('error calling fallback', {
+                                    intended_service,
+                                    model,
+                                    error: e,
+                                });
+                            }
                         }
                     }
                 }
+                
                 ret.result.via_ai_chat_service = true;
                 response_metadata.service_used = service_used;
-
+                
+                // Add flag if we're using the usage-limited service
+                if (service_used === 'usage-limited-chat') {
+                    response_metadata.usage_limited = true;
+                }
+            
                 const username = Context.get('actor').type?.user?.username;
 
                 if (
@@ -639,10 +679,20 @@ class AIChatService extends BaseService {
             permission_options: options,
         };
         await svc_event.emit('ai.prompt.check-usage', event);
-        if ( event.error ) throw event.error;
-        if ( ! event.allowed ) {
-            throw new APIError('forbidden');
+        
+        // If the user has exceeded their usage limit, apply usage-limited-chat which lets them know
+        if ( event.error || ! event.allowed ) {
+            // Instead of throwing an error, modify the intended_service
+            const client_driver_call = Context.get('client_driver_call');
+            client_driver_call.intended_service = 'usage-limited-chat';
+            client_driver_call.response_metadata.usage_limited = true;
+            
+            // Return false to indicate that the user has gone over their limit and service has been changed
+            return false;
         }
+        
+        // Return true if the user has tokens to spend
+        return true;
     }
     
 

src/backend/src/modules/puterai/PuterAIModule.js

@@ -116,6 +116,9 @@ class PuterAIModule extends AdvancedBase {
 
         const{ AITestModeService } = require('./AITestModeService');
         services.registerService('ai-test-mode', AITestModeService);
+
+        const { UsageLimitedChatService } = require('./UsageLimitedChatService');
+        services.registerService('usage-limited-chat', UsageLimitedChatService);
     }
 }
 

src/backend/src/modules/puterai/UsageLimitedChatService.js

@@ -0,0 +1,170 @@
+/*
+ * Copyright (C) 2024-present Puter Technologies Inc.
+ * 
+ * This file is part of Puter.
+ * 
+ * Puter is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ * 
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ * 
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+// METADATA // {"ai-commented":{"service":"claude"}}
+const { default: dedent } = require("dedent");
+const BaseService = require("../../services/BaseService");
+const { PassThrough } = require("stream");
+const { TypedValue } = require("../../services/drivers/meta/Runtime");
+const Streaming = require("./lib/Streaming");
+
+/**
+* UsageLimitedChatService - A specialized chat service that returns resource exhaustion messages.
+* Extends BaseService to provide responses indicating the user has exceeded their usage limits.
+* Follows the same response format as real AI providers but with a custom message about upgrading.
+* Can handle both streaming and non-streaming requests consistently.
+*/
+class UsageLimitedChatService extends BaseService {
+    get_default_model () {
+        return 'usage-limited';
+    }
+    
+    static IMPLEMENTS = {
+        ['puter-chat-completion']: {
+            /**
+            * Returns a list of available model names
+            * @returns {Promise<string[]>} Array containing the single model identifier
+            */
+            async list () {
+                return ['usage-limited'];
+            },
+            
+            /**
+            * Returns model details for the usage-limited model
+            * @returns {Promise<Object[]>} Array containing the model details
+            */
+            async models () {
+                return [{
+                    id: 'usage-limited',
+                    name: 'Usage Limited',
+                    context: 16384,
+                    cost: {
+                        currency: 'usd-cents',
+                        tokens: 1_000_000,
+                        input: 0,
+                        output: 0,
+                    },
+                }];
+            },
+
+            /**
+            * Simulates a chat completion request with a usage limit message
+            * @param {Object} params - The completion parameters
+            * @param {Array} params.messages - Array of chat messages (unused)
+            * @param {boolean} params.stream - Whether to stream the response
+            * @param {string} params.model - The model to use (unused)
+            * @returns {Object|TypedValue} A chat completion response or streamed response
+            */
+            async complete ({ messages, stream, model, customLimitMessage }) {
+                const limitMessage = customLimitMessage || dedent(`
+                    You have reached your AI usage limit for this account.
+                `);
+                
+                // If streaming is requested, return a streaming response
+                if ( stream ) {
+                    const streamObj = new PassThrough();
+                    const retval = new TypedValue({
+                        $: 'stream',
+                        content_type: 'application/x-ndjson',
+                        chunked: true,
+                    }, streamObj);
+                    
+                    const chatStream = new Streaming.AIChatStream({
+                        stream: streamObj,
+                    });
+                    
+                    // Schedule the streaming response
+                    setTimeout(() => {
+                        chatStream.write({
+                            type: 'content_block_start',
+                            index: 0,
+                        });
+                        
+                        chatStream.write({
+                            type: 'content_block_delta',
+                            index: 0,
+                            delta: {
+                                type: 'text',
+                                text: limitMessage,
+                            },
+                        });
+                        
+                        chatStream.write({
+                            type: 'content_block_stop',
+                            index: 0,
+                        });
+                        
+                        chatStream.write({
+                            type: 'message_stop',
+                            stop_reason: 'end_turn',
+                        });
+                        
+                        chatStream.end();
+                    }, 10);
+                    
+                    // Return a TypedValue with usage_promise for proper integration
+                    return new TypedValue({ $: 'ai-chat-intermediate' }, {
+                        stream: true,
+                        init_chat_stream: async ({ chatStream: cs }) => {
+                            // Copy contents from our stream to the provided one
+                            chatStream.pipe(cs.stream);
+                        },
+                        usage_promise: Promise.resolve({
+                            input_tokens: 0,
+                            output_tokens: 1,
+                        }),
+                    });
+                }
+                
+                // Non-streaming response
+                return {
+                    "index": 0,
+                    message: {
+                        "id": "00000000-0000-0000-0000-000000000000",
+                        "type": "message",
+                        "role": "assistant",
+                        "model": "usage-limited",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": limitMessage
+                            }
+                        ],
+                        "stop_reason": "end_turn",
+                        "stop_sequence": null,
+                        "usage": {
+                            "input_tokens": 0,
+                            "output_tokens": 1
+                        }
+                    },
+                    "usage": {
+                        "input_tokens": 0,
+                        "output_tokens": 1
+                    },
+                    "logprobs": null,
+                    "finish_reason": "stop"
+                };
+            }
+        }
+    }
+}
+
+module.exports = {
+    UsageLimitedChatService,
+};