feat: KV testing (#2061)

src/backend/src/modules/puterai/lib/Streaming.js

@@ -1,38 +1,18 @@
-/**
- * Assign the properties of the override object to the original object,
- * like Object.assign, except properties are ordered so override properties
- * are enumerated first.
- *
- * @param {*} original
- * @param {*} override
- */
-const objectAssignTop = (original, override) => {
-    let o = {
-        ...original,
-        ...override,
-    };
-    o = {
-        ...override,
-        ...original,
-    };
-    return o;
-};
-
 class AIChatConstructStream {
     constructor (chatStream, params) {
         this.chatStream = chatStream;
         if ( this._start ) this._start(params);
     }
     end () {
-        if ( this._end ) this._end();
     }
 }
 
 class AIChatTextStream extends AIChatConstructStream {
     addText (text, extra_content) {
         const json = JSON.stringify({
-            type: 'text', text,
-            ...(extra_content?{extra_content}:{})
+            type: 'text',
+            text,
+            ...(extra_content ? { extra_content } : {}),
         });
         this.chatStream.stream.write(`${json }\n`);
     }
@@ -44,10 +24,10 @@ class AIChatTextStream extends AIChatConstructStream {
         this.chatStream.stream.write(`${json }\n`);
     }
 
-    addExtraContent(extra_content) {
+    addExtraContent (extra_content) {
         const json = JSON.stringify({
             type: 'extra_content',
-            extra_content
+            extra_content,
         });
         this.chatStream.stream.write(`${json }\n`);
     }
@@ -61,18 +41,17 @@ class AIChatToolUseStream extends AIChatConstructStream {
     addPartialJSON (partial_json) {
         this.buffer += partial_json;
     }
-    _end () {
+    end () {
         if ( this.buffer.trim() === '' ) {
             this.buffer = '{}';
         }
         if ( process.env.DEBUG ) console.log('BUFFER BEING PARSED', this.buffer);
-        const str = JSON.stringify(objectAssignTop({
+        const str = JSON.stringify({
+            type: 'tool_use',
             ...this.contentBlock,
             input: JSON.parse(this.buffer),
             ...( !this.contentBlock.text ? { text: '' } : {}),
-        }, {
-            type: 'tool_use',
-        }));
+        });
         this.chatStream.stream.write(`${str }\n`);
     }
 }

src/backend/src/services/MeteringService/MeteringService.test.ts

@@ -3,9 +3,11 @@ import { createTestKernel } from '../../../tools/test.mjs';
 import * as config from '../../config';
 import { Actor } from '../auth/Actor';
 import { DBKVServiceWrapper } from '../repositories/DBKVStore/index.mjs';
-import { GLOBAL_APP_KEY } from './consts.js';
+import { GLOBAL_APP_KEY, PERIOD_ESCAPE } from './consts.js';
 import { MeteringService } from './MeteringService';
 import { MeteringServiceWrapper } from './MeteringServiceWrapper.mjs';
+import { COST_MAPS } from './costMaps/index.js';
+import type { EventService } from '../EventService.js';
 
 describe('MeteringService', async () => {
 
@@ -199,4 +201,21 @@ describe('MeteringService', async () => {
         expect(res.total).toBe(12);
         expect(res['aws-polly:standard:character']).toMatchObject({ cost: 12, units: 10, count: 1 });
     });
+
+    it('applies the configured cost map rate for every usage type', async () => {
+        const usageAmount = 2;
+
+        for ( const [usageType, costPerUnit] of Object.entries(COST_MAPS) ) {
+            const actor = makeActor(`cost-map-user-${usageType.replace(/[^a-zA-Z0-9]/g, '-')}`);
+            const result = await testSubject.meteringService.incrementUsage(actor, usageType, usageAmount);
+            const escapedUsageType = usageType.replace(/\./g, PERIOD_ESCAPE);
+
+            expect(result.total).toBe(costPerUnit * usageAmount);
+            expect(result[escapedUsageType]).toMatchObject({
+                cost: costPerUnit * usageAmount,
+                units: usageAmount,
+                count: 1,
+            });
+        }
+    });
 });

src/backend/src/services/repositories/DBKVStore/DBKVStore.test.ts

@@ -0,0 +1,192 @@
+import { describe, expect, it } from 'vitest';
+import { createTestKernel } from '../../../../tools/test.mjs';
+import * as config from '../../../config';
+import { Actor } from '../../auth/Actor';
+import { MeteringServiceWrapper } from '../../MeteringService/MeteringServiceWrapper.mjs';
+import { DBKVServiceWrapper } from './index.mjs';
+
+describe('DBKVStore', async () => {
+
+    config.load_config({
+        'services': {
+            'database': {
+                path: ':memory:',
+            },
+        },
+    });
+
+    const testKernel = await createTestKernel({
+        serviceMap: {
+            meteringService: MeteringServiceWrapper,
+            'puter-kvstore': DBKVServiceWrapper,
+        },
+        initLevelString: 'init',
+        testCore: true,
+    });
+
+    const kvServiceWrapper = testKernel.services!.get('puter-kvstore') as DBKVServiceWrapper;
+    const kvStore = kvServiceWrapper.kvStore;
+    const su = testKernel.services!.get('su');
+
+    const makeActor = (userId: number | string, appUid?: string) => ({
+        type: {
+            user: { id: userId, uuid: String(userId) },
+            ...(appUid ? { app: { uid: appUid } } : {}),
+        },
+    }) as unknown as Actor;
+
+    it('sets and retrieves values for the current actor context', async () => {
+        const actor = makeActor(1);
+        const key = 'greeting';
+        const value = { hello: 'world' };
+
+        await su.sudo(actor, () => kvStore.set({ key, value }));
+        const stored = await su.sudo(actor, () => kvStore.get({ key }));
+
+        expect(stored).toEqual(value);
+    });
+
+    it('scopes data to the app when provided', async () => {
+        const userId = 2;
+        const actorAppOne = makeActor(userId, 'app-one');
+        const actorAppTwo = makeActor(userId, 'app-two');
+        const key = 'scoped-key';
+
+        await su.sudo(actorAppOne, () => kvStore.set({ key, value: 'one' }));
+        await su.sudo(actorAppTwo, () => kvStore.set({ key, value: 'two' }));
+
+        const fromOne = await su.sudo(actorAppOne, () => kvStore.get({ key }));
+        const fromTwo = await su.sudo(actorAppTwo, () => kvStore.get({ key }));
+
+        expect(fromOne).toBe('one');
+        expect(fromTwo).toBe('two');
+    });
+
+    it('increments nested numeric paths and persists the aggregated totals', async () => {
+        const actor = makeActor(3);
+        const key = 'counter-key';
+
+        const first = await su.sudo(actor, () => kvStore.incr({
+            key,
+            pathAndAmountMap: { 'total': 5, 'nested.count': 2 },
+        }));
+        const second = await su.sudo(actor, () => kvStore.incr({
+            key,
+            pathAndAmountMap: { 'total': 1, 'nested.count': 3 },
+        }));
+
+        expect(first).toMatchObject({ total: 5, nested: { count: 2 } });
+        expect(second).toMatchObject({ total: 6, nested: { count: 5 } });
+
+        const persisted = await su.sudo(actor, () => kvStore.get({ key }));
+        expect(persisted).toMatchObject({ total: 6, nested: { count: 5 } });
+    });
+
+    it('decrements numeric paths via decr and keeps values in sync', async () => {
+        const actor = makeActor(4);
+        const key = 'decr-key';
+
+        await su.sudo(actor, () => kvStore.incr({
+            key,
+            pathAndAmountMap: { total: 5, 'nested.count': 4 },
+        }));
+        const afterDecr = await su.sudo(actor, () => kvStore.decr({
+            key,
+            pathAndAmountMap: { total: 2, 'nested.count': 1 },
+        }));
+
+        expect(afterDecr).toMatchObject({ total: 3, nested: { count: 3 } });
+
+        const persisted = await su.sudo(actor, () => kvStore.get({ key }));
+        expect(persisted).toMatchObject({ total: 3, nested: { count: 3 } });
+    });
+
+    it('deletes keys with del', async () => {
+        const actor = makeActor(5);
+        const key = 'delete-me';
+        await su.sudo(actor, () => kvStore.set({ key, value: 'bye' }));
+
+        const res = await su.sudo(actor, () => kvStore.del({ key }));
+        const value = await su.sudo(actor, () => kvStore.get({ key }));
+
+        expect(res).toBe(true);
+        expect(value).toBeNull();
+    });
+
+    it('lists entries, keys, and values while omitting expired rows', async () => {
+        const actor = makeActor(6);
+        await su.sudo(actor, () => kvStore.set({ key: 'k1', value: 'v1' }));
+        await su.sudo(actor, () => kvStore.set({ key: 'expired', value: 'gone', expireAt: Math.floor(Date.now() / 1000) - 10 }));
+
+        const entries = await su.sudo(actor, () => kvStore.list({ as: 'entries' }));
+        const keys = await su.sudo(actor, () => kvStore.list({ as: 'keys' }));
+        const values = await su.sudo(actor, () => kvStore.list({ as: 'values' }));
+
+        expect(entries).toEqual([{ key: 'k1', value: 'v1' }]);
+        expect(keys).toEqual(['k1']);
+        expect(values).toEqual(['v1']);
+    });
+
+    it('rejects invalid list selector', async () => {
+        const actor = makeActor(7);
+        expect(su.sudo(actor, () => kvStore.list({ as: 'bad' as never })))
+            .rejects;
+    });
+
+    it('returns ordered values for arrays and null for expired keys', async () => {
+        const actor = makeActor(8);
+        const now = Math.floor(Date.now() / 1000);
+
+        await su.sudo(actor, () => kvStore.set({ key: 'a', value: 1 }));
+        await su.sudo(actor, () => kvStore.set({ key: 'b', value: 2, expireAt: now - 5 }));
+        await su.sudo(actor, () => kvStore.set({ key: 'c', value: 3 }));
+
+        const results = await su.sudo(actor, () => kvStore.get({ key: ['c', 'b', 'a'] }));
+
+        expect(results).toEqual([3, null, 1]);
+    });
+
+    it('flush clears all keys for the actor/app combination', async () => {
+        const actor = makeActor(9, 'flush-app');
+        await su.sudo(actor, () => kvStore.set({ key: 'one', value: 1 }));
+        await su.sudo(actor, () => kvStore.set({ key: 'two', value: 2 }));
+
+        const res = await su.sudo(actor, () => kvStore.flush());
+        const remaining = await su.sudo(actor, () => kvStore.list({ as: 'entries' }));
+
+        expect(res).toBe(true);
+        expect(remaining).toEqual([]);
+    });
+
+    it('expireAt and expire set timestamps that cause reads to return null', async () => {
+        const actor = makeActor(10);
+        const keyAt = 'expire-at';
+        const keyTtl = 'expire-ttl';
+
+        await su.sudo(actor, () => kvStore.set({ key: keyAt, value: 'keep' }));
+        await su.sudo(actor, () => kvStore.set({ key: keyTtl, value: 'keep' }));
+
+        await su.sudo(actor, () => kvStore.expireAt({ key: keyAt, timestamp: Math.floor(Date.now() / 1000) - 1 }));
+        await su.sudo(actor, () => kvStore.expire({ key: keyTtl, ttl: -1 }));
+
+        const valAt = await su.sudo(actor, () => kvStore.get({ key: keyAt }));
+        const valTtl = await su.sudo(actor, () => kvStore.get({ key: keyTtl }));
+
+        expect(valAt).toBeNull();
+        expect(valTtl).toBeNull();
+    });
+
+    it('enforces key and value size limits', async () => {
+        const actor = makeActor(11);
+        const oversizedKey = 'a'.repeat((config.kv_max_key_size as number) + 1);
+        const oversizedValue = 'b'.repeat((config.kv_max_value_size as number) + 1);
+
+        await expect(su.sudo(actor, () => kvStore.set({ key: oversizedKey, value: 'x' })))
+            .rejects
+            .toThrow(/key is too large/i);
+
+        await expect(su.sudo(actor, () => kvStore.set({ key: 'ok', value: oversizedValue })))
+            .rejects
+            .toThrow(/value is too large/i);
+    });
+});

src/backend/src/services/repositories/DBKVStore/index.mjs

@@ -3,6 +3,7 @@ import { DB_READ } from '../../database/consts.js';
 import { DBKVStore } from './DBKVStore.js';
 
 export class DBKVServiceWrapper extends BaseService {
+    /** @type {DBKVStore} */
     kvStore = undefined;
     _init () {
         /** @type {DBKVStore} */

src/backend/tools/test.mjs

@@ -57,6 +57,9 @@ class TestLogger {
 * Does not include full service initialization or legacy service support
 */
 export class TestKernel extends AdvancedBase {
+
+    /**@type {Context} */
+    root_context;
     constructor () {
         super();