diff --git a/packages/services/service-ai/src/__tests__/agent-chat-quota.test.ts b/packages/services/service-ai/src/__tests__/agent-chat-quota.test.ts new file mode 100644 index 0000000000..8cb40098cf --- /dev/null +++ b/packages/services/service-ai/src/__tests__/agent-chat-quota.test.ts @@ -0,0 +1,197 @@ +// Copyright (c) 2025 ObjectStack. Licensed under the Apache-2.0 license. + +import { describe, it, expect, vi } from 'vitest'; +import type { IDataEngine, Logger } from '@objectstack/spec/contracts'; +import { DailyMessageQuota, type AgentChatQuota } from '../quota/agent-chat-quota.js'; +import { buildAgentRoutes } from '../routes/agent-routes.js'; +import { AIService } from '../ai-service.js'; +import { MemoryLLMAdapter } from '../adapters/memory-adapter.js'; +import { InMemoryConversationService } from '../conversation/in-memory-conversation-service.js'; +import type { AgentRuntime } from '../agent-runtime.js'; + +const silentLogger: Logger = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), +}; + +const FIXED_NOW = () => new Date('2026-06-12T10:00:00.000Z'); + +function mockDataEngine(rows: Record = {}): IDataEngine & { + inserted: unknown[]; + updated: unknown[]; +} { + const inserted: unknown[] = []; + const updated: unknown[] = []; + return { + inserted, + updated, + findOne: vi.fn(async (_obj: string, q?: { where?: { id?: string } }) => { + const id = q?.where?.id; + return id && rows[id] ? { id, ...rows[id] } : null; + }), + insert: vi.fn(async (_obj: string, data: unknown) => { + inserted.push(data); + return data; + }), + update: vi.fn(async (_obj: string, data: unknown, opts: unknown) => { + updated.push({ data, opts }); + return data; + }), + find: vi.fn(async () => []), + delete: vi.fn(async () => ({})), + count: vi.fn(async () => 0), + aggregate: vi.fn(async () => []), + } as unknown as IDataEngine & { inserted: unknown[]; updated: unknown[] }; +} + +// ═══════════════════════════════════════════════════════════════════ +// DailyMessageQuota +// ═══════════════════════════════════════════════════════════════════ + +describe('DailyMessageQuota', () => { + const subject = { userId: 'u1', environmentId: 'env1' }; + const todayId = '2026-06-12:env1:u1'; + + it('allows under the limit and reports remaining + resetAt', async () => { + const quota = new DailyMessageQuota(mockDataEngine({ [todayId]: { messages: 3 } }), 30, FIXED_NOW); + const d = await quota.check(subject); + expect(d.allowed).toBe(true); + expect(d.remaining).toBe(26); // 30 - 3 - this turn + expect(d.resetAt).toBe('2026-06-13T00:00:00.000Z'); + }); + + it('refuses at the limit with honest copy (why + when + way out)', async () => { + const quota = new DailyMessageQuota(mockDataEngine({ [todayId]: { messages: 30 } }), 30, FIXED_NOW); + const d = await quota.check(subject); + expect(d.allowed).toBe(false); + expect(d.resetAt).toBe('2026-06-13T00:00:00.000Z'); + expect(d.message).toContain('30'); + expect(d.message).toContain('恢复'); + expect(d.message).toContain('upgrade'); + }); + + it('consume inserts the first turn of the day, then increments', async () => { + const engine = mockDataEngine(); + const quota = new DailyMessageQuota(engine, 30, FIXED_NOW); + await quota.consume(subject); + expect(engine.inserted).toEqual([ + { id: todayId, day: '2026-06-12', user_id: 'u1', environment_id: 'env1', messages: 1 }, + ]); + + const engine2 = mockDataEngine({ [todayId]: { messages: 5 } }); + const quota2 = new DailyMessageQuota(engine2, 30, FIXED_NOW); + await quota2.consume(subject); + expect(engine2.updated).toEqual([{ data: { messages: 6 }, opts: { where: { id: todayId } } }]); + }); + + it('fails OPEN when the counter store errors — never blocks chat', async () => { + const engine = mockDataEngine(); + (engine.findOne as ReturnType).mockRejectedValue(new Error('db down')); + const quota = new DailyMessageQuota(engine, 1, FIXED_NOW); + await expect(quota.check(subject)).resolves.toMatchObject({ allowed: true }); + await expect(quota.consume(subject)).resolves.toBeUndefined(); + }); + + it('scopes the counter per environment and per user', async () => { + const engine = mockDataEngine({ ['2026-06-12:envA:u1']: { messages: 99 } }); + const quota = new DailyMessageQuota(engine, 10, FIXED_NOW); + // Same user, different environment → independent counter. + expect((await quota.check({ userId: 'u1', environmentId: 'envB' })).allowed).toBe(true); + // Different user, same environment → independent counter. + expect((await quota.check({ userId: 'u2', environmentId: 'envA' })).allowed).toBe(true); + // The exhausted pair stays refused. + expect((await quota.check({ userId: 'u1', environmentId: 'envA' })).allowed).toBe(false); + }); +}); + +// ═══════════════════════════════════════════════════════════════════ +// Agent chat route × quota gate +// ═══════════════════════════════════════════════════════════════════ + +function mockAgentRuntime(): AgentRuntime { + return { + loadAgent: vi.fn(async () => ({ + name: 'data_chat', + label: 'Assistant', + active: true, + })), + resolveActiveSkills: vi.fn(async () => []), + buildSystemMessages: vi.fn(() => [{ role: 'system', content: 'sys' }]), + buildRequestOptions: vi.fn(() => ({})), + listAgents: vi.fn(async () => []), + } as unknown as AgentRuntime; +} + +function chatRoute(quota?: AgentChatQuota) { + const aiService = new AIService({ + adapter: new MemoryLLMAdapter(), + conversationService: new InMemoryConversationService(), + }); + const routes = buildAgentRoutes(aiService, mockAgentRuntime(), silentLogger, { quota }); + const route = routes.find((r) => r.path.endsWith('/chat')); + if (!route) throw new Error('chat route not found'); + return route; +} + +const chatReq = (over: Record = {}) => ({ + params: { agentName: 'data_chat' }, + body: { messages: [{ role: 'user', content: 'hi' }], stream: false, ...over }, + user: { userId: 'u1' }, +}); + +describe('agent chat route quota gate', () => { + it('passes through and consumes exactly once when allowed', async () => { + const quota: AgentChatQuota = { + check: vi.fn(async () => ({ allowed: true })), + consume: vi.fn(async () => {}), + }; + const res = await chatRoute(quota).handler(chatReq() as never); + expect(res.status).toBe(200); + expect(quota.check).toHaveBeenCalledTimes(1); + expect(quota.consume).toHaveBeenCalledTimes(1); + }); + + it('returns 429 + stable code on JSON mode when exhausted (nothing consumed)', async () => { + const quota: AgentChatQuota = { + check: vi.fn(async () => ({ allowed: false, message: '额度已用完', resetAt: '2026-06-13T00:00:00.000Z' })), + consume: vi.fn(async () => {}), + }; + const res = await chatRoute(quota).handler(chatReq() as never); + expect(res.status).toBe(429); + expect(res.body).toMatchObject({ code: 'ai_quota_exhausted', error: '额度已用完', resetAt: '2026-06-13T00:00:00.000Z' }); + expect(quota.consume).not.toHaveBeenCalled(); + }); + + it('streams the refusal as a normal assistant message in stream mode', async () => { + const quota: AgentChatQuota = { + check: vi.fn(async () => ({ allowed: false, message: '今日额度已用完,明日恢复' })), + consume: vi.fn(async () => {}), + }; + const res = await chatRoute(quota).handler(chatReq({ stream: true }) as never); + expect(res.status).toBe(200); + expect(res.stream).toBe(true); + let text = ''; + for await (const chunk of res.events as AsyncIterable) text += chunk; + expect(text).toContain('今日额度已用完'); + expect(text).toContain('"type":"finish"'); + expect(quota.consume).not.toHaveBeenCalled(); + }); + + it('no quota wired → unchanged behavior', async () => { + const res = await chatRoute(undefined).handler(chatReq() as never); + expect(res.status).toBe(200); + }); + + it('forwards the environmentId from chat context to the quota subject', async () => { + const quota: AgentChatQuota = { + check: vi.fn(async () => ({ allowed: true })), + consume: vi.fn(async () => {}), + }; + await chatRoute(quota).handler(chatReq({ context: { environmentId: 'env42' } }) as never); + expect(quota.check).toHaveBeenCalledWith({ userId: 'u1', environmentId: 'env42' }); + expect(quota.consume).toHaveBeenCalledWith({ userId: 'u1', environmentId: 'env42' }); + }); +}); diff --git a/packages/services/service-ai/src/objects/ai-usage-daily.object.ts b/packages/services/service-ai/src/objects/ai-usage-daily.object.ts new file mode 100644 index 0000000000..8d0fdcc27b --- /dev/null +++ b/packages/services/service-ai/src/objects/ai-usage-daily.object.ts @@ -0,0 +1,59 @@ +// Copyright (c) 2025 ObjectStack. Licensed under the Apache-2.0 license. + +import { ObjectSchema, Field } from '@objectstack/spec/data'; + +/** + * ai_usage_daily — per-user, per-day AI chat usage counters + * + * One row per (UTC day, environment, user). The agent chat route increments + * `messages` once per user turn; {@link DailyMessageQuota} reads it to decide + * whether a daily message limit has been reached. + * + * This is a quota counter, not billing-grade metering: increments are + * last-write-wins and a lost update under concurrency only ever under-counts + * by a message — acceptable for an abuse/entitlement gate, not for invoicing. + * + * @namespace ai + */ +export const AiUsageDailyObject = ObjectSchema.create({ + name: 'ai_usage_daily', + label: 'AI Daily Usage', + pluralLabel: 'AI Daily Usage', + icon: 'gauge', + isSystem: true, + description: 'Per-user daily AI chat usage counters (quota enforcement)', + + fields: { + id: Field.text({ + label: 'Usage ID', + required: true, + readonly: true, + description: 'Deterministic key: ::', + }), + + day: Field.text({ + label: 'Day (UTC)', + required: true, + maxLength: 10, + description: 'UTC calendar day, YYYY-MM-DD', + }), + + user_id: Field.text({ + label: 'User ID', + required: true, + maxLength: 255, + }), + + environment_id: Field.text({ + label: 'Environment ID', + required: false, + maxLength: 255, + }), + + messages: Field.number({ + label: 'Messages', + required: true, + description: 'User chat turns consumed this day', + }), + }, +}); diff --git a/packages/services/service-ai/src/objects/index.ts b/packages/services/service-ai/src/objects/index.ts index 91fdb018e6..5a0b43b9b8 100644 --- a/packages/services/service-ai/src/objects/index.ts +++ b/packages/services/service-ai/src/objects/index.ts @@ -12,3 +12,4 @@ export { AiTraceObject } from './ai-trace.object.js'; export { AiPendingActionObject } from './ai-pending-action.object.js'; export { AiEvalCaseObject } from './ai-eval-case.object.js'; export { AiEvalRunObject } from './ai-eval-run.object.js'; +export { AiUsageDailyObject } from './ai-usage-daily.object.js'; diff --git a/packages/services/service-ai/src/plugin.ts b/packages/services/service-ai/src/plugin.ts index ceb578aea3..832a25b824 100644 --- a/packages/services/service-ai/src/plugin.ts +++ b/packages/services/service-ai/src/plugin.ts @@ -14,7 +14,8 @@ import { buildToolRoutes } from './routes/tool-routes.js'; import { buildPendingActionRoutes } from './routes/pending-action-routes.js'; import { buildEvalRoutes } from './routes/eval-routes.js'; import { ObjectQLConversationService } from './conversation/objectql-conversation-service.js'; -import { AiConversationObject, AiMessageObject, AiPendingActionObject, AiTraceObject, AiEvalCaseObject, AiEvalRunObject } from './objects/index.js'; +import { AiConversationObject, AiMessageObject, AiPendingActionObject, AiTraceObject, AiEvalCaseObject, AiEvalRunObject, AiUsageDailyObject } from './objects/index.js'; +import { DailyMessageQuota, type AgentChatQuota } from './quota/agent-chat-quota.js'; import { AiTraceView, AiMessageView, AiPendingActionView, AiEvalCaseView, AiEvalRunView } from './views/index.js'; import { EvalRunner } from './eval/index.js'; import { registerDataTools } from './tools/data-tools.js'; @@ -593,7 +594,7 @@ export class AIServicePlugin implements Plugin { type: 'plugin', scope: 'system', namespace: 'ai', - objects: [AiConversationObject, AiMessageObject, AiTraceObject, AiPendingActionObject, AiEvalCaseObject, AiEvalRunObject], + objects: [AiConversationObject, AiMessageObject, AiTraceObject, AiPendingActionObject, AiEvalCaseObject, AiEvalRunObject, AiUsageDailyObject], views: [AiTraceView, AiMessageView, AiPendingActionView, AiEvalCaseView, AiEvalRunView], }); @@ -857,7 +858,26 @@ export class AIServicePlugin implements Plugin { if (metadataService) { const skillRegistry = new SkillRegistry(metadataService); const agentRuntime = new AgentRuntime(metadataService, skillRegistry); - const agentRoutes = buildAgentRoutes(this.service, agentRuntime, ctx.logger); + + // ── Optional per-turn chat quota (ADR-0040 §5, perception rule) ── + // Mechanism only: the deployment opts in by setting + // AI_DAILY_USER_MESSAGES=. Unset/invalid → no quota, unchanged + // behavior. Plan-tier policies (free vs pro) wire a richer + // AgentChatQuota here later; the gate and counter do not change. + let chatQuota: AgentChatQuota | undefined; + const quotaRaw = process.env.AI_DAILY_USER_MESSAGES; + const quotaLimit = quotaRaw ? Number.parseInt(quotaRaw, 10) : NaN; + if (Number.isFinite(quotaLimit) && quotaLimit > 0) { + const quotaDataEngine = ctx.getService('data'); + if (quotaDataEngine && typeof quotaDataEngine.findOne === 'function') { + chatQuota = new DailyMessageQuota(quotaDataEngine, quotaLimit); + ctx.logger.info(`[AI] Daily chat quota enabled (${quotaLimit} user turns/user/day)`); + } else { + ctx.logger.warn('[AI] AI_DAILY_USER_MESSAGES set but IDataEngine unavailable — quota disabled'); + } + } + + const agentRoutes = buildAgentRoutes(this.service, agentRuntime, ctx.logger, { quota: chatQuota }); routes.push(...agentRoutes); ctx.logger.info(`[AI] Agent routes registered (${agentRoutes.length} routes)`); diff --git a/packages/services/service-ai/src/quota/agent-chat-quota.ts b/packages/services/service-ai/src/quota/agent-chat-quota.ts new file mode 100644 index 0000000000..3d775ea649 --- /dev/null +++ b/packages/services/service-ai/src/quota/agent-chat-quota.ts @@ -0,0 +1,121 @@ +// Copyright (c) 2025 ObjectStack. Licensed under the Apache-2.0 license. + +import type { IDataEngine } from '@objectstack/spec/contracts'; + +const USAGE_OBJECT = 'ai_usage_daily'; + +/** Identity a quota decision is made for. */ +export interface QuotaSubject { + userId: string; + environmentId?: string; +} + +/** Outcome of a quota check. */ +export interface AgentChatQuotaDecision { + allowed: boolean; + /** Turns left today (after this one), when the implementation knows. */ + remaining?: number; + /** ISO timestamp when the quota resets (next UTC midnight for daily). */ + resetAt?: string; + /** + * Honest, user-facing refusal copy (perception rule, ADR-0040 §5): a limit + * that can only manifest as refusal must say why, when it recovers, and + * what the way out is — never degrade silently. + */ + message?: string; +} + +/** + * Pluggable per-turn quota for the agent chat route. + * + * The route calls {@link check} before dispatching a user turn and + * {@link consume} exactly once when the turn is admitted. Implementations + * decide the policy (daily counters, plan entitlements, token buckets); + * the route only enforces the decision. No quota wired → no behavior change. + */ +export interface AgentChatQuota { + check(subject: QuotaSubject): Promise; + consume(subject: QuotaSubject): Promise; +} + +/** UTC calendar day (YYYY-MM-DD) for `now`. */ +function utcDay(now: Date): string { + return now.toISOString().slice(0, 10); +} + +/** ISO timestamp of the next UTC midnight after `now`. */ +function nextUtcMidnight(now: Date): string { + const next = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate() + 1)); + return next.toISOString(); +} + +/** Deterministic counter row id — one row per (day, environment, user). */ +function usageId(day: string, subject: QuotaSubject): string { + return `${day}:${subject.environmentId ?? '-'}:${subject.userId}`; +} + +/** + * DailyMessageQuota — N user turns per user per environment per UTC day. + * + * Backed by the `ai_usage_daily` platform object through {@link IDataEngine}, + * so the count survives restarts and is shared across instances. Counter + * semantics (read → write, last-write-wins): a concurrent lost update only + * under-counts a turn, which is fine for an entitlement gate and NOT fit for + * billing. Failures are fail-open — a broken counter must never take chat + * down with it. + */ +export class DailyMessageQuota implements AgentChatQuota { + constructor( + private readonly dataEngine: IDataEngine, + private readonly dailyLimit: number, + /** Injectable clock for tests. */ + private readonly now: () => Date = () => new Date(), + ) {} + + async check(subject: QuotaSubject): Promise { + const at = this.now(); + const day = utcDay(at); + let used = 0; + try { + const row = await this.dataEngine.findOne(USAGE_OBJECT, { where: { id: usageId(day, subject) } }); + used = typeof row?.messages === 'number' ? row.messages : 0; + } catch { + return { allowed: true }; // fail-open: never block chat on a counter error + } + if (used < this.dailyLimit) { + return { allowed: true, remaining: this.dailyLimit - used - 1, resetAt: nextUtcMidnight(at) }; + } + const resetAt = nextUtcMidnight(at); + return { + allowed: false, + remaining: 0, + resetAt, + message: + `今日 AI 助手额度(${this.dailyLimit} 条/天)已用完,将于明日(UTC)自动恢复;如需立即继续,请联系管理员或升级套餐。 ` + + `Daily AI assistant limit reached (${this.dailyLimit} messages/day). It resets at ${resetAt}; ` + + 'to continue now, contact your administrator or upgrade your plan.', + }; + } + + async consume(subject: QuotaSubject): Promise { + const day = utcDay(this.now()); + const id = usageId(day, subject); + try { + const row = await this.dataEngine.findOne(USAGE_OBJECT, { where: { id } }); + if (row) { + const used = typeof row.messages === 'number' ? row.messages : 0; + await this.dataEngine.update(USAGE_OBJECT, { messages: used + 1 }, { where: { id } }); + } else { + await this.dataEngine.insert(USAGE_OBJECT, { + id, + day, + user_id: subject.userId, + environment_id: subject.environmentId ?? null, + messages: 1, + }); + } + } catch { + /* fail-open: a lost increment under-counts one turn; never break chat */ + } + } +} diff --git a/packages/services/service-ai/src/routes/agent-routes.ts b/packages/services/service-ai/src/routes/agent-routes.ts index 1582ee84d0..30465f1760 100644 --- a/packages/services/service-ai/src/routes/agent-routes.ts +++ b/packages/services/service-ai/src/routes/agent-routes.ts @@ -2,9 +2,11 @@ import type { ModelMessage } from '@objectstack/spec/contracts'; import type { Logger } from '@objectstack/spec/contracts'; +import type { TextStreamPart, ToolSet } from 'ai'; import type { AIService } from '../ai-service.js'; import type { AgentRuntime, AgentChatContext } from '../agent-runtime.js'; import type { RouteDefinition } from './ai-routes.js'; +import type { AgentChatQuota } from '../quota/agent-chat-quota.js'; import { normalizeMessage, validateMessageContent } from './message-utils.js'; import { encodeVercelDataStream } from '../stream/vercel-stream-encoder.js'; @@ -33,6 +35,25 @@ function validateAgentMessage(raw: unknown): string | null { return validateMessageContent(msg, { allowEmptyContent: allowEmpty }); } +/** Optional behaviors for {@link buildAgentRoutes}. */ +export interface AgentRouteOptions { + /** + * Per-turn chat quota. Checked before each user turn is dispatched and + * consumed exactly once when admitted. Absent → no quota (unchanged + * behavior). Policy lives in the implementation; the route only enforces. + */ + quota?: AgentChatQuota; +} + +/** + * A quota refusal as a well-formed UI message stream: the honest copy arrives + * as an ordinary assistant message (perception rule, ADR-0040 §5), so no + * client change is needed and the chat never shows a raw transport error. + */ +async function* quotaRefusalParts(message: string): AsyncIterable> { + yield { type: 'text-delta', text: message } as TextStreamPart; +} + /** * Build agent-specific REST routes. * @@ -45,6 +66,7 @@ export function buildAgentRoutes( aiService: AIService, agentRuntime: AgentRuntime, logger: Logger, + options?: AgentRouteOptions, ): RouteDefinition[] { return [ // ── List active agents ────────────────────────────────────── @@ -116,6 +138,45 @@ export function buildAgentRoutes( return { status: 403, body: { error: `Agent "${agentName}" is not active` } }; } + // ── Per-turn quota gate (optional) ─────────────────────── + // Refusal is HONEST at the moment of impact (ADR-0040 §5): why, when + // it recovers, and the way out. Streaming clients get the copy as a + // normal assistant message; JSON clients get a 429 with a stable code. + const wantStreamMode = body.stream !== false; + if (options?.quota && req.user) { + const subject = { + userId: req.user.userId, + environmentId: + typeof chatContext?.environmentId === 'string' ? chatContext.environmentId : undefined, + }; + const decision = await options.quota.check(subject); + if (!decision.allowed) { + const message = + decision.message ?? 'Daily AI message limit reached. Please try again tomorrow.'; + if (wantStreamMode) { + return { + status: 200, + stream: true, + vercelDataStream: true, + contentType: 'text/event-stream', + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + 'Connection': 'keep-alive', + 'x-vercel-ai-ui-message-stream': 'v1', + }, + events: encodeVercelDataStream(quotaRefusalParts(message)), + }; + } + return { + status: 429, + body: { error: message, code: 'ai_quota_exhausted', resetAt: decision.resetAt }, + }; + } + // Admitted: count the turn now (not per tool call, not per token). + await options.quota.consume(subject); + } + try { // Resolve active skills for this agent in the current context const activeSkills = await agentRuntime.resolveActiveSkills(agent, chatContext);