tf_docs/orchestrator.js

'use strict';

/**
 * orchestrator.js — связующий обработчик POST /<stand>/_api/assist.
 *
 * Склеивает пайплайн:
 *   validateInput → rate-limit → кэш/дедуп → этап 1 → поиск/объединение → этап 2 → ответ.
 *
 * Реализует решения брифа:
 *   - ASSIST_TOTAL_DEADLINE_MS = 40000 (общий дедлайн запроса);
 *   - повтор этапа 1 только если запас времени > 15 с;
 *   - req.on('close') → прерывание LLM через AbortController;
 *   - семафор ASSIST_MAX_CONCURRENT (сверх → 503 + Retry-After);
 *   - объединение одинаковых запросов (ключ = stand + нормализованный запрос);
 *   - ответ 503 {error:"busy"} при перегрузке.
 */

const ENV = require('./assist.js').ENV;
const { validateInput, createRateLimiter, createSemaphore, normalizeQuery, jsonError } = require('./assist.js');
const { loadIndex, search } = require('./search.js');
const { stage1ExtractResourceType } = require('./stage1.js');
const { mergeSearchResults, buildContext, collectSources, stage2Generate } = require('./stage2.js');

const HTTP = require('http');
const HTTPS = require('https');

const TOTAL_DEADLINE_MS = ENV.ASSIST_TOTAL_DEADLINE_MS || 40000;
const RETRY_STAGE1_MIN_BUDGET_MS = 15000;
const STAND_RE = /^[A-Za-z0-9._-]+$/;
const INDEX_TTL_MS = Number(process.env.ASSIST_INDEX_TTL_MS) || 600000;

// Кэш индекса по stand (в памяти): экономит HTTP-запрос и парсинг на каждом вопросе.
const indexCache = new Map();

/**
 * Загружает search_index.json с VM_DOCS_BASE_URL/<stand>/search/search_index.json.
 * Мемоизирует результат по stand с TTL. При сбое возвращает свежий кэш, если есть.
 *
 * @param {string} stand
 * @param {string} vmDocsBaseUrl
 * @returns {Promise<object|null>} - index из loadIndex() или null.
 */
async function fetchIndex(stand, vmDocsBaseUrl) {
    const now = Date.now();
    const cached = indexCache.get(stand);
    if (cached && (now - cached.at) < INDEX_TTL_MS) {
        return cached.index;
    }
    try {
        const url = `${vmDocsBaseUrl.replace(/\/+$/, '')}/${stand}/search/search_index.json`;
        const mod = url.startsWith('https:') ? HTTPS : HTTP;
        const text = await httpGet(mod, url);
        const index = loadIndex(text);
        indexCache.set(stand, { index, at: now });
        return index;
    } catch (e) {
        // Свежий кэш лучше ошибки: продолжаем работу на устаревшем индексе.
        if (cached) return cached.index;
        return null;
    }
}

/** Простой GET по http/https с возвратом тела. */
function httpGet(mod, url) {
    return new Promise((resolve, reject) => {
        const req = mod.get(url, (res) => {
            if (res.statusCode !== 200) {
                res.resume();
                reject(new Error('index status ' + res.statusCode));
                return;
            }
            let data = '';
            res.on('data', (c) => { data += c; });
            res.on('end', () => resolve(data));
        });
        req.on('error', reject);
        req.setTimeout(10000, () => req.destroy(new Error('index timeout')));
    });
}

/**
 * Главный обработчик POST /<stand>/_api/assist.
 *
 * @param {import('http').IncomingMessage} req
 * @param {import('http').ServerResponse} res
 * @param {object} config
 * @param {string} config.stand
 * @param {string} config.vmDocsBaseUrl
 * @param {object} [config.shared] - разделяемые состояния (семафор, кэш).
 */
