fix(chat): bound attachment encoding and recover rejected sends
This commit is contained in:
7 files changed
+657
-96
No files matched your search
@@ -0,0 +1,35 @@
|
||||
package com.m57.hermescontrol.ui.chat
|
||||
|
||||
import java.io.InputStream
|
||||
import java.io.OutputStream
|
||||
import java.util.Base64
|
||||
|
||||
/** Keeps the encoded JSON-RPC frame safely below the WebSocket's 16 MiB limit. */
|
||||
internal const val MAX_CHAT_ATTACHMENT_BYTES = 10L * 1024L * 1024L
|
||||
|
||||
internal enum class AttachmentSizeResult {
|
||||
WITHIN_LIMIT,
|
||||
TOO_LARGE,
|
||||
}
|
||||
|
||||
/** Streams Base64 to [output] while reading at most [maxBytes] of raw input. */
|
||||
internal fun encodeAttachmentBase64(
|
||||
input: InputStream,
|
||||
output: OutputStream,
|
||||
maxBytes: Long = MAX_CHAT_ATTACHMENT_BYTES,
|
||||
): AttachmentSizeResult {
|
||||
require(maxBytes >= 0) { "maxBytes must be non-negative" }
|
||||
|
||||
val buffer = ByteArray(DEFAULT_BUFFER_SIZE)
|
||||
var totalBytes = 0L
|
||||
|
||||
Base64.getEncoder().wrap(output).use { encodedOutput ->
|
||||
while (true) {
|
||||
val bytesRead = input.read(buffer)
|
||||
if (bytesRead == -1) return AttachmentSizeResult.WITHIN_LIMIT
|
||||
totalBytes += bytesRead
|
||||
if (totalBytes > maxBytes) return AttachmentSizeResult.TOO_LARGE
|
||||
encodedOutput.write(buffer, 0, bytesRead)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -46,6 +46,20 @@ object ChatInputPolicy {
|
||||
*/
|
||||
fun commandFieldValue(command: String): TextFieldValue = TextFieldValue(command, TextRange(command.length))
|
||||
|
||||
/** Restores a rejected prompt without discarding text typed since dispatch. */
|
||||
fun restoreRejectedText(
|
||||
rejectedText: String,
|
||||
current: TextFieldValue,
|
||||
): TextFieldValue {
|
||||
val restored =
|
||||
when {
|
||||
current.text.isEmpty() -> rejectedText
|
||||
rejectedText.isEmpty() -> current.text
|
||||
else -> "$rejectedText\n${current.text}"
|
||||
}
|
||||
return commandFieldValue(restored)
|
||||
}
|
||||
|
||||
/**
|
||||
* Rank slash-command suggestions by how often the user has dispatched them
|
||||
* (issue #865): most-used first, ties broken by the current (catalog)
|
||||
|
||||
@@ -298,6 +298,12 @@ fun ChatScreen(
|
||||
viewModel.consumePendingPrefill()
|
||||
}
|
||||
}
|
||||
LaunchedEffect(state.composerTextToRestore) {
|
||||
state.composerTextToRestore?.let { text ->
|
||||
inputFieldValue = ChatInputPolicy.restoreRejectedText(text, inputFieldValue)
|
||||
viewModel.consumeComposerTextRestore()
|
||||
}
|
||||
}
|
||||
var isListening by rememberSaveable { mutableStateOf(false) }
|
||||
var lastAnimatedMessageId by rememberSaveable { mutableStateOf<String?>(null) }
|
||||
var showReloginDialog by rememberSaveable { mutableStateOf(false) }
|
||||
|
||||
@@ -2,8 +2,6 @@ package com.m57.hermescontrol.ui.chat
|
||||
|
||||
import android.app.Application
|
||||
import android.net.Uri
|
||||
import android.util.Base64
|
||||
import android.util.Base64OutputStream
|
||||
import android.util.Log
|
||||
import androidx.lifecycle.AndroidViewModel
|
||||
import androidx.lifecycle.viewModelScope
|
||||
@@ -43,7 +41,6 @@ import kotlinx.coroutines.flow.stateIn
|
||||
import kotlinx.coroutines.flow.update
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.withContext
|
||||
import kotlinx.coroutines.yield
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.decodeFromJsonElement
|
||||
import kotlinx.serialization.json.intOrNull
|
||||
@@ -52,12 +49,27 @@ import okhttp3.MediaType.Companion.toMediaType
|
||||
import okhttp3.Request
|
||||
import okhttp3.RequestBody.Companion.toRequestBody
|
||||
import org.json.JSONObject
|
||||
import java.io.ByteArrayOutputStream
|
||||
import java.io.File
|
||||
import java.util.concurrent.ConcurrentHashMap
|
||||
|
||||
private const val TAG = "ChatViewModel"
|
||||
private const val MESSAGE_PAGE_SIZE = 150
|
||||
|
||||
private data class PreparedAttachment(
|
||||
val attachment: Attachment,
|
||||
val encodedFile: File,
|
||||
)
|
||||
|
||||
private sealed interface PrepareAttachmentResult {
|
||||
data class Success(
|
||||
val prepared: PreparedAttachment,
|
||||
) : PrepareAttachmentResult
|
||||
|
||||
data object TooLarge : PrepareAttachmentResult
|
||||
|
||||
data object Unreadable : PrepareAttachmentResult
|
||||
}
|
||||
|
||||
/**
|
||||
* Canonical comparison key for a tool message's result payload.
|
||||
*
|
||||
@@ -349,6 +361,8 @@ data class ChatUiState(
|
||||
val compressionCount: Int? = null,
|
||||
// Attachment state
|
||||
val pendingAttachments: List<Attachment> = emptyList(),
|
||||
/** One-shot composer recovery after an attachment is rejected before send. */
|
||||
val composerTextToRestore: String? = null,
|
||||
// Reaction animation — set when a reaction WS event arrives, auto-clears
|
||||
val reactionKind: String? = null,
|
||||
/** Monotonic trigger ID so consecutive same-kind reactions re-animate. */
|
||||
@@ -519,6 +533,7 @@ class ChatViewModel(
|
||||
val text: String,
|
||||
val attachments: List<Attachment>,
|
||||
val wasStreaming: Boolean,
|
||||
val userMessage: ChatMessage,
|
||||
)
|
||||
|
||||
private var pendingInitialPrompt: PendingPrompt? = null
|
||||
@@ -1132,6 +1147,7 @@ class ChatViewModel(
|
||||
wasStreaming = pending.wasStreaming,
|
||||
storageSessionId = storageId,
|
||||
agentSessionId = runtimeId,
|
||||
userMessage = pending.userMessage,
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -1436,6 +1452,20 @@ class ChatViewModel(
|
||||
return
|
||||
}
|
||||
|
||||
val oversizedAttachment =
|
||||
_uiState.value.pendingAttachments.firstOrNull {
|
||||
it.size > MAX_CHAT_ATTACHMENT_BYTES
|
||||
}
|
||||
if (oversizedAttachment != null) {
|
||||
_uiState.update {
|
||||
it.copy(
|
||||
errorMessage = attachmentTooLargeMessage(oversizedAttachment),
|
||||
composerTextToRestore = text,
|
||||
)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Snapshot + clear attachments so the input bar empties immediately
|
||||
val attachments = _uiState.value.pendingAttachments.toList()
|
||||
clearAttachments()
|
||||
@@ -1463,7 +1493,7 @@ class ChatViewModel(
|
||||
if (storageSessionId == null || agentSessionId == null) {
|
||||
// Issue #969: Session creation is still in-flight. Hold the prompt
|
||||
// so it is dispatched automatically the moment SESSION_CREATE lands.
|
||||
pendingInitialPrompt = PendingPrompt(text, attachments, wasStreaming)
|
||||
pendingInitialPrompt = PendingPrompt(text, attachments, wasStreaming, userMessage)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1485,6 +1515,7 @@ class ChatViewModel(
|
||||
agentSessionId: String,
|
||||
userMessage: ChatMessage? = null,
|
||||
) {
|
||||
val dispatchGeneration = sessionGeneration
|
||||
val msgToPersist =
|
||||
userMessage ?: ChatMessage(
|
||||
role = MessageRole.USER,
|
||||
@@ -1492,128 +1523,222 @@ class ChatViewModel(
|
||||
attachments = if (attachments.isNotEmpty()) attachments else null,
|
||||
)
|
||||
|
||||
// Persist under the original Desktop session ID.
|
||||
viewModelScope.launch(ioDispatcher) {
|
||||
repo.persistMessage(msgToPersist, storageSessionId)
|
||||
}
|
||||
|
||||
// Upload attachments then submit prompt
|
||||
viewModelScope.launch(ioDispatcher) {
|
||||
val fileRefs = mutableListOf<String>()
|
||||
val preparedAttachments = mutableListOf<PreparedAttachment>()
|
||||
|
||||
for (attachment in attachments) {
|
||||
val b64 = readContentUriBase64(attachment.uri)
|
||||
if (b64 == null) {
|
||||
Log.w(TAG, "Skipping unreadable attachment: ${attachment.name}")
|
||||
continue
|
||||
try {
|
||||
// Snapshot every attachment before the first RPC. This closes the
|
||||
// content-URI TOCTOU window without retaining multiple Base64
|
||||
// strings in the heap: encoded snapshots live in private cache.
|
||||
for (attachment in attachments) {
|
||||
when (val result = prepareAttachment(attachment)) {
|
||||
is PrepareAttachmentResult.Success -> {
|
||||
preparedAttachments += result.prepared
|
||||
}
|
||||
|
||||
PrepareAttachmentResult.TooLarge -> {
|
||||
rejectOversizedAttachment(
|
||||
attachment = attachment,
|
||||
attachments = attachments,
|
||||
message = msgToPersist,
|
||||
wasStreaming = wasStreaming,
|
||||
generation = dispatchGeneration,
|
||||
)
|
||||
return@launch
|
||||
}
|
||||
|
||||
PrepareAttachmentResult.Unreadable -> {
|
||||
Log.w(TAG, "Skipping unreadable attachment: ${attachment.name}")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Persist only after every readable attachment has a bounded,
|
||||
// immutable snapshot and no partial server upload can occur.
|
||||
try {
|
||||
if (attachment.isImage) {
|
||||
// Await so the backend stages the image into
|
||||
// session["attached_images"] BEFORE prompt.submit runs
|
||||
// (a fire-and-forget send raced prompt.submit and the
|
||||
// image was dropped). Requires session_id or the gateway
|
||||
// 4001s "session not found" (desktop passes it too).
|
||||
val result =
|
||||
repo.persistMessage(msgToPersist, storageSessionId)
|
||||
} catch (e: Exception) {
|
||||
if (e is kotlinx.coroutines.CancellationException) throw e
|
||||
Log.e(TAG, "Failed to persist outgoing message", e)
|
||||
if (dispatchGeneration == sessionGeneration) {
|
||||
_uiState.update { state ->
|
||||
state.copy(
|
||||
messages = state.messages.filterNot { it.id == msgToPersist.id },
|
||||
pendingAttachments = attachments + state.pendingAttachments,
|
||||
composerTextToRestore = msgToPersist.content,
|
||||
errorMessage = "Failed to save message",
|
||||
isAgentTyping = wasStreaming,
|
||||
)
|
||||
}
|
||||
}
|
||||
return@launch
|
||||
}
|
||||
|
||||
for ((attachment, encodedFile) in preparedAttachments) {
|
||||
try {
|
||||
// Materialize one bounded Base64 payload at a time. The
|
||||
// awaited RPC completes before the next snapshot is read.
|
||||
val b64 = encodedFile.readText(Charsets.US_ASCII)
|
||||
|
||||
if (attachment.isImage) {
|
||||
// Await so the backend stages the image into
|
||||
// session["attached_images"] BEFORE prompt.submit runs
|
||||
// (a fire-and-forget send raced prompt.submit and the
|
||||
// image was dropped). Requires session_id or the gateway
|
||||
// 4001s "session not found" (desktop passes it too).
|
||||
val result =
|
||||
sendRpcAndAwait(
|
||||
method = WsMethods.IMAGE_ATTACH_BYTES,
|
||||
params =
|
||||
mapOf(
|
||||
"session_id" to agentSessionId,
|
||||
"content_base64" to "data:${attachment.mimeType};base64,$b64",
|
||||
"filename" to attachment.name,
|
||||
"ext" to attachment.fileExtension,
|
||||
),
|
||||
)
|
||||
if (result != null) {
|
||||
@Suppress("UNCHECKED_CAST")
|
||||
val ok = (result as? Map<String, Any?>)?.get("attached") as? Boolean
|
||||
if (ok != true) {
|
||||
Log.w(TAG, "Image attach for ${attachment.name} returned non-ok: $result")
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Await the @file: ref text so we can embed it in the prompt.
|
||||
// file.attach also requires session_id or the gateway 4001s
|
||||
// "session not found" (same resolver as image.attach_bytes).
|
||||
sendRpcAndAwait(
|
||||
method = WsMethods.IMAGE_ATTACH_BYTES,
|
||||
method = WsMethods.FILE_ATTACH,
|
||||
params =
|
||||
mapOf(
|
||||
"session_id" to agentSessionId,
|
||||
"content_base64" to "data:${attachment.mimeType};base64,$b64",
|
||||
"filename" to attachment.name,
|
||||
"ext" to attachment.fileExtension,
|
||||
"data_url" to "data:${attachment.mimeType};base64,$b64",
|
||||
"name" to attachment.name,
|
||||
),
|
||||
)
|
||||
if (result != null) {
|
||||
@Suppress("UNCHECKED_CAST")
|
||||
val ok = (result as? Map<String, Any?>)?.get("attached") as? Boolean
|
||||
if (ok != true) {
|
||||
Log.w(TAG, "Image attach for ${attachment.name} returned non-ok: $result")
|
||||
)?.let { result ->
|
||||
@Suppress("UNCHECKED_CAST")
|
||||
val refText =
|
||||
(result as? Map<String, Any?>)?.get("ref_text") as? String
|
||||
if (!refText.isNullOrBlank()) fileRefs.add(refText)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Await the @file: ref text so we can embed it in the prompt.
|
||||
// file.attach also requires session_id or the gateway 4001s
|
||||
// "session not found" (same resolver as image.attach_bytes).
|
||||
sendRpcAndAwait(
|
||||
method = WsMethods.FILE_ATTACH,
|
||||
params =
|
||||
mapOf(
|
||||
"session_id" to agentSessionId,
|
||||
"data_url" to "data:${attachment.mimeType};base64,$b64",
|
||||
"name" to attachment.name,
|
||||
),
|
||||
)?.let { result ->
|
||||
@Suppress("UNCHECKED_CAST")
|
||||
val refText =
|
||||
(result as? Map<String, Any?>)?.get("ref_text") as? String
|
||||
if (!refText.isNullOrBlank()) fileRefs.add(refText)
|
||||
} catch (e: Exception) {
|
||||
if (e is kotlinx.coroutines.CancellationException) throw e
|
||||
Log.e(TAG, "Failed to upload attachment ${attachment.name}", e)
|
||||
if (dispatchGeneration == sessionGeneration) {
|
||||
_uiState.update {
|
||||
it.copy(errorMessage = "Upload failed: ${attachment.name}")
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (e: Exception) {
|
||||
Log.e(TAG, "Failed to upload attachment ${attachment.name}", e)
|
||||
_uiState.update {
|
||||
it.copy(errorMessage = "Upload failed: ${attachment.name}")
|
||||
}
|
||||
|
||||
// Build prompt text — prepend @file: refs for non-image files
|
||||
val fullText =
|
||||
if (fileRefs.isEmpty()) {
|
||||
text
|
||||
} else {
|
||||
fileRefs.joinToString("\n") +
|
||||
if (text.isNotBlank()) "\n\n$text" else ""
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build prompt text — prepend @file: refs for non-image files
|
||||
val fullText =
|
||||
if (fileRefs.isEmpty()) {
|
||||
text
|
||||
// While a turn is actively streaming and this is a plain text prompt
|
||||
// (no attachments — session.redirect carries text only), steer the
|
||||
// in-flight turn via session.redirect instead of queueing a fresh
|
||||
// prompt.submit. The backend rewrites the live turn when it can, or
|
||||
// queues the correction as the next turn otherwise (issue #710).
|
||||
if (dispatchGeneration == sessionGeneration) {
|
||||
ActiveSessionHolder.set(agentSessionId, storageSessionId)
|
||||
}
|
||||
if (wasStreaming && attachments.isEmpty()) {
|
||||
wsClient.sendRedirect(
|
||||
agentSessionId,
|
||||
fullText,
|
||||
onSent = { id ->
|
||||
trackSessionRequest(
|
||||
id = id,
|
||||
method = WsMethods.SESSION_REDIRECT,
|
||||
generation = dispatchGeneration,
|
||||
sessionId = storageSessionId,
|
||||
)
|
||||
},
|
||||
)
|
||||
} else {
|
||||
fileRefs.joinToString("\n") +
|
||||
if (text.isNotBlank()) "\n\n$text" else ""
|
||||
wsClient.sendMessage(
|
||||
agentSessionId,
|
||||
fullText,
|
||||
onSent = { id ->
|
||||
trackSessionRequest(
|
||||
id = id,
|
||||
method = WsMethods.PROMPT_SUBMIT,
|
||||
generation = dispatchGeneration,
|
||||
sessionId = storageSessionId,
|
||||
)
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// While a turn is actively streaming and this is a plain text prompt
|
||||
// (no attachments — session.redirect carries text only), steer the
|
||||
// in-flight turn via session.redirect instead of queueing a fresh
|
||||
// prompt.submit. The backend rewrites the live turn when it can, or
|
||||
// queues the correction as the next turn otherwise (issue #710).
|
||||
ActiveSessionHolder.set(agentSessionId, storageSessionId)
|
||||
if (wasStreaming && attachments.isEmpty()) {
|
||||
wsClient.sendRedirect(
|
||||
agentSessionId,
|
||||
fullText,
|
||||
onSent = { id -> trackRequest(id, WsMethods.SESSION_REDIRECT) },
|
||||
)
|
||||
} else {
|
||||
wsClient.sendMessage(
|
||||
agentSessionId,
|
||||
fullText,
|
||||
onSent = { id -> trackRequest(id, WsMethods.PROMPT_SUBMIT) },
|
||||
)
|
||||
} finally {
|
||||
preparedAttachments.forEach { it.encodedFile.delete() }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Read and encode a `content://` or `file://` URI to Base64 via ContentResolver, avoiding large allocations. */
|
||||
private suspend fun readContentUriBase64(uriString: String): String? =
|
||||
try {
|
||||
/** Snapshot one URI to private cache while enforcing the outbound frame limit. */
|
||||
private suspend fun prepareAttachment(attachment: Attachment): PrepareAttachmentResult {
|
||||
var encodedFile: File? = null
|
||||
return try {
|
||||
val context = getApplication<Application>()
|
||||
val uri = Uri.parse(uriString)
|
||||
context.contentResolver.openInputStream(uri)?.use { stream ->
|
||||
val baos = ByteArrayOutputStream()
|
||||
val b64os = Base64OutputStream(baos, Base64.NO_WRAP)
|
||||
val buffer = ByteArray(1024 * 128) // 128KB chunk
|
||||
var bytesRead: Int
|
||||
val uri = Uri.parse(attachment.uri)
|
||||
encodedFile = File.createTempFile("chat-attachment-", ".b64", context.cacheDir)
|
||||
val result =
|
||||
context.contentResolver.openInputStream(uri)?.use { input ->
|
||||
encodedFile.outputStream().buffered().use { output ->
|
||||
encodeAttachmentBase64(input, output)
|
||||
}
|
||||
} ?: return PrepareAttachmentResult.Unreadable.also { encodedFile.delete() }
|
||||
|
||||
while (stream.read(buffer).also { bytesRead = it } != -1) {
|
||||
b64os.write(buffer, 0, bytesRead)
|
||||
yield() // Prevent blocking the thread during large reads
|
||||
when (result) {
|
||||
AttachmentSizeResult.WITHIN_LIMIT -> {
|
||||
PrepareAttachmentResult.Success(PreparedAttachment(attachment, encodedFile))
|
||||
}
|
||||
|
||||
AttachmentSizeResult.TOO_LARGE -> {
|
||||
PrepareAttachmentResult.TooLarge.also { encodedFile.delete() }
|
||||
}
|
||||
b64os.close()
|
||||
baos.toString("UTF-8")
|
||||
}
|
||||
} catch (e: Exception) {
|
||||
encodedFile?.delete()
|
||||
if (e is kotlinx.coroutines.CancellationException) throw e
|
||||
Log.e(TAG, "Failed to read and encode attachment: ${e.message}", e)
|
||||
null
|
||||
Log.e(TAG, "Failed to prepare attachment: ${e.message}", e)
|
||||
PrepareAttachmentResult.Unreadable
|
||||
}
|
||||
}
|
||||
|
||||
private fun rejectOversizedAttachment(
|
||||
attachment: Attachment,
|
||||
attachments: List<Attachment>,
|
||||
message: ChatMessage,
|
||||
wasStreaming: Boolean,
|
||||
generation: Long,
|
||||
) {
|
||||
Log.w(TAG, "Rejecting oversized attachment: ${attachment.name}")
|
||||
if (generation != sessionGeneration) return
|
||||
_uiState.update { state ->
|
||||
state.copy(
|
||||
messages = state.messages.filterNot { it.id == message.id },
|
||||
pendingAttachments = attachments + state.pendingAttachments,
|
||||
composerTextToRestore = message.content,
|
||||
errorMessage = attachmentTooLargeMessage(attachment),
|
||||
isAgentTyping = wasStreaming,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
private fun attachmentTooLargeMessage(attachment: Attachment): String =
|
||||
"Attachment too large: ${attachment.name} (maximum 10 MB)"
|
||||
|
||||
/**
|
||||
* Send a JSON-RPC call and suspend until the response arrives, delegating
|
||||
@@ -1656,6 +1781,10 @@ class ChatViewModel(
|
||||
_uiState.update { it.copy(openError = null) }
|
||||
}
|
||||
|
||||
fun consumeComposerTextRestore() {
|
||||
_uiState.update { it.copy(composerTextToRestore = null) }
|
||||
}
|
||||
|
||||
fun gatewayPathFor(attachment: Attachment): String = mediaDelegate.gatewayPathFor(attachment)
|
||||
|
||||
fun clearAttachments() = attachmentsDelegate.clearAttachments()
|
||||
@@ -2794,6 +2923,7 @@ class ChatViewModel(
|
||||
contextBreakdown = null,
|
||||
compressionCount = null,
|
||||
pendingAttachments = emptyList(),
|
||||
composerTextToRestore = null,
|
||||
reactionKind = null,
|
||||
subagentIndicators = emptyList(),
|
||||
todos = emptyList(),
|
||||
|
||||
Reference in new issue
Block a user