learnpress / inc / AI / Assistant / TokenQuotaGuard.php

TokenQuotaGuard.php in LearnPress – WordPress LMS Plugin for Create and Sell Online Courses 4.4.4, at inc/AI/Assistant/TokenQuotaGuard.php

375 lines 10.3 KB

Raw Download Zip

1	<?php
2
3	namespace LearnPress\AI\Assistant;
4
5	use LearnPress\Services\OpenAiService;
6	use LP_Settings;
7
8	/**
9	* TokenQuotaGuard — per-user daily token quota tracking and enforcement.
10	*
11	* Wraps OpenAI chat requests with a quota check so that Agent.php is not
12	* responsible for any quota bookkeeping.
13	*
14	* @package LearnPress\AI\Assistant
15	* @since 4.3.5
16	*/
17	class TokenQuotaGuard {
18
19	private const USER_META_DAILY_TOKEN_USAGE = '_lp_ai_assistant_daily_token_usage';
20
21	/**
22	* Prefix of the per-user/per-day advisory lock option.
23	*/
24	private const LOCK_OPTION_PREFIX = '_lp_ai_assistant_quota_lock_';
25
26	/**
27	* Seconds after which a held lock is considered abandoned.
28	*
29	* Must exceed the slowest realistic OpenAI round trip, otherwise a second request
30	* could steal the lock while the first is still awaiting a response.
31	*/
32	private const LOCK_TTL = 120;
33
34	/**
35	* Human-readable block message set after the first exhausted-quota call.
36	*/
37	private string $block_message = '';
38
39	/**
40	* Send an OpenAI chat request guarded by the daily token quota.
41	*
42	* Check-call-record runs inside a per-user/per-day lock. Without it, N concurrent
43	* requests all read the same pre-call usage figure, all pass the check, and all
44	* spend tokens — the quota is bypassed by exactly the amount of concurrency.
45	*
46	* On quota exhaustion, or when the lock cannot be taken, the guard sets a block
47	* message and returns an empty array instead of calling the API (fail closed).
48	*
49	* @param OpenAiService $service OpenAI service instance.
50	* @param array $messages Chat messages payload.
51	* @param int $user_id Current user ID.
52	*
53	* @return array OpenAI response, or empty array when blocked.
54	* @throws \Throwable Re-throws underlying OpenAI errors.
55	*/
56	public function send_chat_with_guard( OpenAiService $service, array $messages, int $user_id ): array {
57
58	// 0 means unlimited: nothing to serialize, so do not pay for the lock.
59	if ( $this->get_daily_token_limit() <= 0 ) {
60	$response = $service->send_chat_request( array( 'messages' => $messages ) );
61	$this->track_token_usage_from_response( $user_id, $response );
62
63	return $response;
64	}
65
66	if ( ! $this->acquire_quota_lock( $user_id ) ) {
67	$this->block_message = $this->build_busy_message();
68
69	return array();
70	}
71
72	try {
73	// Re-read under the lock: another request may have consumed the remaining
74	// budget between our first look and acquiring it.
75	if ( $this->has_reached_daily_token_limit( $user_id, true ) ) {
76	$this->block_message = $this->build_block_message();
77
78	return array();
79	}
80
81	$response = $service->send_chat_request( array( 'messages' => $messages ) );
82
83	// Record actual usage before releasing, so the next waiter reads a current total.
84	$this->track_token_usage_from_response( $user_id, $response );
85
86	if ( $this->has_reached_daily_token_limit( $user_id ) ) {
87	$this->block_message = $this->build_block_message();
88	}
89
90	return $response;
91	} finally {
92	$this->release_quota_lock( $user_id );
93	}
94	}
95
96	/**
97	* Whether the last send_chat_with_guard call was blocked by quota.
98	*
99	* @return bool
100	*/
101	public function is_blocked(): bool {
102	return $this->block_message !== '';
103	}
104
105	/**
106	* Human-readable block message (empty string when not blocked).
107	*
108	* @return string
109	*/
110	public function get_block_message(): string {
111	return $this->block_message;
112	}
113
114	/**
115	* Reset block state between runs.
116	*
117	* @return void
118	*/
119	public function reset(): void {
120	$this->block_message = '';
121	}
122
123	// ----------------------------------------------------------------
124	// Quota lock
125	// ----------------------------------------------------------------
126
127	/**
128	* Option name of the per-user/per-day lock.
129	*
130	* Scoped by date so a lock abandoned on a previous day can never block today.
131	*
132	* @param int $user_id Current user ID.
133	*
134	* @return string
135	*/
136	private function get_lock_option_name( int $user_id ): string {
137	return self::LOCK_OPTION_PREFIX . $user_id . '_' . $this->get_local_current_date();
138	}
139
140	/**
141	* Acquire the quota lock for a user.
142	*
143	* Uses INSERT IGNORE against the unique option_name index, which is atomic across
144	* concurrent PHP workers. add_option()/get_option() cannot be used here: they
145	* read-then-write, leaving exactly the race this lock exists to close. Mirrors the
146	* approach in WP core's WP_Upgrader::create_lock().
147	*
148	* @param int $user_id Current user ID.
149	*
150	* @return bool True when the lock is held by this request.
151	*/
152	private function acquire_quota_lock( int $user_id ): bool {
153	global $wpdb;
154
155	if ( $user_id <= 0 ) {
156	return false;
157	}
158
159	$lock_option = $this->get_lock_option_name( $user_id );
160	$now = time();
161
162	$acquired = $wpdb->query(
163	$wpdb->prepare(
164	"INSERT IGNORE INTO `$wpdb->options` ( `option_name`, `option_value`, `autoload` ) VALUES ( %s, %s, 'no' )",
165	$lock_option,
166	(string) $now
167	)
168	);
169
170	if ( 1 === (int) $acquired ) {
171	return true;
172	}
173
174	// Someone holds it. Recover only if it is older than the TTL, i.e. the holder
175	// died before releasing (fatal error, timeout, killed worker).
176	$held_since = (int) $wpdb->get_var(
177	$wpdb->prepare( "SELECT `option_value` FROM `$wpdb->options` WHERE `option_name` = %s LIMIT 1", $lock_option )
178	);
179
180	if ( $held_since > 0 && ( $now - $held_since ) < self::LOCK_TTL ) {
181	return false;
182	}
183
184	/**
185	* Stale (or unreadable) lock: delete and re-attempt exactly once. The retry is
186	* still an atomic INSERT IGNORE, so if several requests detect the same stale
187	* lock simultaneously only one of them wins.
188	*/
189	$wpdb->delete( $wpdb->options, array( 'option_name' => $lock_option ) );
190	wp_cache_delete( $lock_option, 'options' );
191
192	$acquired = $wpdb->query(
193	$wpdb->prepare(
194	"INSERT IGNORE INTO `$wpdb->options` ( `option_name`, `option_value`, `autoload` ) VALUES ( %s, %s, 'no' )",
195	$lock_option,
196	(string) $now
197	)
198	);
199
200	return 1 === (int) $acquired;
201	}
202
203	/**
204	* Release the quota lock for a user.
205	*
206	* @param int $user_id Current user ID.
207	*
208	* @return void
209	*/
210	private function release_quota_lock( int $user_id ): void {
211	global $wpdb;
212
213	if ( $user_id <= 0 ) {
214	return;
215	}
216
217	$lock_option = $this->get_lock_option_name( $user_id );
218
219	$wpdb->delete( $wpdb->options, array( 'option_name' => $lock_option ) );
220	wp_cache_delete( $lock_option, 'options' );
221	}
222
223	// ----------------------------------------------------------------
224	// Private helpers
225	// ----------------------------------------------------------------
226
227	/**
228	* Track token usage from an OpenAI response usage payload.
229	*
230	* @param int $user_id Current user ID.
231	* @param array $response Raw OpenAI response.
232	*
233	* @return void
234	*/
235	private function track_token_usage_from_response( int $user_id, array $response ): void {
236
237	$usage = $response['usage'] ?? array();
238	$total_tokens = absint( $usage['total_tokens'] ?? 0 );
239
240	if ( $total_tokens <= 0 ) {
241	return;
242	}
243
244	$this->increase_daily_token_usage( $user_id, $total_tokens );
245	}
246
247	/**
248	* Check whether the learner has reached their daily token limit.
249	*
250	* @param int $user_id Current user ID.
251	* @param bool $bypass_cache Re-read usage from the database instead of the object
252	* cache. Required for the recheck under the lock, where a
253	* value cached earlier in this request would be stale.
254	*
255	* @return bool
256	*/
257	private function has_reached_daily_token_limit( int $user_id, bool $bypass_cache = false ): bool {
258
259	$limit = $this->get_daily_token_limit();
260	if ( $limit <= 0 ) {
261	return false;
262	}
263
264	if ( $bypass_cache && $user_id > 0 ) {
265	wp_cache_delete( $user_id, 'user_meta' );
266	}
267
268	return $this->get_daily_token_usage( $user_id ) >= $limit;
269	}
270
271	/**
272	* Get the configured daily token limit from plugin settings.
273	*
274	* @return int 0 means unlimited.
275	*/
276	private function get_daily_token_limit(): int {
277	return absint( LP_Settings::get_option( 'ai_assistant_max_usage_tokens_per_day', 0 ) );
278	}
279
280	/**
281	* Read learner daily token usage from user meta.
282	*
283	* @param int $user_id Current user ID.
284	*
285	* @return int
286	*/
287	private function get_daily_token_usage( int $user_id ): int {
288
289	if ( $user_id <= 0 ) {
290	return 0;
291	}
292
293	$payload = get_user_meta( $user_id, self::USER_META_DAILY_TOKEN_USAGE, true );
294	if ( ! is_array( $payload ) ) {
295	return 0;
296	}
297
298	$current_date = $this->get_local_current_date();
299	$stored_date = (string) ( $payload['date'] ?? '' );
300	if ( $stored_date !== $current_date ) {
301	return 0;
302	}
303
304	return absint( $payload['total_tokens'] ?? 0 );
305	}
306
307	/**
308	* Increase learner daily token usage and persist to user meta.
309	*
310	* @param int $user_id Current user ID.
311	* @param int $tokens Tokens consumed this call.
312	*
313	* @return void
314	*/
315	private function increase_daily_token_usage( int $user_id, int $tokens ): void {
316
317	if ( $user_id <= 0 \|\| $tokens <= 0 ) {
318	return;
319	}
320
321	// Read through to the database: the caller holds the quota lock, and a value
322	// cached earlier in this request would undercount the running total.
323	wp_cache_delete( $user_id, 'user_meta' );
324
325	$current_total = $this->get_daily_token_usage( $user_id );
326	$next_total = $current_total + $tokens;
327
328	update_user_meta(
329	$user_id,
330	self::USER_META_DAILY_TOKEN_USAGE,
331	array(
332	'date' => $this->get_local_current_date(),
333	'total_tokens' => $next_total,
334	)
335	);
336	}
337
338	/**
339	* Get local current date key for daily usage bucket.
340	*
341	* @return string
342	*/
343	private function get_local_current_date(): string {
344	return (string) current_time( 'Y-m-d' );
345	}
346
347	/**
348	* Build user-facing message for a request that could not take the quota lock.
349	*
350	* Reached when another request from the same learner is mid-flight, or a stale lock
351	* has not yet aged past its TTL. Failing closed keeps the quota authoritative.
352	*
353	* @return string
354	*/
355	private function build_busy_message(): string {
356	return __( 'Another AI Assistant request is still running. Please wait a moment and try again.', 'learnpress' );
357	}
358
359	/**
360	* Build user-facing quota exceeded message.
361	*
362	* @return string
363	*/
364	private function build_block_message(): string {
365
366	$limit = $this->get_daily_token_limit();
367
368	return sprintf(
369	/* translators: %d: max usage tokens per learner per day. */
370	__( 'Daily AI usage limit reached (%d tokens). Please try again tomorrow or contact the site administrator.', 'learnpress' ),
371	$limit
372	);
373	}
374	}
375