Skip to content

Commit 57abe7a

Browse files
os-zhuangclaude
andauthored
feat(service-ai): optional per-turn daily chat quota on the agent chat route (ADR-0040 §5) (#1760)
Mechanism only, default OFF — no behavior change unless a deployment sets AI_DAILY_USER_MESSAGES=<N>: - AgentChatQuota hook on buildAgentRoutes: checked before each user turn, consumed exactly once when admitted (per turn, not per tool call). - Refusal is honest at the moment of impact (perception rule): streaming clients receive the copy as a normal assistant message (no client change, no raw transport error); JSON clients get 429 + code 'ai_quota_exhausted' + resetAt. - DailyMessageQuota: N user turns / user / environment / UTC day, backed by the new ai_usage_daily platform object via IDataEngine. Counter semantics (lost update under-counts one turn) — entitlement gate, not billing. Fail-open on store errors: a broken counter never takes chat down. - Plan-tier policies (free vs pro) later implement AgentChatQuota and plug into the same gate; route and counter unchanged. 10 new tests (quota policy + route gate, both modes, env scoping, fail-open). Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent 17ffc74 commit 57abe7a

6 files changed

Lines changed: 462 additions & 3 deletions

File tree

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
// Copyright (c) 2025 ObjectStack. Licensed under the Apache-2.0 license.
2+
3+
import { describe, it, expect, vi } from 'vitest';
4+
import type { IDataEngine, Logger } from '@objectstack/spec/contracts';
5+
import { DailyMessageQuota, type AgentChatQuota } from '../quota/agent-chat-quota.js';
6+
import { buildAgentRoutes } from '../routes/agent-routes.js';
7+
import { AIService } from '../ai-service.js';
8+
import { MemoryLLMAdapter } from '../adapters/memory-adapter.js';
9+
import { InMemoryConversationService } from '../conversation/in-memory-conversation-service.js';
10+
import type { AgentRuntime } from '../agent-runtime.js';
11+
12+
const silentLogger: Logger = {
13+
debug: vi.fn(),
14+
info: vi.fn(),
15+
warn: vi.fn(),
16+
error: vi.fn(),
17+
fatal: vi.fn(),
18+
};
19+
20+
const FIXED_NOW = () => new Date('2026-06-12T10:00:00.000Z');
21+
22+
function mockDataEngine(rows: Record<string, { messages: number }> = {}): IDataEngine & {
23+
inserted: unknown[];
24+
updated: unknown[];
25+
} {
26+
const inserted: unknown[] = [];
27+
const updated: unknown[] = [];
28+
return {
29+
inserted,
30+
updated,
31+
findOne: vi.fn(async (_obj: string, q?: { where?: { id?: string } }) => {
32+
const id = q?.where?.id;
33+
return id && rows[id] ? { id, ...rows[id] } : null;
34+
}),
35+
insert: vi.fn(async (_obj: string, data: unknown) => {
36+
inserted.push(data);
37+
return data;
38+
}),
39+
update: vi.fn(async (_obj: string, data: unknown, opts: unknown) => {
40+
updated.push({ data, opts });
41+
return data;
42+
}),
43+
find: vi.fn(async () => []),
44+
delete: vi.fn(async () => ({})),
45+
count: vi.fn(async () => 0),
46+
aggregate: vi.fn(async () => []),
47+
} as unknown as IDataEngine & { inserted: unknown[]; updated: unknown[] };
48+
}
49+
50+
// ═══════════════════════════════════════════════════════════════════
51+
// DailyMessageQuota
52+
// ═══════════════════════════════════════════════════════════════════
53+
54+
describe('DailyMessageQuota', () => {
55+
const subject = { userId: 'u1', environmentId: 'env1' };
56+
const todayId = '2026-06-12:env1:u1';
57+
58+
it('allows under the limit and reports remaining + resetAt', async () => {
59+
const quota = new DailyMessageQuota(mockDataEngine({ [todayId]: { messages: 3 } }), 30, FIXED_NOW);
60+
const d = await quota.check(subject);
61+
expect(d.allowed).toBe(true);
62+
expect(d.remaining).toBe(26); // 30 - 3 - this turn
63+
expect(d.resetAt).toBe('2026-06-13T00:00:00.000Z');
64+
});
65+
66+
it('refuses at the limit with honest copy (why + when + way out)', async () => {
67+
const quota = new DailyMessageQuota(mockDataEngine({ [todayId]: { messages: 30 } }), 30, FIXED_NOW);
68+
const d = await quota.check(subject);
69+
expect(d.allowed).toBe(false);
70+
expect(d.resetAt).toBe('2026-06-13T00:00:00.000Z');
71+
expect(d.message).toContain('30');
72+
expect(d.message).toContain('恢复');
73+
expect(d.message).toContain('upgrade');
74+
});
75+
76+
it('consume inserts the first turn of the day, then increments', async () => {
77+
const engine = mockDataEngine();
78+
const quota = new DailyMessageQuota(engine, 30, FIXED_NOW);
79+
await quota.consume(subject);
80+
expect(engine.inserted).toEqual([
81+
{ id: todayId, day: '2026-06-12', user_id: 'u1', environment_id: 'env1', messages: 1 },
82+
]);
83+
84+
const engine2 = mockDataEngine({ [todayId]: { messages: 5 } });
85+
const quota2 = new DailyMessageQuota(engine2, 30, FIXED_NOW);
86+
await quota2.consume(subject);
87+
expect(engine2.updated).toEqual([{ data: { messages: 6 }, opts: { where: { id: todayId } } }]);
88+
});
89+
90+
it('fails OPEN when the counter store errors — never blocks chat', async () => {
91+
const engine = mockDataEngine();
92+
(engine.findOne as ReturnType<typeof vi.fn>).mockRejectedValue(new Error('db down'));
93+
const quota = new DailyMessageQuota(engine, 1, FIXED_NOW);
94+
await expect(quota.check(subject)).resolves.toMatchObject({ allowed: true });
95+
await expect(quota.consume(subject)).resolves.toBeUndefined();
96+
});
97+
98+
it('scopes the counter per environment and per user', async () => {
99+
const engine = mockDataEngine({ ['2026-06-12:envA:u1']: { messages: 99 } });
100+
const quota = new DailyMessageQuota(engine, 10, FIXED_NOW);
101+
// Same user, different environment → independent counter.
102+
expect((await quota.check({ userId: 'u1', environmentId: 'envB' })).allowed).toBe(true);
103+
// Different user, same environment → independent counter.
104+
expect((await quota.check({ userId: 'u2', environmentId: 'envA' })).allowed).toBe(true);
105+
// The exhausted pair stays refused.
106+
expect((await quota.check({ userId: 'u1', environmentId: 'envA' })).allowed).toBe(false);
107+
});
108+
});
109+
110+
// ═══════════════════════════════════════════════════════════════════
111+
// Agent chat route × quota gate
112+
// ═══════════════════════════════════════════════════════════════════
113+
114+
function mockAgentRuntime(): AgentRuntime {
115+
return {
116+
loadAgent: vi.fn(async () => ({
117+
name: 'data_chat',
118+
label: 'Assistant',
119+
active: true,
120+
})),
121+
resolveActiveSkills: vi.fn(async () => []),
122+
buildSystemMessages: vi.fn(() => [{ role: 'system', content: 'sys' }]),
123+
buildRequestOptions: vi.fn(() => ({})),
124+
listAgents: vi.fn(async () => []),
125+
} as unknown as AgentRuntime;
126+
}
127+
128+
function chatRoute(quota?: AgentChatQuota) {
129+
const aiService = new AIService({
130+
adapter: new MemoryLLMAdapter(),
131+
conversationService: new InMemoryConversationService(),
132+
});
133+
const routes = buildAgentRoutes(aiService, mockAgentRuntime(), silentLogger, { quota });
134+
const route = routes.find((r) => r.path.endsWith('/chat'));
135+
if (!route) throw new Error('chat route not found');
136+
return route;
137+
}
138+
139+
const chatReq = (over: Record<string, unknown> = {}) => ({
140+
params: { agentName: 'data_chat' },
141+
body: { messages: [{ role: 'user', content: 'hi' }], stream: false, ...over },
142+
user: { userId: 'u1' },
143+
});
144+
145+
describe('agent chat route quota gate', () => {
146+
it('passes through and consumes exactly once when allowed', async () => {
147+
const quota: AgentChatQuota = {
148+
check: vi.fn(async () => ({ allowed: true })),
149+
consume: vi.fn(async () => {}),
150+
};
151+
const res = await chatRoute(quota).handler(chatReq() as never);
152+
expect(res.status).toBe(200);
153+
expect(quota.check).toHaveBeenCalledTimes(1);
154+
expect(quota.consume).toHaveBeenCalledTimes(1);
155+
});
156+
157+
it('returns 429 + stable code on JSON mode when exhausted (nothing consumed)', async () => {
158+
const quota: AgentChatQuota = {
159+
check: vi.fn(async () => ({ allowed: false, message: '额度已用完', resetAt: '2026-06-13T00:00:00.000Z' })),
160+
consume: vi.fn(async () => {}),
161+
};
162+
const res = await chatRoute(quota).handler(chatReq() as never);
163+
expect(res.status).toBe(429);
164+
expect(res.body).toMatchObject({ code: 'ai_quota_exhausted', error: '额度已用完', resetAt: '2026-06-13T00:00:00.000Z' });
165+
expect(quota.consume).not.toHaveBeenCalled();
166+
});
167+
168+
it('streams the refusal as a normal assistant message in stream mode', async () => {
169+
const quota: AgentChatQuota = {
170+
check: vi.fn(async () => ({ allowed: false, message: '今日额度已用完,明日恢复' })),
171+
consume: vi.fn(async () => {}),
172+
};
173+
const res = await chatRoute(quota).handler(chatReq({ stream: true }) as never);
174+
expect(res.status).toBe(200);
175+
expect(res.stream).toBe(true);
176+
let text = '';
177+
for await (const chunk of res.events as AsyncIterable<string>) text += chunk;
178+
expect(text).toContain('今日额度已用完');
179+
expect(text).toContain('"type":"finish"');
180+
expect(quota.consume).not.toHaveBeenCalled();
181+
});
182+
183+
it('no quota wired → unchanged behavior', async () => {
184+
const res = await chatRoute(undefined).handler(chatReq() as never);
185+
expect(res.status).toBe(200);
186+
});
187+
188+
it('forwards the environmentId from chat context to the quota subject', async () => {
189+
const quota: AgentChatQuota = {
190+
check: vi.fn(async () => ({ allowed: true })),
191+
consume: vi.fn(async () => {}),
192+
};
193+
await chatRoute(quota).handler(chatReq({ context: { environmentId: 'env42' } }) as never);
194+
expect(quota.check).toHaveBeenCalledWith({ userId: 'u1', environmentId: 'env42' });
195+
expect(quota.consume).toHaveBeenCalledWith({ userId: 'u1', environmentId: 'env42' });
196+
});
197+
});
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
// Copyright (c) 2025 ObjectStack. Licensed under the Apache-2.0 license.
2+
3+
import { ObjectSchema, Field } from '@objectstack/spec/data';
4+
5+
/**
6+
* ai_usage_daily — per-user, per-day AI chat usage counters
7+
*
8+
* One row per (UTC day, environment, user). The agent chat route increments
9+
* `messages` once per user turn; {@link DailyMessageQuota} reads it to decide
10+
* whether a daily message limit has been reached.
11+
*
12+
* This is a quota counter, not billing-grade metering: increments are
13+
* last-write-wins and a lost update under concurrency only ever under-counts
14+
* by a message — acceptable for an abuse/entitlement gate, not for invoicing.
15+
*
16+
* @namespace ai
17+
*/
18+
export const AiUsageDailyObject = ObjectSchema.create({
19+
name: 'ai_usage_daily',
20+
label: 'AI Daily Usage',
21+
pluralLabel: 'AI Daily Usage',
22+
icon: 'gauge',
23+
isSystem: true,
24+
description: 'Per-user daily AI chat usage counters (quota enforcement)',
25+
26+
fields: {
27+
id: Field.text({
28+
label: 'Usage ID',
29+
required: true,
30+
readonly: true,
31+
description: 'Deterministic key: <day>:<environment|->:<user>',
32+
}),
33+
34+
day: Field.text({
35+
label: 'Day (UTC)',
36+
required: true,
37+
maxLength: 10,
38+
description: 'UTC calendar day, YYYY-MM-DD',
39+
}),
40+
41+
user_id: Field.text({
42+
label: 'User ID',
43+
required: true,
44+
maxLength: 255,
45+
}),
46+
47+
environment_id: Field.text({
48+
label: 'Environment ID',
49+
required: false,
50+
maxLength: 255,
51+
}),
52+
53+
messages: Field.number({
54+
label: 'Messages',
55+
required: true,
56+
description: 'User chat turns consumed this day',
57+
}),
58+
},
59+
});

packages/services/service-ai/src/objects/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,4 @@ export { AiTraceObject } from './ai-trace.object.js';
1212
export { AiPendingActionObject } from './ai-pending-action.object.js';
1313
export { AiEvalCaseObject } from './ai-eval-case.object.js';
1414
export { AiEvalRunObject } from './ai-eval-run.object.js';
15+
export { AiUsageDailyObject } from './ai-usage-daily.object.js';

packages/services/service-ai/src/plugin.ts

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@ import { buildToolRoutes } from './routes/tool-routes.js';
1414
import { buildPendingActionRoutes } from './routes/pending-action-routes.js';
1515
import { buildEvalRoutes } from './routes/eval-routes.js';
1616
import { ObjectQLConversationService } from './conversation/objectql-conversation-service.js';
17-
import { AiConversationObject, AiMessageObject, AiPendingActionObject, AiTraceObject, AiEvalCaseObject, AiEvalRunObject } from './objects/index.js';
17+
import { AiConversationObject, AiMessageObject, AiPendingActionObject, AiTraceObject, AiEvalCaseObject, AiEvalRunObject, AiUsageDailyObject } from './objects/index.js';
18+
import { DailyMessageQuota, type AgentChatQuota } from './quota/agent-chat-quota.js';
1819
import { AiTraceView, AiMessageView, AiPendingActionView, AiEvalCaseView, AiEvalRunView } from './views/index.js';
1920
import { EvalRunner } from './eval/index.js';
2021
import { registerDataTools } from './tools/data-tools.js';
@@ -593,7 +594,7 @@ export class AIServicePlugin implements Plugin {
593594
type: 'plugin',
594595
scope: 'system',
595596
namespace: 'ai',
596-
objects: [AiConversationObject, AiMessageObject, AiTraceObject, AiPendingActionObject, AiEvalCaseObject, AiEvalRunObject],
597+
objects: [AiConversationObject, AiMessageObject, AiTraceObject, AiPendingActionObject, AiEvalCaseObject, AiEvalRunObject, AiUsageDailyObject],
597598
views: [AiTraceView, AiMessageView, AiPendingActionView, AiEvalCaseView, AiEvalRunView],
598599
});
599600

