Jyotisha/frontend/tests/database-consultation-free-completion.test.ts

import assert from "node:assert/strict";
import { spawnSync } from "node:child_process";
import { fileURLToPath } from "node:url";
import test from "node:test";
import pg from "pg";
import { startPostgresFixture } from "./helpers/postgres-fixture.ts";

const runner = fileURLToPath(new URL("../scripts/db-migrate.mjs", import.meta.url));
const docker = spawnSync("docker", ["version", "--format", "{{.Server.Version}}"], { stdio: "ignore" }).status === 0;
const session = "99999999-9999-4999-8999-999999999971";
const otherSession = "99999999-9999-4999-8999-999999999972";

test("free completion atomically persists, refunds exact credits, records actual cost and guards trust/idempotency", { skip: docker ? false : "docker unavailable" }, async () => {
  const fixture = startPostgresFixture();
  try {
    const migration = spawnSync(process.execPath, [runner], { encoding: "utf8", env: { ...process.env, SCHEMA_DATABASE_URL: fixture.connectionUrl("schema_owner", "schema-owner-test-password") } });
    assert.equal(migration.status, 0, migration.stderr);
    fixture.psqlAs("identity_runtime", "identity-runtime-test-password", `insert into identity.users(name,email,email_verified,email_verified_at) values ('Fictional Free','free@example.com',true,now()),('Fictional Other','other-free@example.com',true,now());`);
    const user = fixture.psql("select id from identity.users where email='free@example.com'");
    const other = fixture.psql("select id from identity.users where email='other-free@example.com'");
    fixture.psql(`update public.profiles set credits=20 where id='${user}';
      insert into public.chat_sessions(id,user_id,title,theme,model_id,messages,session_type,updated_at) values
      ('${session}','${user}','test','general','test-model','[]','consultation',now()),
      ('${otherSession}','${other}','test','general','test-model','[]','consultation',now());`);
    const reserve = (id: string) => fixture.psql(`select success from public.reserve_consultation_usage('${user}','${id}','${session}','test-model',3)`);
    const balance = () => fixture.psql(`select credits from public.profiles where id='${user}'`);
    const completeSql = (id: string, reply = "你好", owner = user, sid = session, usage = '{"inputTokens":23,"outputTokens":7,"costMicrousd":91,"actualModelId":"test-model","modelConfigVersion":1}') =>
      `select success::text || ':' || coalesce(error_code,'null') from public.complete_consultation_free('${owner}','${id}','${sid}','{"role":"assistant","text":"${reply}","responseKind":"smalltalk"}', '${usage}')`;
    for (const role of ["authenticated", "anon", "app_runtime"]) {
      assert.equal(fixture.psql(`select has_function_privilege('${role}','public.complete_consultation_free(uuid,text,uuid,jsonb,jsonb)','execute')`), "f");
    }
    assert.equal(fixture.psql("select has_function_privilege('service_role','public.complete_consultation_free(uuid,text,uuid,jsonb,jsonb)','execute')"), "t");
    assert.equal(reserve("free-1"), "t"); assert.equal(balance(), "17");
    assert.equal(fixture.psql(completeSql("free-1", "你好", other)), "f:request_missing");
    assert.equal(fixture.psql(completeSql("free-1", "你好", user, otherSession)), "f:request_missing");
    assert.equal(balance(), "17");
    assert.equal(fixture.psql(`set role service_role; ${completeSql("free-1")}`).split("\n").at(-1), "true:null");
    assert.equal(balance(), "20");
    assert.equal(fixture.psql(completeSql("free-1")), "true:null");
    assert.equal(fixture.psql(completeSql("free-1", "再见")), "f:response_conflict");
    assert.equal(balance(), "20");
    assert.equal(fixture.psql(`select jsonb_array_length(messages) from public.chat_sessions where id='${session}'`), "1");
    assert.equal(fixture.psql(`select status || ':' || (response_message->>'responseKind') from public.consultation_requests where user_id='${user}' and request_id='free-1'`), "completed:smalltalk");
    assert.equal(fixture.psql(`select input_tokens || ':' || output_tokens || ':' || cost_microusd || ':' || (metadata->>'freeCompletion') from public.usage_ledger where user_id='${user}' and request_id='free-1'`), "23:7:91:true");
    assert.equal(fixture.psql(`select count(*) || ':' || sum(amount) from public.credit_transactions where user_id='${user}' and request_id='free-1' and transaction_type='refund'`), "1:3");
    assert.equal(fixture.psql(`select status from public.usage_reservations where user_id='${user}' and request_id='free-1'`), "released");
    assert.equal(fixture.psql(`select success::text || ':' || coalesce(error_code,'null') from public.cancel_consultation_credit('${user}','free-1')`), "f:request_completed");

    assert.equal(reserve("cancelled"), "t");
    fixture.psql(`select * from public.cancel_consultation_credit('${user}','cancelled')`);
    assert.equal(fixture.psql(completeSql("cancelled")), "f:request_cancelled"); assert.equal(balance(), "20");

    assert.equal(reserve("paid"), "t");
    fixture.psql(`select * from public.complete_consultation_response('${user}','paid','${session}','{"role":"assistant","text":"paid consultation"}','{}')`);
    assert.equal(fixture.psql(completeSql("paid")), "f:response_conflict"); assert.equal(balance(), "17");
    assert.equal(fixture.psql(`select count(*) from public.credit_transactions where user_id='${user}' and request_id='paid' and transaction_type='refund'`), "0");

    assert.equal(reserve("bad-usage"), "t");
    assert.throws(() => fixture.psql(completeSql("bad-usage", "你好", user, session, '{"inputTokens":-1}')));
    assert.equal(balance(), "14");
    assert.equal(fixture.psql(`select status from public.consultation_requests where user_id='${user}' and request_id='bad-usage'`), "reserved");
    assert.equal(fixture.psql(`select jsonb_array_length(messages) from public.chat_sessions where id='${session}'`), "2");
    // A successful retry after a transaction rollback refunds once and persists once.
    assert.equal(fixture.psql(completeSql("bad-usage")), "true:null"); assert.equal(balance(), "17");

    // A subscription-shaped reservation has no credit debit/refund, and releases
    // the same quota counter (reserved/completed rows) while preserving cost.
    assert.equal(reserve("subscription-free"), "t");
    fixture.psql(`update public.profiles set credits=credits+3 where id='${user}';
      update public.usage_reservations set source='subscription',credit_amount=0 where user_id='${user}' and request_id='subscription-free';`);
    assert.equal(fixture.psql(completeSql("subscription-free")), "true:null");
    assert.equal(balance(), "17");
    assert.equal(fixture.psql(`select count(*) from public.credit_transactions where user_id='${user}' and request_id='subscription-free' and transaction_type='refund'`), "0");
    assert.equal(fixture.psql(`select count(*) from public.usage_reservations where user_id='${user}' and request_id='subscription-free' and status in ('reserved','completed')`), "0");

    // Genuine concurrent connections: cancellation/free completion share one lock.
    assert.equal(reserve("race"), "t");
    const clients = [new pg.Client({ connectionString: fixture.connectionUrl("schema_owner", "schema-owner-test-password") }), new pg.Client({ connectionString: fixture.connectionUrl("schema_owner", "schema-owner-test-password") })];
    try {
      await Promise.all(clients.map((client) => client.connect()));
      await Promise.all([
        clients[0]!.query(completeSql("race")),
        clients[1]!.query(`select * from public.cancel_consultation_credit('${user}','race')`),
      ]);
      assert.equal(balance(), "17");
      assert.equal(fixture.psql(`select count(*) from public.credit_transactions where user_id='${user}' and request_id='race' and transaction_type='refund'`), "1");
      assert.match(fixture.psql(`select status from public.consultation_requests where user_id='${user}' and request_id='race'`), /^(completed|cancelled)$/);
    } finally { await Promise.all(clients.map((client) => client.end())); }
  } finally { fixture.stop(); }
});