chore: restructure llm plugin to be more coherent & fix various bugs in llm ui

Plugins/ZShell/Llm/generation.cpp

@@ -1,37 +1,50 @@
 #include "generation.hpp"
 
 #include "llmclient.hpp"
-#include "messagemodel.hpp"
+#include "message.hpp"
 #include "session.hpp"
 
 #include <QDateTime>
 
 namespace ZShell::llm {
 
-ChatGeneration::ChatGeneration(qint64 timestamp, QObject* parent)
-	: QObject(parent), m_timestamp(timestamp) {
+ChatGeneration::ChatGeneration(
+	ChatMessage* message, qint64 timestamp, QObject* parent)
+	: QObject(parent), m_message(message), m_timestamp(timestamp) {
 	m_timer.setParent(this);
 	m_timer.setInterval(500);
 	m_timer.setTimerType(Qt::CoarseTimer);
 	connect(&m_timer, &QTimer::timeout, this, [this]() {
-		bool anyRunning = false;
+		bool runningText = false;
+		bool runningTool = false;
 		for (auto* segment : m_segments) {
 			if (!segment->running()) continue;
-			anyRunning = true;
-			segment->elapsedMsChanged();
+			Q_EMIT segment->elapsedMsChanged();
+			if (segment->type() == LlmSegment::Type::ToolCall)
+				runningTool = true;
+			else
+				runningText = true;
 		}
-		if (anyRunning) Q_EMIT elapsedMsChanged();
-		if (!anyRunning && !m_streaming) m_timer.stop();
+		if (runningText) {
+			Q_EMIT reasoningElapsedMsChanged();
+			Q_EMIT contentElapsedMsChanged();
+		}
+		if (runningTool) Q_EMIT toolsElapsedMsChanged();
+		if (!runningText && !runningTool && !m_streaming) m_timer.stop();
 	});
 }
 
+ChatMessage* ChatGeneration::message() const {
+	return m_message;
+}
+
 QString ChatGeneration::content() const {
 	QStringList parts;
 	for (const auto* segment : m_segments) {
-		if (segment->type() != LlmSegment::Type::Content ||
-			segment->text().isEmpty())
-			continue;
-		parts.append(segment->text());
+		if (segment->type() != LlmSegment::Type::Content) continue;
+		const auto* text = qobject_cast<const LlmTextSegment*>(segment);
+		if (!text || text->text().isEmpty()) continue;
+		parts.append(text->text());
 	}
 	return parts.join(QStringLiteral("\n\n"));
 }
@@ -39,10 +52,10 @@ QString ChatGeneration::content() const {
 QString ChatGeneration::reasoning() const {
 	QStringList parts;
 	for (const auto* segment : m_segments) {
-		if (segment->type() != LlmSegment::Type::Reasoning ||
-			segment->text().isEmpty())
-			continue;
-		parts.append(segment->text());
+		if (segment->type() != LlmSegment::Type::Reasoning) continue;
+		const auto* text = qobject_cast<const LlmTextSegment*>(segment);
+		if (!text || text->text().isEmpty()) continue;
+		parts.append(text->text());
 	}
 	return parts.join(QStringLiteral("\n\n"));
 }
@@ -77,25 +90,27 @@ qint64 ChatGeneration::toolsElapsedMs() const {
 	return total;
 }
 
+QList<LlmToolCall*> ChatGeneration::toolCalls() const {
+	QList<LlmToolCall*> calls;
+	for (auto* segment : m_segments)
+		if (auto* call = qobject_cast<LlmToolCall*>(segment))
+			calls.append(call);
+	return calls;
+}
+
 int ChatGeneration::toolCallCount() const {
-	int count = 0;
-	for (const auto* segment : m_segments)
-		if (segment->type() == LlmSegment::Type::ToolCall) ++count;
-	return count;
+	return static_cast<int>(toolCalls().size());
 }
 
 bool ChatGeneration::hasRunningTool() const {
-	for (const auto* segment : m_segments)
-		if (segment->type() == LlmSegment::Type::ToolCall && segment->running())
-			return true;
+	for (const auto* call : toolCalls())
+		if (call->running()) return true;
 	return false;
 }
 
 bool ChatGeneration::hasPendingTool() const {
-	for (const auto* segment : m_segments)
-		if (segment->type() == LlmSegment::Type::ToolCall &&
-			segment->status() == LlmSegment::Status::Pending)
-			return true;
+	for (const auto* call : toolCalls())
+		if (call->status() == LlmToolCall::Status::Pending) return true;
 	return false;
 }
 
@@ -103,30 +118,37 @@ void ChatGeneration::setApprovalPending(bool value) {
 	if (m_approvalPending == value) return;
 	m_approvalPending = value;
 	Q_EMIT toolApprovalPendingChanged();
+	Q_EMIT pendingToolCallsChanged();
 }
 
 QVariantList ChatGeneration::pendingToolCalls() const {
 	QVariantList out;
 	if (!m_approvalPending) return out;
-	for (const auto* segment : m_segments)
-		if (segment->type() == LlmSegment::Type::ToolCall &&
-			segment->status() == LlmSegment::Status::Pending)
-			out.append(QVariant::fromValue(segment));
+	for (const auto* call : toolCalls())
+		if (call->status() == LlmToolCall::Status::Pending)
+			out.append(QVariant::fromValue(call));
 	return out;
 }
 
 void ChatGeneration::approveTools() {
-	if (auto* llmClient = client()) llmClient->approveTools();
+	if (auto* llmClient = client()) llmClient->approveTools(this);
 }
 
 void ChatGeneration::denyTools() {
-	if (auto* llmClient = client()) llmClient->denyTools();
+	if (auto* llmClient = client()) llmClient->denyTools(this);
+}
+
+void ChatGeneration::dismissError() {
+	if (!m_hasError && m_errorMessage.isEmpty()) return;
+	m_hasError = false;
+	m_errorMessage.clear();
+	Q_EMIT errorChanged();
+	Q_EMIT errorMessageChanged();
 }
 
 LlmClient* ChatGeneration::client() const {
-	if (auto* message = qobject_cast<ChatMessage*>(parent()))
-		if (auto* model = qobject_cast<ChatMessageModel*>(message->parent()))
-			if (auto* session = model->session()) return session->client();
+	if (m_message && m_message->session())
+		return m_message->session()->client();
 	return nullptr;
 }
 
@@ -138,21 +160,19 @@ void ChatGeneration::updateReasoningActive() {
 }
 
 void ChatGeneration::setContent(const QString& value) {
-	LlmSegment* first = nullptr;
+	LlmTextSegment* first = nullptr;
 	for (auto* segment : m_segments) {
 		if (segment->type() != LlmSegment::Type::Content) continue;
+		auto* text = qobject_cast<LlmTextSegment*>(segment);
+		if (!text) continue;
 		if (!first)
-			first = segment;
+			first = text;
 		else
-			segment->setText(QString());
+			text->setText(QString());
 	}
 	if (!first) {
 		if (value.isEmpty()) return;
-		first = new LlmSegment(
-			LlmSegment::Type::Content,
-			QDateTime::currentMSecsSinceEpoch(),
-			this);
-		addSegment(first);
+		first = openTextSegment(LlmSegment::Type::Content);
 	}
 	first->setText(value);
 }
@@ -164,7 +184,7 @@ void ChatGeneration::appendContent(const QString& piece) {
 			segment->running())
 			segment->close();
 	}
-	openContentSegment()->appendText(piece);
+	openTextSegment(LlmSegment::Type::Content)->appendText(piece);
 }
 
 void ChatGeneration::appendReasoning(const QString& piece) {
@@ -173,7 +193,7 @@ void ChatGeneration::appendReasoning(const QString& piece) {
 		if (segment->type() == LlmSegment::Type::Content && segment->running())
 			segment->close();
 	}
-	openReasoningSegment()->appendText(piece);
+	openTextSegment(LlmSegment::Type::Reasoning)->appendText(piece);
 }
 
 void ChatGeneration::setError(const QString& value) {
@@ -182,13 +202,7 @@ void ChatGeneration::setError(const QString& value) {
 	m_hasError = hasError;
 	m_errorMessage = value;
 	Q_EMIT errorChanged();
-}
-
-void ChatGeneration::dismissError() {
-	if (!m_hasError && m_errorMessage.isEmpty()) return;
-	m_hasError = false;
-	m_errorMessage.clear();
-	Q_EMIT errorChanged();
+	Q_EMIT errorMessageChanged();
 }
 
 void ChatGeneration::setStreaming(bool value) {
@@ -201,70 +215,67 @@ void ChatGeneration::setStreaming(bool value) {
 		closeOpenSegments();
 		m_timer.stop();
 	}
-	Q_EMIT elapsedMsChanged();
+	Q_EMIT reasoningElapsedMsChanged();
+	Q_EMIT contentElapsedMsChanged();
+	Q_EMIT toolsElapsedMsChanged();
 	updateReasoningActive();
 }
 
-LlmSegment* ChatGeneration::openContentSegment() {
-	for (auto* segment : m_segments)
-		if (segment->type() == LlmSegment::Type::Content && segment->running())
-			return segment;
-	auto* segment = new LlmSegment(
-		LlmSegment::Type::Content, QDateTime::currentMSecsSinceEpoch(), this);
-	segment->begin();
-	addSegment(segment);
-	return segment;
+LlmTextSegment* ChatGeneration::openTextSegment(LlmSegment::Type type) {
+	for (auto* segment : m_segments) {
+		if (segment->type() != type) continue;
+		auto* text = qobject_cast<LlmTextSegment*>(segment);
+		if (text && text->running()) return text;
+	}
+	auto* text =
+		new LlmTextSegment(type, QDateTime::currentMSecsSinceEpoch(), this);
+	text->begin();
+	addSegment(text);
+	return text;
 }
 
-LlmSegment* ChatGeneration::openReasoningSegment() {
-	for (auto* segment : m_segments)
-		if (segment->type() == LlmSegment::Type::Reasoning &&
-			segment->running())
-			return segment;
-	auto* segment = new LlmSegment(
-		LlmSegment::Type::Reasoning, QDateTime::currentMSecsSinceEpoch(), this);
-	segment->begin();
-	addSegment(segment);
-	return segment;
-}
-
-LlmSegment* ChatGeneration::beginToolCall(
+LlmToolCall* ChatGeneration::beginToolCall(
 	const QString& name, const QString& toolCallId) {
 	closeOpenSegments();
-	auto* segment = new LlmSegment(
-		LlmSegment::Type::ToolCall, QDateTime::currentMSecsSinceEpoch(), this);
-	segment->setName(name);
-	segment->setToolCallId(toolCallId);
-	segment->setStatus(LlmSegment::Status::Pending);
-	segment->begin();
-	addSegment(segment);
-	Q_EMIT toolStateChanged();
-	return segment;
+	auto* call = new LlmToolCall(QDateTime::currentMSecsSinceEpoch(), this);
+	call->setName(name);
+	call->setToolCallId(toolCallId);
+	call->setStatus(LlmToolCall::Status::Pending);
+	call->begin();
+	addSegment(call);
+	Q_EMIT toolCallCountChanged();
+	return call;
 }
 
 void ChatGeneration::addSegment(LlmSegment* segment) {
 	if (!segment || m_segments.contains(segment)) return;
 	segment->setParent(this);
-	connect(segment, &LlmSegment::textChanged, this, [this, segment]() {
-		if (segment->type() == LlmSegment::Type::Reasoning)
-			Q_EMIT reasoningChanged();
-		else if (segment->type() == LlmSegment::Type::Content)
-			Q_EMIT contentChanged();
-		updateReasoningActive();
-	});
-	connect(segment, &LlmSegment::statusChanged, this, [this]() {
-		Q_EMIT toolStateChanged();
-	});
-	connect(segment, &LlmSegment::resultChanged, this, [this]() {
-		Q_EMIT toolStateChanged();
-	});
-	connect(segment, &LlmSegment::runningChanged, this, [this]() {
-		Q_EMIT elapsedMsChanged();
-		Q_EMIT toolStateChanged();
+	connect(segment, &LlmSegment::runningChanged, this, [this, segment]() {
+		Q_EMIT toolCallCountChanged();
+		if (segment->type() == LlmSegment::Type::ToolCall)
+			Q_EMIT hasRunningToolChanged();
 		updateReasoningActive();
 	});
+	if (auto* text = qobject_cast<LlmTextSegment*>(segment))
+		connect(text, &LlmTextSegment::textChanged, this, [this, text]() {
+			if (text->type() == LlmSegment::Type::Reasoning)
+				Q_EMIT reasoningChanged();
+			else
+				Q_EMIT contentChanged();
+			updateReasoningActive();
+		});
+	if (auto* call = qobject_cast<LlmToolCall*>(segment)) {
+		connect(call, &LlmToolCall::statusChanged, this, [this, call]() {
+			Q_EMIT hasRunningToolChanged();
+			Q_EMIT pendingToolCallsChanged();
+			if (call->status() == LlmToolCall::Status::Pending)
+				Q_EMIT toolApprovalPendingChanged();
+			updateReasoningActive();
+		});
+	}
 	m_segments.append(segment);
 	Q_EMIT segmentsChanged();
+	Q_EMIT toolCallCountChanged();
 }
 
 void ChatGeneration::closeOpenSegments() {