async function handleAssist(req, res, config) {
    const stand = config.stand;
    const vmDocsBaseUrl = config.vmDocsBaseUrl;

    // 1) Валидация входа.
    const input = await validateInput(req);
    if (!input.ok) {
        jsonError(res, statusForCode(input.code), input.code, input.message);
        return;
    }
    const question = input.question;
    const normKey = `${stand}\u0000${normalizeQuery(question)}`;

    // 2) Rate-limit (по IP).
    const ip = clientIp(req);
    const limiter = config.shared.limiter;
    const lim = limiter.check(ip);
    if (!lim.ok) {
        jsonError(res, 429, 'rate_limited', 'too many requests', { 'Retry-After': String(lim.retryAfterSec) });
        return;
    }

    // 3) Кэш.
    const cache = config.shared.cache;
    if (cache.has(normKey)) {
        sendJson(res, 200, cache.get(normKey));
        return;
    }

    // 4) Дедуп: одинаковый запрос ждёт результат первого.
    if (config.shared.inflight.has(normKey)) {
        await waitForInflight(normKey, config.shared, res);
        return;
    }

    // 5) Семафор.
    const sem = config.shared.semaphore;
    if (!sem.tryAcquire()) {
        jsonError(res, 503, 'busy', 'server is busy, retry later', { 'Retry-After': '5' });
        return;
    }

    // Регистрируем как «держащий вызов».
    const inflightEntry = { holderCount: 1, waiters: [] };
    config.shared.inflight.set(normKey, inflightEntry);

    const controller = new AbortController();
    const start = Date.now();
    let holderActive = true;

    // Общий бюджет времени: по истечении прерываем вызовы ЛЛМ.
    const deadlineTimer = setTimeout(() => controller.abort(), TOTAL_DEADLINE_MS);

    // Прерывание LLM при закрытии соединения клиентом.
    const onClose = () => {
        inflightEntry.holderCount -= 1;
        if (inflightEntry.holderCount <= 0 && holderActive) {
            controller.abort();
        }
    };
    req.on('close', onClose);

    const release = () => {
        clearTimeout(deadlineTimer);
        try { sem.release(); } catch { }
        req.removeListener('close', onClose);
        config.shared.inflight.delete(normKey);
        // пробуждаем ожидающих результатом (результат упадёт — отдаём 502/503 им).
        for (const w of inflightEntry.waiters) w.res(res);
    };

    try {
        const result = await runPipeline(question, stand, vmDocsBaseUrl, controller, start);
        holderActive = false;
        const body = { answer: result.answer, sources: result.sources, meta: result.meta };
        cache.set(normKey, body);
        sendJson(res, 200, body);
        for (const w of inflightEntry.waiters) w.succeed(body);
        logAssist('ok', stand, question.length, result.meta, Date.now() - start);
    } catch (err) {
        holderActive = false;
        const code = err && err.code ? err.code : 'internal_error';
        const status = mapPipelineError(code);
        for (const w of inflightEntry.waiters) w.fail(status, code);
        jsonError(res, status, code, err && err.message ? err.message : 'internal error');
        logAssist('error', stand, question.length, { error: code }, Date.now() - start);
    } finally {
        release();
    }
}

/** Однострочный JSON-лог без текста запроса (7.11 брифа). */
function logAssist(status, stand, questionLen, meta, totalMs) {
    try {
        const line = JSON.stringify({
            t: new Date().toISOString(),
            status,
            stand,
            question_len: questionLen,
            total_ms: totalMs,
            ...meta,
        });
        process.stdout.write(line + '\n');
    } catch { /* лог не должен ронять */ }
}

/** Возвращает IP клиента (7.3: X-Forwarded-For только для доверенных proxy). */
function clientIp(req) {
    const trusted = (ENV.ASSIST_TRUSTED_PROXY_CIDR || '').trim();
    const remote = req.socket && req.socket.remoteAddress ? req.socket.remoteAddress : '';
    if (trusted && remoteIsInCidr(remote, trusted)) {
        const xff = req.headers['x-forwarded-for'];
        if (xff) {
            const parts = String(xff).split(',').map((s) => s.trim()).filter(Boolean);
            if (parts.length) return parts[parts.length - 1];
        }
    }
    return remote || 'unknown';
}

/** Простейшая проверка IP == CIDR (точное совпадение или префикс). v1: только равенство/CIDR /32. */
function remoteIsInCidr(ip, cidr) {
    // Нормализация IPv6-префикса ::ffff:1.2.3.4 → 1.2.3.4
    const norm = String(ip).replace(/^::ffff:/, '');
    const parts = cidr.split(',');
    for (const p of parts) {
        const t = p.trim();
        if (!t) continue;
        if (!t.includes('/')) {
            if (t === norm) return true;
            continue;
        }
        const [base, bitsStr] = t.split('/');
        if (t === norm) return true;
        const bits = Number.parseInt(bitsStr, 10);
        if (Number.isFinite(bits) && bits >= 32 && base === norm) return true;
    }
    return false;
}

