z-bar-qt/Plugins/ZShell/Llm/segment.cpp

#include "segment.hpp"

#include "markdownparser.hpp"

#include <QCoreApplication>
#include <QDateTime>
#include <QPointer>
#include <QThreadPool>

namespace ZShell::llm {

namespace {
// Re-parse cadence while a segment streams; content between refreshes is
// at most this stale.
constexpr int kMarkdownRefreshMs = 150;
} // namespace

LlmSegment::LlmSegment(Type type, qint64 timestamp, QObject* parent)
: QObject(parent), m_type(type), m_timestamp(timestamp) {
	m_markdownTimer.setInterval(kMarkdownRefreshMs);
	connect(
		&m_markdownTimer, &QTimer::timeout, this, [this]() {
			if (!m_markdownDirty)
				return;
			// A parse may still be running; leave the dirty flag set so
			// it re-parses the newest text when that one completes.
			if (parseMarkdown())
				m_markdownDirty = false;
		});
}

qint64 LlmSegment::elapsedMs() const {
	if (m_startedAt <= 0)
		return 0;
	const qint64 end = m_endedAt > 0 ? m_endedAt
	                                  : QDateTime::currentMSecsSinceEpoch();
	return end - m_startedAt;
}

void LlmSegment::begin() {
	if (m_running)
		return;
	m_running = true;
	m_startedAt = QDateTime::currentMSecsSinceEpoch();
	m_endedAt = 0;
	Q_EMIT runningChanged();
	Q_EMIT elapsedMsChanged();
}

void LlmSegment::close() {
	if (m_running) {
		m_running = false;
		m_endedAt = QDateTime::currentMSecsSinceEpoch();
		Q_EMIT runningChanged();
		Q_EMIT elapsedMsChanged();
	}
	// The segment is final; parse immediately so the UI does not wait
	// out the debounce.
	if (m_type == Type::Content && m_markdownDirty) {
		m_markdownTimer.stop();
		if (parseMarkdown())
			m_markdownDirty = false;
	}
}

void LlmSegment::appendText(const QString& piece) {
	if (piece.isEmpty())
		return;
	m_text += piece;
	Q_EMIT textChanged();
	scheduleMarkdown();
}

void LlmSegment::setText(const QString& value) {
	if (m_text == value)
		return;
	m_text = value;
	Q_EMIT textChanged();
	scheduleMarkdown();
}

void LlmSegment::setName(const QString& value) {
	if (m_name == value)
		return;
	m_name = value;
	Q_EMIT nameChanged();
}

void LlmSegment::setToolCallId(const QString& value) {
	if (m_toolCallId == value)
		return;
	m_toolCallId = value;
	Q_EMIT toolCallIdChanged();
}

void LlmSegment::appendArguments(const QString& piece) {
	if (piece.isEmpty())
		return;
	m_arguments += piece;
	Q_EMIT argumentsChanged();
}

void LlmSegment::setResult(const QString& value) {
	if (m_result == value)
		return;
	m_result = value;
	Q_EMIT resultChanged();
}

void LlmSegment::setStatus(Status value) {
	if (m_status == value)
		return;
	m_status = value;
	Q_EMIT statusChanged();
}

void LlmSegment::finishTool(const QString& resultText, bool success) {
	setResult(resultText);
	setStatus(success ? Status::Success : Status::Error);
	close();
}

void LlmSegment::restore(qint64 elapsedMs) {
	m_startedAt = m_timestamp;
	m_endedAt = m_timestamp + qMax<qint64>(0, elapsedMs);
}

void LlmSegment::scheduleMarkdown() {
	// Content segments only; reasoning/tool output is never parsed.
	// (User content is parsed too, so it can later be rendered as blocks
	// as well; the QML currently only does that for assistant messages.)
	if (m_type != Type::Content)
		return;
	m_markdownDirty = true;
	if (!m_markdownTimer.isActive())
		m_markdownTimer.start();
}

bool LlmSegment::parseMarkdown() {
	if (m_parseInFlight)
		return false;
	m_parseInFlight = true;
	const QString text = m_text;
	QThreadPool::globalInstance()->start([this, text]() {
		const QVariantList blocks = MarkdownParser::parse(text);
		// Deliver through qApp (never destroyed) and re-check the
		// pointer on the GUI thread: posting to `this` from the pool
		// thread would race with its destruction.
		QPointer<LlmSegment> guard(this);
		QMetaObject::invokeMethod(
			QCoreApplication::instance(),
			[guard, blocks]() {
				LlmSegment* seg = guard;
				if (!seg)
					return;
				seg->m_parseInFlight = false;
				seg->m_markdown = blocks;
				Q_EMIT seg->markdownChanged();
				// Text arrived while the worker was running; parse it.
				if (seg->m_markdownDirty)
					seg->parseMarkdown();
			},
			Qt::QueuedConnection);
	});
	return true;
}

} // namespace ZShell::llm