fix: truncate webfetch content to fit context

Plugins/ZShell/Llm/llmclient.cpp

@@ -17,6 +17,81 @@
 
 namespace ZShell::llm {
 
+namespace {
+
+// Minimum prompt budget kept even when the transcript alone nearly fills
+// the window, so the model always sees the triggering message.
+constexpr int kMinContextTokens = 512;
+
+qsizetype jsonValueChars(const QJsonValue& value) {
+	switch (value.type()) {
+	case QJsonValue::String:
+		return value.toString().size();
+	case QJsonValue::Array: {
+		qsizetype total = 0;
+		for (const QJsonValue& item : value.toArray())
+			total += jsonValueChars(item);
+		return total;
+	}
+	case QJsonValue::Object: {
+		qsizetype total = 0;
+		const QJsonObject obj = value.toObject();
+		for (auto it = obj.constBegin(); it != obj.constEnd(); ++it)
+			total += jsonValueChars(it.value());
+		return total;
+	}
+	default:
+		return 2;
+	}
+}
+
+int estimateMessageTokens(const QJsonObject& message) {
+	return int(jsonValueChars(message) / LlmTool::CharsPerToken) + 4;
+}
+
+// One message model row and the JSON messages it produces. Trimming works
+// in whole units so an assistant turn and its tool results always travel
+// together.
+struct ContextUnit {
+	QList<QJsonObject> messages;
+	int tokens = 0;
+};
+
+// Shrink the free-text fields of a unit until it fits the budget. Used
+// when a single unit (e.g. one huge user message) alone exceeds it.
+void shrinkUnitToBudget(QList<QJsonObject>& unit, int budgetTokens) {
+	static const QStringList kTextKeys = {
+		QStringLiteral("content"), QStringLiteral("reasoning_content")};
+	for (int pass = 0; pass < 8; ++pass) {
+		int tokens = 0;
+		for (const QJsonObject& message : unit)
+			tokens += estimateMessageTokens(message);
+		if (tokens <= budgetTokens) return;
+		qsizetype bestMessage = -1;
+		QString bestKey;
+		qsizetype bestLength = 0;
+		for (qsizetype i = 0; i < unit.size(); ++i) {
+			const QJsonObject message = unit.at(i);
+			for (const QString& key : kTextKeys) {
+				const QJsonValue value = message.value(key);
+				if (value.isString() && value.toString().size() > bestLength) {
+					bestMessage = i;
+					bestKey = key;
+					bestLength = value.toString().size();
+				}
+			}
+		}
+		if (bestMessage < 0) return;
+		QString text = unit.at(bestMessage).value(bestKey).toString();
+		text.truncate(qMax<qsizetype>(32, text.size() * 3 / 4));
+		const QString marker = QStringLiteral("\n[…]");
+		if (!text.endsWith(marker)) text += marker;
+		unit[bestMessage][bestKey] = text;
+	}
+}
+
+} // namespace
+
 QString LlmClient::completionsPath(
 	const QString& endpoint, const QString& subpath) {
 	QString base = endpoint.trimmed();
@@ -97,7 +172,7 @@ void LlmClient::startGeneration(ChatSession* session, ChatGeneration* target) {
 	if (m_busy || !session || !target) return;
 	m_active = session;
 	m_streaming = target;
-	m_streaming->setStreaming(true);
+	target->setStreaming(true);
 	setBusy(true);
 	setStreamingChatId(session->id());
 
@@ -125,6 +200,7 @@ void LlmClient::startGeneration(ChatSession* session, ChatGeneration* target) {
 	m_round = 0;
 	m_contentMark = 0;
 	m_reasoningMark = 0;
+	setApprovalPending(false);
 
 	sendRound();
 }
@@ -146,9 +222,11 @@ void LlmClient::sendRound() {
 		return;
 	}
 
-	const auto* model = m_active->messagesModel();
+	ChatSession* active = m_active;
+	ChatGeneration* streaming = m_streaming;
+	const auto* model = active->messagesModel();
 	const int targetRow =
-		model->rowOf(qobject_cast<ChatMessage*>(m_streaming->parent()));
+		model->rowOf(qobject_cast<ChatMessage*>(streaming->parent()));
 	if (targetRow < 0) {
 		fail(QStringLiteral(
 			"Internal error: generation target is not in the "
@@ -156,7 +234,13 @@ void LlmClient::sendRound() {
 		return;
 	}
 
-	QJsonArray messages = buildContextMessages(m_active, targetRow);
+	// The transcript (this generation so far) is always sent in full, so
+	// reserve its estimated size out of the context budget.
+	int transcriptTokens = 0;
+	for (const QJsonValue& value : m_transcript)
+		transcriptTokens += estimateMessageTokens(value.toObject());
+	QJsonArray messages =
+		buildContextMessages(m_active, targetRow, transcriptTokens);
 	for (const QJsonValue& value : m_transcript)
 		messages.append(value);
 
@@ -205,67 +289,103 @@ void LlmClient::sendRound() {
 }
 
 QJsonArray LlmClient::buildContextMessages(
-	ChatSession* session, int stopBeforeRow) const {
+	ChatSession* session, int stopBeforeRow, int extraReserveTokens) const {
 	const auto* model = session->messagesModel();
-	QJsonArray messages;
+	QList<ContextUnit> units;
 	for (int row = model->rowCount() - 1; row > stopBeforeRow; --row) {
 		const auto* message = model->at(row);
 		const auto* generation = message->activeGeneration();
 		if (!generation) continue;
 
+		ContextUnit unit;
 		if (message->role() == ChatMessage::Role::User) {
 			QJsonObject user;
 			user[QStringLiteral("role")] = QStringLiteral("user");
 			user[QStringLiteral("content")] = generation->content();
-			messages.append(user);
-			continue;
-		}
+			unit.messages.append(user);
+		} else {
+			QList<const LlmSegment*> toolSegments;
+			for (const auto* segment : generation->segments())
+				if (segment->type() == LlmSegment::Type::ToolCall)
+					toolSegments.append(segment);
 
-		QList<const LlmSegment*> toolSegments;
-		for (const auto* segment : generation->segments())
-			if (segment->type() == LlmSegment::Type::ToolCall)
-				toolSegments.append(segment);
-
-		QJsonObject assistant;
-		assistant[QStringLiteral("role")] = QStringLiteral("assistant");
-		if (toolSegments.isEmpty())
-			assistant[QStringLiteral("content")] = generation->content();
-		else if (!generation->content().isEmpty())
-			assistant[QStringLiteral("content")] = generation->content();
-		if (!generation->reasoning().isEmpty())
-			assistant[QStringLiteral("reasoning_content")] =
-				generation->reasoning();
-
-		if (toolSegments.isEmpty()) {
-			messages.append(assistant);
-			continue;
-		}
-		QJsonArray calls;
-		for (const auto* segment : toolSegments) {
-			QJsonObject function;
-			function[QStringLiteral("name")] = segment->name();
-			function[QStringLiteral("arguments")] = segment->arguments();
-			QJsonObject call;
-			call[QStringLiteral("id")] = segment->toolCallId();
-			call[QStringLiteral("type")] = QStringLiteral("function");
-			call[QStringLiteral("function")] = function;
-			calls.append(call);
-		}
-		assistant[QStringLiteral("tool_calls")] = calls;
-		messages.append(assistant);
-		for (const auto* segment : toolSegments) {
-			QJsonObject toolMessage;
-			toolMessage[QStringLiteral("role")] = QStringLiteral("tool");
-			toolMessage[QStringLiteral("tool_call_id")] = segment->toolCallId();
-			toolMessage[QStringLiteral("content")] = segment->result();
-			messages.append(toolMessage);
+			QJsonObject assistant;
+			assistant[QStringLiteral("role")] = QStringLiteral("assistant");
+			if (toolSegments.isEmpty())
+				assistant[QStringLiteral("content")] = generation->content();
+			else if (!generation->content().isEmpty())
+				assistant[QStringLiteral("content")] = generation->content();
+			if (!generation->reasoning().isEmpty())
+				assistant[QStringLiteral("reasoning_content")] =
+					generation->reasoning();
+			if (!toolSegments.isEmpty()) {
+				QJsonArray calls;
+				for (const auto* segment : toolSegments) {
+					QJsonObject function;
+					function[QStringLiteral("name")] = segment->name();
+					function[QStringLiteral("arguments")] =
+						segment->arguments();
+					QJsonObject call;
+					call[QStringLiteral("id")] = segment->toolCallId();
+					call[QStringLiteral("type")] = QStringLiteral("function");
+					call[QStringLiteral("function")] = function;
+					calls.append(call);
+				}
+				assistant[QStringLiteral("tool_calls")] = calls;
+			}
+			unit.messages.append(assistant);
+			for (const auto* segment : toolSegments) {
+				QJsonObject toolMessage;
+				toolMessage[QStringLiteral("role")] = QStringLiteral("tool");
+				toolMessage[QStringLiteral("tool_call_id")] =
+					segment->toolCallId();
+				toolMessage[QStringLiteral("content")] = segment->result();
+				unit.messages.append(toolMessage);
+			}
 		}
+		for (const QJsonObject& messageJson : unit.messages)
+			unit.tokens += estimateMessageTokens(messageJson);
+		units.append(unit);
 	}
+
+	if (m_contextSize > 0) {
+		// Keep headroom for the model's own reply so the prompt cannot
+		// fill the whole window.
+		const int completionReserve = qBound(1024, m_contextSize / 8, 8192);
+		int budget = m_contextSize - completionReserve - extraReserveTokens;
+		budget = qMax(budget, kMinContextTokens);
+
+		int total = 0;
+		for (const ContextUnit& unit : units)
+			total += unit.tokens;
+		// Drop oldest units until the prompt fits the budget.
+		while (units.size() > 1 && total > budget) {
+			total -= units.first().tokens;
+			units.removeFirst();
+		}
+		// A conversation cannot start with an assistant turn.
+		while (
+			units.size() > 1 &&
+			units.first().messages.first()[QStringLiteral("role")].toString() ==
+				QLatin1String("assistant")) {
+			total -= units.first().tokens;
+			units.removeFirst();
+		}
+		// A single oversized unit still gets sent, shrunk to fit.
+		if (units.size() == 1 && units.first().tokens > budget)
+			shrinkUnitToBudget(units.first().messages, budget);
+	}
+
+	QJsonArray messages;
+	for (const ContextUnit& unit : units)
+		for (const QJsonObject& messageJson : unit.messages)
+			messages.append(messageJson);
 	return messages;
 }
 
 void LlmClient::applyToolCallDelta(const QJsonObject& call) {
-	if (!m_streaming) return;
+	ChatGeneration* streaming = m_streaming;
+	if (!streaming) return;
 	const int index = call[QStringLiteral("index")].toInt(-1);
 	if (index < 0) return;
 	while (m_callBuilders.size() <= index)
@@ -281,8 +401,8 @@ void LlmClient::applyToolCallDelta(const QJsonObject& call) {
 	if (!arguments.isEmpty()) builder.arguments += arguments;
 
 	if (!builder.segment) {
-		m_streaming->closeOpenSegments();
-		builder.segment = m_streaming->beginToolCall(builder.name, builder.id);
+		streaming->closeOpenSegments();
+		builder.segment = streaming->beginToolCall(builder.name, builder.id);
 	}
 	builder.segment->setName(builder.name);
 	builder.segment->setToolCallId(builder.id);
@@ -291,6 +411,7 @@ void LlmClient::applyToolCallDelta(const QJsonObject& call) {
 
 void LlmClient::roundFinished() {
 	if (m_roundDone || !m_streaming) return;
+	ChatGeneration* streaming = m_streaming;
 	m_roundDone = true;
 
 	bool hasCalls = false;
@@ -306,10 +427,10 @@ void LlmClient::roundFinished() {
 		return;
 	}
 
-	m_streaming->closeOpenSegments();
+	streaming->closeOpenSegments();
 
-	const QString content = m_streaming->content();
-	const QString reasoning = m_streaming->reasoning();
+	const QString content = streaming->content();
+	const QString reasoning = streaming->reasoning();
 	QJsonObject assistant;
 	assistant[QStringLiteral("role")] = QStringLiteral("assistant");
 	if (content.size() > m_contentMark) {
@@ -337,7 +458,9 @@ void LlmClient::roundFinished() {
 	m_transcript.append(assistant);
 	m_round++;
 
-	executeAllCalls();
+	// The whole batch of this round waits for a single user decision;
+	// approveTools() runs all of it, denyTools() ends the turn.
+	setApprovalPending(true);
 }
 
 void LlmClient::executeAllCalls() {
@@ -375,7 +498,7 @@ void LlmClient::executeAllCalls() {
 		}
 
 		++m_pendingCalls;
-		tool->execute(args, [this, i, call](const QJsonObject& result) {
+		tool->execute(call.id, args, [this, i, call](const QJsonObject& result) {
 			if (!m_streaming) return;
 			const bool success = result.contains(QStringLiteral("output"));
 			const QString content =
@@ -407,12 +530,60 @@ void LlmClient::flushCallResults() {
 	sendRound();
 }
 
+void LlmClient::setApprovalPending(bool value) {
+	if (m_approvalPending == value) return;
+	m_approvalPending = value;
+	// Mirror onto the streaming generation; QML reads and answers from
+	// there, so only the affected ProcessBlock reacts.
+	if (ChatGeneration* streaming = m_streaming)
+		streaming->setApprovalPending(value);
+}
+
+void LlmClient::approveTools() {
+	if (!m_approvalPending) return;
+	ChatGeneration* streaming = m_streaming;
+	if (!streaming) return;
+	setApprovalPending(false);
+	for (auto* segment : streaming->segments()) {
+		if (segment->type() != LlmSegment::Type::ToolCall ||
+			segment->status() != LlmSegment::Status::Pending)
+			continue;
+		// Restart the segment timer so elapsed time covers execution, not
+		// the wait for approval.
+		segment->setStatus(LlmSegment::Status::Running);
+		segment->begin();
+	}
+	executeAllCalls();
+}
+
+void LlmClient::denyTools() {
+	if (!m_approvalPending || !m_streaming) return;
+	setApprovalPending(false);
+	finishPendingCalls(QStringLiteral("Denied by user"));
+	finishTurn();
+}
+
+void LlmClient::finishPendingCalls(const QString& resultText) {
+	ChatGeneration* streaming = m_streaming;
+	if (!streaming) return;
+	for (auto* segment : streaming->segments())
+		if (segment->type() == LlmSegment::Type::ToolCall &&
+			segment->status() == LlmSegment::Status::Pending)
+			segment->finishTool(resultText, false);
+}
+
 void LlmClient::stop() {
 	if (!m_busy) return;
+	if (m_approvalPending) {
+		setApprovalPending(false);
+		finishPendingCalls(QStringLiteral("Cancelled"));
+		finishTurn();
+		return;
+	}
 	if (m_toolPhase) {
 		m_tools->cancelAll();
-		if (m_streaming) {
-			for (auto* segment : m_streaming->segments()) {
+		if (ChatGeneration* streaming = m_streaming) {
+			for (auto* segment : streaming->segments()) {
 				if (segment->type() == LlmSegment::Type::ToolCall &&
 					segment->running())
 					segment->finishTool(QStringLiteral("Cancelled"), false);
@@ -426,11 +597,19 @@ void LlmClient::stop() {
 
 void LlmClient::endStream() {
 	if (!m_streaming) return;
-	auto* generation = m_streaming;
-	auto* session = m_active;
+	ChatGeneration* generation = m_streaming;
+	ChatSession* session = m_active;
 	m_streaming = nullptr;
 	m_active = nullptr;
 	generation->setStreaming(false);
+	m_approvalPending = false;
+	generation->setApprovalPending(false);
+	// Finalize tool calls that never received an approval (turn cancelled,
+	// round limit reached, ...).
+	for (auto* segment : generation->segments())
+		if (segment->type() == LlmSegment::Type::ToolCall &&
+			segment->status() == LlmSegment::Status::Pending)
+			segment->finishTool(QStringLiteral("Cancelled"), false);
 	if (generation->content().isEmpty() && generation->reasoning().isEmpty() &&
 		generation->toolCallCount() == 0) {
 		if (auto* message = qobject_cast<ChatMessage*>(generation->parent())) {
@@ -570,6 +749,7 @@ void LlmClient::refreshModels() {
 void LlmClient::setContextSize(int size) {
 	if (size <= 0 || m_contextSize == size) return;
 	m_contextSize = size;
+	m_tools->setContextSize(size);
 	Q_EMIT contextSizeChanged();
 }
 
@@ -611,12 +791,13 @@ void LlmClient::probeContextSize() {
 }
 
 void LlmClient::updateTokenUsage(const QJsonObject& data) {
-	if (!m_active || m_contextSize <= 0) return;
+	ChatSession* active = m_active;
+	if (!active || m_contextSize <= 0) return;
 	const QJsonObject usage = data["usage"].toObject();
 	if (usage.isEmpty()) return;
 	const double used = usage.value("prompt_tokens").toDouble() +
 						usage.value("completion_tokens").toDouble();
-	if (used > 0) m_active->setLastTokenCount(static_cast<int>(used));
+	if (used > 0) active->setLastTokenCount(static_cast<int>(used));
 }
 
 void LlmClient::shortRequest(