diff --git a/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/GeminiBackend.kt b/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/GeminiBackend.kt index 2d06e877..3243b20d 100644 --- a/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/GeminiBackend.kt +++ b/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/GeminiBackend.kt @@ -298,35 +298,39 @@ class GeminiBackend( val fullText = StringBuilder() var chunkCount = 0 - val conn = openConnection(getModelName(), METHOD_STREAM_GENERATE_CONTENT, sse = true, apiKey = apiKey) - val cancelHandle = coroutineContext[Job]?.invokeOnCompletion { cause -> - if (cause != null) conn.disconnect() - } - try { - writeBody(conn, body) - checkResponse(conn) - // SSE: each chunk arrives as a `data: {json}` line; parse text as it streams. - conn.inputStream.bufferedReader().useLines { lines -> - for (line in lines) { - ensureActive() - if (!line.startsWith("data:")) continue - val payload = line.substringAfter("data:").trim() - if (payload.isEmpty() || payload == "[DONE]") continue - // A malformed/non-JSON chunk must not abort the whole stream; skip it. - val chunk = runCatching { extractText(JSONObject(payload)) }.getOrElse { - context.logger.warn("GeminiBackend: skipping malformed SSE chunk: ${it.message}") - "" - } - if (chunk.isNotEmpty()) { - chunkCount++ - fullText.append(chunk) - callback.onToken(chunk) + // Read outside the tagged block: `coroutineContext` is only reachable from suspend code. + val requestJob = coroutineContext[Job] + withTrafficTag(NetworkTags.INFERENCE) { + val conn = openConnection(getModelName(), METHOD_STREAM_GENERATE_CONTENT, sse = true, apiKey = apiKey) + val cancelHandle = requestJob?.invokeOnCompletion { cause -> + if (cause != null) conn.disconnect() + } + try { + writeBody(conn, body) + checkResponse(conn) + // SSE: each chunk arrives as a `data: {json}` line; parse text as it streams. + conn.inputStream.bufferedReader().useLines { lines -> + for (line in lines) { + ensureActive() + if (!line.startsWith("data:")) continue + val payload = line.substringAfter("data:").trim() + if (payload.isEmpty() || payload == "[DONE]") continue + // A malformed/non-JSON chunk must not abort the whole stream; skip it. + val chunk = runCatching { extractText(JSONObject(payload)) }.getOrElse { + context.logger.warn("GeminiBackend: skipping malformed SSE chunk: ${it.message}") + "" + } + if (chunk.isNotEmpty()) { + chunkCount++ + fullText.append(chunk) + callback.onToken(chunk) + } } } + } finally { + cancelHandle?.dispose() + conn.disconnect() } - } finally { - cancelHandle?.dispose() - conn.disconnect() } val finalText = fullText.toString() @@ -476,55 +480,58 @@ class GeminiBackend( /** * Fetch and parse the ListModels catalog, following pagination, keeping only models that - * support [METHOD_GENERATE_CONTENT]. Runs on the caller's (IO) coroutine. The + * support [METHOD_GENERATE_CONTENT]. Runs on the caller's (IO) coroutine, sockets tagged + * [NetworkTags.CATALOG] across every page. The * `ListModels HTTP ` message is a cross-plugin contract — keep that shape if you reword. */ private fun fetchAvailableModels(apiKey: String): List { val names = mutableListOf() - var pageToken: String? = null - - do { - val url = buildString { - append(MODELS_BASE_URL) - append("?pageSize=1000") - pageToken?.let { append("&pageToken=").append(java.net.URLEncoder.encode(it, "UTF-8")) } - } + withTrafficTag(NetworkTags.CATALOG) { + var pageToken: String? = null + + do { + val url = buildString { + append(MODELS_BASE_URL) + append("?pageSize=1000") + pageToken?.let { append("&pageToken=").append(java.net.URLEncoder.encode(it, "UTF-8")) } + } - val conn = (java.net.URL(url).openConnection() as java.net.HttpURLConnection).apply { - requestMethod = "GET" - connectTimeout = 15_000 - readTimeout = 15_000 - // Pass the API key as a header, never in the URL query string: query - // strings leak into logs, proxies, and crash reports. - setRequestProperty("x-goog-api-key", apiKey) - } + val conn = (java.net.URL(url).openConnection() as java.net.HttpURLConnection).apply { + requestMethod = "GET" + connectTimeout = 15_000 + readTimeout = 15_000 + // Pass the API key as a header, never in the URL query string: query + // strings leak into logs, proxies, and crash reports. + setRequestProperty("x-goog-api-key", apiKey) + } - val body = try { - val code = conn.responseCode - if (code !in 200..299) { - val err = conn.errorStream?.bufferedReader()?.use { it.readText() }.orEmpty() - throw java.io.IOException("ListModels HTTP $code: $err") + val body = try { + val code = conn.responseCode + if (code !in 200..299) { + val err = conn.errorStream?.bufferedReader()?.use { it.readText() }.orEmpty() + throw java.io.IOException("ListModels HTTP $code: $err") + } + conn.inputStream.bufferedReader().use { it.readText() } + } finally { + conn.disconnect() } - conn.inputStream.bufferedReader().use { it.readText() } - } finally { - conn.disconnect() - } - val json = org.json.JSONObject(body) - val models = json.optJSONArray("models") - if (models != null) { - for (i in 0 until models.length()) { - val model = models.getJSONObject(i) - val methods = model.optJSONArray("supportedGenerationMethods") ?: continue - val supportsChat = (0 until methods.length()) - .any { methods.optString(it) == METHOD_GENERATE_CONTENT } - if (!supportsChat) continue - val name = model.optString("name").removePrefix("models/") - if (name.isNotBlank()) names.add(name) + val json = org.json.JSONObject(body) + val models = json.optJSONArray("models") + if (models != null) { + for (i in 0 until models.length()) { + val model = models.getJSONObject(i) + val methods = model.optJSONArray("supportedGenerationMethods") ?: continue + val supportsChat = (0 until methods.length()) + .any { methods.optString(it) == METHOD_GENERATE_CONTENT } + if (!supportsChat) continue + val name = model.optString("name").removePrefix("models/") + if (name.isNotBlank()) names.add(name) + } } - } - pageToken = json.optString("nextPageToken").takeIf { it.isNotBlank() } - } while (pageToken != null) + pageToken = json.optString("nextPageToken").takeIf { it.isNotBlank() } + } while (pageToken != null) + } return names.distinct() } @@ -580,21 +587,24 @@ User: $userPrompt""" /** * POST [body] to a model method and return the concatenated response text. * + * The socket is tagged [NetworkTags.INFERENCE]. + * * @param model model name (without the `models/` prefix) * @param apiKey Gemini API key * @param body request payload built by [buildRequestJson] * @return the response text, or "" when the API returned no candidates/parts */ - private fun requestText(model: String, apiKey: String, body: JSONObject): String { - val conn = openConnection(model, METHOD_GENERATE_CONTENT, sse = false, apiKey = apiKey) - return try { - writeBody(conn, body) - checkResponse(conn) - extractText(JSONObject(conn.inputStream.bufferedReader().use { it.readText() })) - } finally { - conn.disconnect() + private fun requestText(model: String, apiKey: String, body: JSONObject): String = + withTrafficTag(NetworkTags.INFERENCE) { + val conn = openConnection(model, METHOD_GENERATE_CONTENT, sse = false, apiKey = apiKey) + try { + writeBody(conn, body) + checkResponse(conn) + extractText(JSONObject(conn.inputStream.bufferedReader().use { it.readText() })) + } finally { + conn.disconnect() + } } - } /** * Open a POST connection to `.../models/{model}:{method}`. diff --git a/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTags.kt b/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTags.kt new file mode 100644 index 00000000..1bdb7c1e --- /dev/null +++ b/ai-agent-gemini/src/main/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTags.kt @@ -0,0 +1,34 @@ +package com.itsaky.androidide.plugins.aiagentgemini.backend + +import android.net.TrafficStats + +/** + * Thread stats tags for this plugin's sockets: ASCII bytes, following the host's convention, so a + * raw `dumpsys netstats` dump stays readable. Keep-alive pooling makes the split approximate — a + * pooled socket keeps the tag it was born with, and inference and ListModels share a host. + */ +internal object NetworkTags { + + /** Generation, streaming and not — `"GEIN"`. */ + const val INFERENCE = 0x4745494E + + /** The ListModels catalog — `"GECT"`. */ + const val CATALOG = 0x47454354 +} + +/** + * Run [block] with [tag] on this thread's sockets, clearing the tag in a `finally` so it never + * leaks past the request onto these shared `Dispatchers.IO` threads. Not `inline` on purpose: the + * tag is thread-local, so [block] must not suspend. + * + * @param tag one of [NetworkTags] + * @return whatever [block] returned + */ +internal fun withTrafficTag(tag: Int, block: () -> T): T { + TrafficStats.setThreadStatsTag(tag) + return try { + block() + } finally { + TrafficStats.clearThreadStatsTag() + } +} diff --git a/ai-agent-gemini/src/test/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTagsTest.kt b/ai-agent-gemini/src/test/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTagsTest.kt new file mode 100644 index 00000000..ca086b63 --- /dev/null +++ b/ai-agent-gemini/src/test/kotlin/com/itsaky/androidide/plugins/aiagentgemini/backend/NetworkTagsTest.kt @@ -0,0 +1,76 @@ +package com.itsaky.androidide.plugins.aiagentgemini.backend + +import android.net.TrafficStats +import io.mockk.every +import io.mockk.mockkStatic +import io.mockk.unmockkStatic +import org.junit.After +import org.junit.Assert.* +import org.junit.Before +import org.junit.Test + +/** + * [TrafficStats] is stubbed rather than letting `isReturnDefaultValues` mock every Android API, + * which would silently change behaviour for the rest of the suite. + */ +class NetworkTagsTest { + + private var currentTag = UNTAGGED + + @Before + fun setup() { + mockkStatic(TrafficStats::class) + every { TrafficStats.getThreadStatsTag() } answers { currentTag } + every { TrafficStats.setThreadStatsTag(any()) } answers { currentTag = firstArg() } + every { TrafficStats.clearThreadStatsTag() } answers { currentTag = UNTAGGED } + } + + @After + fun tearDown() { + unmockkStatic(TrafficStats::class) + } + + @Test + fun givenATag_whenRunningABlock_thenTheTagIsSetForTheBlockAndItsValueReturned() { + var tagDuringBlock = UNTAGGED + + val result = withTrafficTag(NetworkTags.INFERENCE) { + tagDuringBlock = TrafficStats.getThreadStatsTag() + "done" + } + + assertEquals(NetworkTags.INFERENCE, tagDuringBlock) + assertEquals("done", result) + } + + @Test + fun givenABlockThatCompletes_whenItReturns_thenTheThreadStatsTagIsCleared() { + // A tag left behind would be charged to whatever this shared Dispatchers.IO thread + // does next, which is the very defect this helper exists to fix. + withTrafficTag(NetworkTags.CATALOG) {} + + assertEquals(UNTAGGED, currentTag) + } + + @Test + fun givenABlockThatThrows_whenRunningIt_thenTheThreadStatsTagIsStillCleared() { + assertThrows(IllegalStateException::class.java) { + withTrafficTag(NetworkTags.INFERENCE) { throw IllegalStateException("boom") } + } + + assertEquals(UNTAGGED, currentTag) + } + + @Test + fun givenTheDeclaredTags_whenInspected_thenTheyAreDistinctAndNonZero() { + // A zero tag means "untagged" to the kernel, so it would defeat the whole point. + assertNotEquals(0, NetworkTags.INFERENCE) + assertNotEquals(0, NetworkTags.CATALOG) + assertNotEquals(NetworkTags.INFERENCE, NetworkTags.CATALOG) + } + + private companion object { + /** What `clearThreadStatsTag()` leaves behind — the platform's "no tag" value. */ + const val UNTAGGED = -1 + } +} diff --git a/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/NetworkTags.kt b/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/NetworkTags.kt new file mode 100644 index 00000000..84086774 --- /dev/null +++ b/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/NetworkTags.kt @@ -0,0 +1,34 @@ +package com.itsaky.androidide.plugins.aiagentopenai.backend + +import android.net.TrafficStats + +/** + * Thread stats tags for this plugin's sockets: ASCII bytes, following the host's convention, so a + * raw `dumpsys netstats` dump stays readable. Keep-alive pooling makes the split approximate — a + * pooled socket keeps the tag it was born with, and inference and the catalog share a base URL. + */ +internal object NetworkTags { + + /** Chat completions, streaming and not — `"OAIN"`. */ + const val INFERENCE = 0x4F41494E + + /** The model catalog — `"OACT"`. */ + const val CATALOG = 0x4F414354 +} + +/** + * Run [block] with [tag] on this thread's sockets, clearing the tag in a `finally` so it never + * leaks past the request onto these shared `Dispatchers.IO` threads. Not `inline` on purpose: the + * tag is thread-local, so [block] must not suspend. + * + * @param tag one of [NetworkTags] + * @return whatever [block] returned + */ +internal fun withTrafficTag(tag: Int, block: () -> T): T { + TrafficStats.setThreadStatsTag(tag) + return try { + block() + } finally { + TrafficStats.clearThreadStatsTag() + } +} diff --git a/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/OpenAiHttpClient.kt b/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/OpenAiHttpClient.kt index 42f28229..d7e581d3 100644 --- a/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/OpenAiHttpClient.kt +++ b/ai-agent-openai/src/main/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/OpenAiHttpClient.kt @@ -13,6 +13,10 @@ import org.json.JSONObject * `okhttp3` resolves to the host's older OkHttp and an SDK bundling its own copy crashes with a * NoSuchMethodError. Kept apart from the backend so the backend is about generating, not sockets. * + * Both entry points tag their sockets ([NetworkTags]): the host installs + * `StrictMode.VmPolicy.detectAll()` process-wide, and an untagged socket trips + * `detectUntaggedSockets()` from inside plugin code. + * * @param connectTimeoutMs how long to wait for the connection itself * @param readTimeoutMs how long a generation may take to answer */ @@ -31,7 +35,8 @@ internal class OpenAiHttpClient( /** * POST [body] to [url] and hand the response's reader to [readResponse]. * - * The connection is closed before this returns, whatever [readResponse] did with it. + * The connection is closed before this returns, whatever [readResponse] did with it. Its + * socket is tagged [NetworkTags.INFERENCE]. * * @param apiKey bearer token, or blank for a server that needs none * @param sse true to ask for the server-sent-events stream @@ -46,7 +51,7 @@ internal class OpenAiHttpClient( sse: Boolean = false, onConnected: (HttpURLConnection) -> Unit = {}, readResponse: (BufferedReader) -> T, - ): T { + ): T = withTrafficTag(NetworkTags.INFERENCE) { val conn = open(url, "POST", apiKey).apply { readTimeout = readTimeoutMs doOutput = true @@ -54,7 +59,7 @@ internal class OpenAiHttpClient( if (sse) setRequestProperty("Accept", "text/event-stream") } onConnected(conn) - return try { + try { conn.outputStream.use { it.write(body.toString().toByteArray(Charsets.UTF_8)) } conn.failIfNotOk() conn.inputStream.bufferedReader().use(readResponse) @@ -64,14 +69,14 @@ internal class OpenAiHttpClient( } /** - * GET [url] and return its response body. + * GET [url] and return its response body, over a socket tagged [NetworkTags.CATALOG]. * * @param apiKey bearer token, or blank for a server that needs none * @throws OpenAiHttpException on a non-2xx answer, carrying the server's error body */ - fun get(url: String, apiKey: String): String { + fun get(url: String, apiKey: String): String = withTrafficTag(NetworkTags.CATALOG) { val conn = open(url, "GET", apiKey) - return try { + try { conn.failIfNotOk() conn.inputStream.bufferedReader().use { it.readText() } } finally { diff --git a/ai-agent-openai/src/test/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/NetworkTagsTest.kt b/ai-agent-openai/src/test/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/NetworkTagsTest.kt new file mode 100644 index 00000000..5665b363 --- /dev/null +++ b/ai-agent-openai/src/test/kotlin/com/itsaky/androidide/plugins/aiagentopenai/backend/NetworkTagsTest.kt @@ -0,0 +1,76 @@ +package com.itsaky.androidide.plugins.aiagentopenai.backend + +import android.net.TrafficStats +import io.mockk.every +import io.mockk.mockkStatic +import io.mockk.unmockkStatic +import org.junit.After +import org.junit.Assert.* +import org.junit.Before +import org.junit.Test + +/** + * [TrafficStats] is stubbed rather than letting `isReturnDefaultValues` mock every Android API, + * which would silently change behaviour for the rest of the suite. + */ +class NetworkTagsTest { + + private var currentTag = UNTAGGED + + @Before + fun setup() { + mockkStatic(TrafficStats::class) + every { TrafficStats.getThreadStatsTag() } answers { currentTag } + every { TrafficStats.setThreadStatsTag(any()) } answers { currentTag = firstArg() } + every { TrafficStats.clearThreadStatsTag() } answers { currentTag = UNTAGGED } + } + + @After + fun tearDown() { + unmockkStatic(TrafficStats::class) + } + + @Test + fun givenATag_whenRunningABlock_thenTheTagIsSetForTheBlockAndItsValueReturned() { + var tagDuringBlock = UNTAGGED + + val result = withTrafficTag(NetworkTags.INFERENCE) { + tagDuringBlock = TrafficStats.getThreadStatsTag() + "done" + } + + assertEquals(NetworkTags.INFERENCE, tagDuringBlock) + assertEquals("done", result) + } + + @Test + fun givenABlockThatCompletes_whenItReturns_thenTheThreadStatsTagIsCleared() { + // A tag left behind would be charged to whatever this shared Dispatchers.IO thread + // does next, which is the very defect this helper exists to fix. + withTrafficTag(NetworkTags.CATALOG) {} + + assertEquals(UNTAGGED, currentTag) + } + + @Test + fun givenABlockThatThrows_whenRunningIt_thenTheThreadStatsTagIsStillCleared() { + assertThrows(IllegalStateException::class.java) { + withTrafficTag(NetworkTags.INFERENCE) { throw IllegalStateException("boom") } + } + + assertEquals(UNTAGGED, currentTag) + } + + @Test + fun givenTheDeclaredTags_whenInspected_thenTheyAreDistinctAndNonZero() { + // A zero tag means "untagged" to the kernel, so it would defeat the whole point. + assertNotEquals(0, NetworkTags.INFERENCE) + assertNotEquals(0, NetworkTags.CATALOG) + assertNotEquals(NetworkTags.INFERENCE, NetworkTags.CATALOG) + } + + private companion object { + /** What `clearThreadStatsTag()` leaves behind — the platform's "no tag" value. */ + const val UNTAGGED = -1 + } +}