|
| 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 | +}); |
0 commit comments