/** Основной пайплайн (этап 1 → поиск → объединение → этап 2). */
async function runPipeline(question, stand, vmDocsBaseUrl, controller, start) {
    const t1 = Date.now();
    const index = await fetchIndex(stand, vmDocsBaseUrl);
    if (!index) {
        throw { code: 'index_unavailable', message: 'search index unavailable' };
    }
    const resourceTypes = index.listResourceTypes();
    const tIndex = Date.now();

    // Этап 1 (resource_type).
    let resourceType = null;
    let stage1Ok = false;
    const usage = { prompt: 0, completion: 0, total: 0 };
    const budget = TOTAL_DEADLINE_MS - (Date.now() - start);
    if (budget > 0 && resourceTypes.size > 0) {
        const r1 = await stage1ExtractResourceType(question, {
            resourceTypesSet: resourceTypes,
            controller,
            usageRef: usage,
        });
        stage1Ok = r1.ok;
        if (r1.ok && r1.resourceType) resourceType = r1.resourceType;
    }
    const tStage1 = Date.now();

    // Поиск по токенам исходного запроса — всегда (параллельно уже можно, но для простоты
    // объединяем в mergeSearchResults, который сам делает поиск по query и по resource_type).
    const merged = mergeSearchResults(index, question, resourceType, { topK: 5 });
    const context = buildContext(merged, 1500);
    const tSearch = Date.now();

    // Этап 2.
    const r2 = await stage2Generate({
        question,
        context,
        baseUrl: ENV.LLM_BASE_URL,
        apiKey: ENV.LLM_API_KEY,
        model: ENV.LLM_MODEL,
        maxTokens: 2048,
        controller,
        usageRef: usage,
    });
    const tStage2 = Date.now();

    if (!r2.ok) {
        throw { code: r2.error, message: r2.error === 'bad_llm_output' ? 'bad llm output' : 'llm error' };
    }

    const sources = collectSources(merged, stand);

    return {
        answer: r2.answer,
        sources,
        meta: {
            model: ENV.LLM_MODEL,
            resource_type: resourceType,
            stage1_ok: stage1Ok,
            elapsed_ms: Date.now() - t1,
            index_ms: tIndex - t1,
            stage1_ms: tStage1 - tIndex,
            search_ms: tSearch - tStage1,
            stage2_ms: tStage2 - tSearch,
            tokens: usage.total,
            prompt_tokens: usage.prompt,
            completion_tokens: usage.completion,
        },
    };
}

/** Централизованный семафор/кэш/inflight-состояния (инициализируется один раз). */
function createSharedState() {
    return {
        semaphore: createSemaphore(ENV.ASSIST_MAX_CONCURRENT),
        limiter: createRateLimiter({}),
        cache: new Map(),
        inflight: new Map(),
    };
}

/** Ожидание результата уже выполняющегося одинакового запроса. */
function waitForInflight(normKey, shared, res) {
    return new Promise((resolvePromise) => {
        const entry = shared.inflight.get(normKey);
        if (!entry) {
            // запрос завершился между проверками — просто повторим обработку
            return handleAssistAgain(res); // fallback (не используется в v1)
        }
        entry.waiters.push({
            succeed(body) { sendJson(res, 200, body); resolvePromise(); },
            fail(status, code) { jsonError(res, status, code, 'upstream failed'); resolvePromise(); },
            res() { },
        });
    });
}

/** Заглушка fallback (не используется). */
function handleAssistAgain() { }

function statusForCode(code) {
    switch (code) {
        case 400: return 400;
        case 413: return 413;
        case 415: return 415;
        default: return 400;
    }
}

function mapPipelineError(code) {
    switch (code) {
        case 'bad_llm_output': return 502;
        case 'index_unavailable': return 503;
        case 'empty_content': return 502;
        case 'aborted': return 499;
        case 'llm_timeout': return 504;
        case 'rate_limited': return 503;
        default: return 502;
    }
}

function sendJson(res, status, obj) {
    const body = JSON.stringify(obj);
    if (!res.headersSent) {
        res.writeHead(status, {
            'Content-Type': 'application/json; charset=utf-8',
            'Content-Length': Buffer.byteLength(body),
        });
    }
    res.end(body);
}

module.exports = {
    handleAssist,
    createSharedState,
    runPipeline,
    fetchIndex,
    clientIp,
    STAND_RE,
};