@@ -857,7 +858,26 @@ export class AIServicePlugin implements Plugin {
857858
if (metadataService) {
858859
const skillRegistry = new SkillRegistry(metadataService);
859860
const agentRuntime = new AgentRuntime(metadataService, skillRegistry);
860-
const agentRoutes = buildAgentRoutes(this.service, agentRuntime, ctx.logger);
861+
862+
// ── Optional per-turn chat quota (ADR-0040 §5, perception rule) ──
863+
// Mechanism only: the deployment opts in by setting
864+
// AI_DAILY_USER_MESSAGES=<N>. Unset/invalid → no quota, unchanged
865+
// behavior. Plan-tier policies (free vs pro) wire a richer
866+
// AgentChatQuota here later; the gate and counter do not change.
867+
let chatQuota: AgentChatQuota | undefined;
868+
const quotaRaw = process.env.AI_DAILY_USER_MESSAGES;
869+
const quotaLimit = quotaRaw ? Number.parseInt(quotaRaw, 10) : NaN;
870+
if (Number.isFinite(quotaLimit) && quotaLimit > 0) {
871+
const quotaDataEngine = ctx.getService<IDataEngine>('data');
872+
if (quotaDataEngine && typeof quotaDataEngine.findOne === 'function') {
873+
chatQuota = new DailyMessageQuota(quotaDataEngine, quotaLimit);
874+
ctx.logger.info(`[AI] Daily chat quota enabled (${quotaLimit} user turns/user/day)`);
875+
} else {
876+
ctx.logger.warn('[AI] AI_DAILY_USER_MESSAGES set but IDataEngine unavailable — quota disabled');
877+
}
878+
}
879+
880+
const agentRoutes = buildAgentRoutes(this.service, agentRuntime, ctx.logger, { quota: chatQuota });
861881
routes.push(...agentRoutes);
862882
ctx.logger.info(`[AI] Agent routes registered (${agentRoutes.length} routes)`);
863883

0 commit comments

Comments
 (0)