QodeAssist/sources/Session/Session.cpp

// Copyright (C) 2024-2026 Petr Mironychev
// SPDX-License-Identifier: GPL-3.0-or-later

#include "Session.hpp"

#include <LLMQore/BaseClient.hpp>

#include <QDebug>
#include <QJsonObject>
#include <QUrl>

#include <algorithm>

#include "Agent.hpp"
#include "AgentConfig.hpp"
#include "ContextData.hpp"
#include "Message.hpp"
#include "PluginBlocks.hpp"
#include "PromptTemplate.hpp"
#include "Provider.hpp"
#include "ResponseRouter.hpp"
#include "SystemPromptBuilder.hpp"

namespace QodeAssist {

namespace {

QString roleToLegacyString(Message::Role role)
{
    switch (role) {
    case Message::Role::System: return QStringLiteral("system");
    case Message::Role::User: return QStringLiteral("user");
    case Message::Role::Assistant: return QStringLiteral("assistant");
    }
    return QStringLiteral("user");
}

} // namespace

Session::Session(QObject *parent)
    : QObject(parent)
    , m_history(new ConversationHistory(this))
    , m_systemPrompt(new SystemPromptBuilder(this))
{
    m_invalidReason = QStringLiteral("Session: no agent attached");
}

Session::Session(Agent *agent, QObject *parent)
    : Session(agent, /*externalHistory=*/nullptr, parent)
{}

Session::Session(Agent *agent, ConversationHistory *externalHistory, QObject *parent)
    : QObject(parent)
    , m_agent(agent)
    , m_history(externalHistory ? externalHistory : new ConversationHistory(this))
    , m_systemPrompt(new SystemPromptBuilder(this))
{
    if (!m_agent) {
        m_invalidReason = QStringLiteral("Session: agent is null");
        return;
    }
    m_agent->setParent(this);

    if (!m_agent->isValid()) {
        m_invalidReason = m_agent->invalidReason();
        return;
    }

    auto *provider = m_agent->provider();
    auto *client = provider ? provider->client() : nullptr;
    if (!client) {
        m_invalidReason = QStringLiteral("Session: provider has no live client");
        return;
    }
    if (!m_agent->promptTemplate()) {
        m_invalidReason
            = QStringLiteral("Session: agent has no inline prompt template");
        return;
    }

    m_router = new ResponseRouter(client, m_history, this);
    connect(m_router, &ResponseRouter::event, this, &Session::onRouterEvent);

    m_systemPrompt->setLayer(QStringLiteral("agent.role"), m_agent->config().role);
}

Session::~Session()
{
    if (isInFlight())
        cancel();
}

bool Session::isValid() const noexcept
{
    return m_invalidReason.isEmpty();
}

QString Session::invalidReason() const
{
    return m_invalidReason;
}

bool Session::isInFlight() const noexcept
{
    return !m_inFlight.isEmpty();
}

void Session::setContentLoader(ContentLoader loader)
{
    m_contentLoader = std::move(loader);
}

void Session::setContextBindings(Templates::ContextRenderer::Bindings bindings)
{
    m_contextBindings = std::move(bindings);
}

QString Session::renderAgentContext() const
{
    if (!m_agent)
        return {};
    const auto &cfg = m_agent->config();
    if (cfg.context.isEmpty())
        return {};
    QString err;
    QString rendered = Templates::ContextRenderer::render(cfg.context, m_contextBindings, &err);
    if (!err.isEmpty())
        qWarning("[QodeAssist] agent.context render failed: %s", qUtf8Printable(err));
    return rendered;
}

LLMQore::RequestID Session::sendText(const QString &text)
{
    std::vector<std::unique_ptr<LLMQore::ContentBlock>> blocks;
    if (!text.isEmpty())
        blocks.push_back(std::make_unique<LLMQore::TextContent>(text));
    return send(std::move(blocks));
}

LLMQore::RequestID Session::send(
    std::vector<std::unique_ptr<LLMQore::ContentBlock>> userBlocks,
    std::optional<bool> toolsOverride)
{
    if (!isValid() || userBlocks.empty())
        return {};
    if (!m_history)
        return {};

    if (isInFlight())
        cancel();

    Message msg(Message::Role::User);
    for (auto &b : userBlocks)
        msg.appendBlock(std::move(b));
    m_history->append(std::move(msg));

    return dispatch(toolsOverride);
}

void Session::cancel()
{
    if (m_inFlight.isEmpty())
        return;

    const auto id = m_inFlight;
    m_inFlight.clear();
    if (m_router)
        m_router->endRequest();
    if (m_agent && m_agent->provider())
        m_agent->provider()->cancelRequest(id);
    emit failed(id, QStringLiteral("Cancelled by user"));
}

LLMQore::RequestID Session::sendCompletion(Templates::ContextData ctx)
{
    if (!isValid())
        return {};
    if (isInFlight())
        cancel();

    if (m_history)
        m_history->clear();

    auto *provider = m_agent->provider();
    auto *tmpl = m_agent->promptTemplate();
    const auto &cfg = m_agent->config();

    QJsonObject payload{{QStringLiteral("model"), cfg.model}};
    if (!provider->prepareRequest(payload, tmpl, ctx, /*tools=*/false, /*thinking=*/false))
        return {};

    const auto id = provider->sendRequest(QUrl(provider->url()), payload, cfg.endpoint);
    if (id.isEmpty())
        return {};

    m_inFlight = id;
    if (m_router)
        m_router->beginRequest(id);
    emit started(id);
    return id;
}

LLMQore::RequestID Session::dispatch(std::optional<bool> toolsOverride)
{
    auto *provider = m_agent->provider();
    auto *tmpl = m_agent->promptTemplate();
    const auto &cfg = m_agent->config();

    const QString renderedContext = renderAgentContext();
    if (renderedContext.isEmpty())
        m_systemPrompt->clearLayer(QStringLiteral("agent.context"));
    else
        m_systemPrompt->setLayer(QStringLiteral("agent.context"), renderedContext);

    Templates::ContextData ctx = toLegacyContext();
    QJsonObject payload{{QStringLiteral("model"), cfg.model}};

    const bool tools = toolsOverride.value_or(cfg.enableTools);
    if (!provider->prepareRequest(payload, tmpl, ctx, tools, cfg.enableThinking))
        return {};

    const auto id = provider->sendRequest(QUrl(provider->url()), payload, cfg.endpoint);
    if (id.isEmpty())
        return {};

    m_inFlight = id;
    if (m_router)
        m_router->beginRequest(id);
    emit started(id);
    return id;
}

Templates::ContextData Session::toLegacyContext() const
{
    if (!m_history)
        return {};
    return buildLegacyContext(m_history->messages(), m_systemPrompt->compose(), m_contentLoader);
}

Templates::ContextData Session::buildLegacyContext(
    const std::vector<Message> &history,
    const QString &systemPrompt,
    const ContentLoader &loader)
{
    using Templates::ContentBlockEntry;
    using Templates::ContextData;
    using LegacyMessage = Templates::Message;

    ContextData ctx;
    if (!systemPrompt.isEmpty())
        ctx.systemPrompt = systemPrompt;

    QSet<QString> resolvedToolUseIds;
    QSet<QString> declaredToolUseIds;
    for (const auto &m : history) {
        for (const auto &blockPtr : m.blocks()) {
            if (auto *tr = dynamic_cast<LLMQore::ToolResultContent *>(blockPtr.get()))
                resolvedToolUseIds.insert(tr->toolUseId());
            if (auto *tu = dynamic_cast<LLMQore::ToolUseContent *>(blockPtr.get()))
                declaredToolUseIds.insert(tu->id());
        }
    }

    QVector<LegacyMessage> hist;

    for (const auto &m : history) {
        QVector<ContentBlockEntry> blockEntries;

        for (const auto &blockPtr : m.blocks()) {
            auto *block = blockPtr.get();
            if (!block)
                continue;

            if (auto *t = dynamic_cast<LLMQore::TextContent *>(block)) {
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::Text;
                e.text = t->text();
                blockEntries.append(std::move(e));
            } else if (auto *img = dynamic_cast<LLMQore::ImageContent *>(block)) {
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::Image;
                e.imageData = img->data();
                e.mediaType = img->mediaType();
                e.isImageUrl
                    = (img->sourceType() == LLMQore::ImageContent::ImageSourceType::Url);
                blockEntries.append(std::move(e));
            } else if (auto *si = dynamic_cast<StoredImageContent *>(block)) {
                if (!loader)
                    continue;
                const QString base64 = loader(si->storedPath());
                if (base64.isEmpty())
                    continue;
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::Image;
                e.imageData = base64;
                e.mediaType = si->mediaType();
                e.isImageUrl = false;
                blockEntries.append(std::move(e));
            } else if (auto *sa = dynamic_cast<StoredAttachmentContent *>(block)) {
                if (!loader)
                    continue;
                const QString text = loader(sa->storedPath());
                if (text.isEmpty())
                    continue;
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::Text;
                e.text = QStringLiteral("File: %1\n```\n%2\n```")
                             .arg(sa->fileName(), text);
                blockEntries.append(std::move(e));
            } else if (auto *th = dynamic_cast<LLMQore::ThinkingContent *>(block)) {
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::Thinking;
                e.thinking = th->thinking();
                e.signature = th->signature();
                blockEntries.append(std::move(e));
            } else if (auto *rth = dynamic_cast<LLMQore::RedactedThinkingContent *>(block)) {
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::RedactedThinking;
                e.signature = rth->signature();
                blockEntries.append(std::move(e));
            } else if (auto *tu = dynamic_cast<LLMQore::ToolUseContent *>(block)) {
                if (!resolvedToolUseIds.contains(tu->id()))
                    continue;
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::ToolUse;
                e.toolUseId = tu->id();
                e.toolName = tu->name();
                e.toolInput = tu->input();
                blockEntries.append(std::move(e));
            } else if (auto *tr = dynamic_cast<LLMQore::ToolResultContent *>(block)) {
                if (!declaredToolUseIds.contains(tr->toolUseId()))
                    continue;
                ContentBlockEntry e;
                e.kind = ContentBlockEntry::Kind::ToolResult;
                e.toolUseId = tr->toolUseId();
                e.result = tr->result();
                blockEntries.append(std::move(e));
            }
        }

        if (blockEntries.isEmpty())
            continue;

        const bool hasNonThinking = std::any_of(
            blockEntries.begin(), blockEntries.end(), [](const ContentBlockEntry &e) {
                return e.kind != ContentBlockEntry::Kind::Thinking
                       && e.kind != ContentBlockEntry::Kind::RedactedThinking;
            });
        if (!hasNonThinking)
            continue;

        LegacyMessage lm;
        lm.role = roleToLegacyString(m.role());
        lm.blocks = std::move(blockEntries);
        hist.append(std::move(lm));
    }

    if (!hist.isEmpty())
        ctx.history = std::move(hist);

    return ctx;
}

void Session::onRouterEvent(const ResponseEvent &ev)
{
    if (m_inFlight.isEmpty())
        return; // stale events after cancel

    emit event(ev);

    if (ev.kind() == ResponseEvent::Kind::MessageStop) {
        const auto *stop = ev.as<ResponseEvents::MessageStop>();
        const QString reason = stop ? stop->stopReason : QString();
        const auto id = m_inFlight;
        m_inFlight.clear();
        emit finished(id, reason);
    } else if (ev.kind() == ResponseEvent::Kind::Error) {
        const auto *err = ev.as<ResponseEvents::Error>();
        const QString msg = err ? err->message : QStringLiteral("unknown error");
        const auto id = m_inFlight;
        m_inFlight.clear();
        emit failed(id, msg);
    }
}

} // namespace QodeAssist