All files / services/crawler-service-worker/src/activities track-usage.js

52.45% Statements 32/61
70% Branches 7/10
66.66% Functions 2/3
52.45% Lines 32/61

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 621x 1x 1x 1x 1x 1x 1x 3x 3x 3x 1x                                                   1x 1x 1x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x 3x         3x 3x 3x 3x 1x 1x  
const http = require('http');
const url = require('url');
 
const CONFIG_SERVICE_URL = process.env.CONFIG_SERVICE_URL || 'http://config-service:8000';
const parsed = url.parse(CONFIG_SERVICE_URL);
const { addInternalKey } = require('../../shared/internalAuth');
 
function internalHeaders() {
  return addInternalKey({ 'Content-Type': 'application/json' }, 'crawler-service-worker');
}
 
async function trackUsage(organizationId, llmCost, ocrCost) {
  return new Promise((resolve, reject) => {
    const body = JSON.stringify({ organizationId, llmCost, ocrCost });
    const req = http.request({
      hostname: parsed.hostname,
      port: parsed.port || 8000,
      path: '/internal/usage/track',
      method: 'POST',
      headers: {
        ...internalHeaders(),
        'Content-Length': Buffer.byteLength(body),
      },
    }, res => {
      let data = '';
      res.on('data', c => data += c);
      res.on('end', () => {
        if (res.statusCode === 200) resolve(JSON.parse(data));
        else reject(new Error(`Track usage failed: ${res.statusCode}`));
      });
    });
    req.on('error', reject);
    req.write(body);
    req.end();
  });
}
 
// Les activités de budget RETOURNENT uniquement l'état (pas de throw).
// La politique d'application (blocage) est gérée par le workflow via ApplicationFailure.
async function checkLlmBudget(organizationId) {
  return new Promise((resolve, reject) => {
    http.get({
      hostname: parsed.hostname,
      port: parsed.port || 8000,
      path: `/internal/usage/can-consume?organizationId=${organizationId}&type=llm`,
      headers: internalHeaders(),
    }, res => {
      let data = '';
      res.on('data', c => data += c);
      res.on('end', () => {
        if (res.statusCode === 200) resolve({ allowed: true });
        else if (res.statusCode === 429) {
          try { const b = JSON.parse(data); resolve({ allowed: false, used: b.used, limit: b.limit }); }
          catch { resolve({ allowed: false }); }
        } else reject(new Error(`Check budget failed: ${res.statusCode}`));
      });
    }).on('error', reject);
  });
}
 
module.exports = { trackUsage, checkLlmBudget };