ADFA-5254 | Tag AI network sockets with TrafficStats

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 <code>` message is a cross-plugin contract — keep that shape if you reword.
      */
     private fun fetchAvailableModels(apiKey: String): List<String> {
         val names = mutableListOf<String>()
-        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}`.

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 <T> withTrafficTag(tag: Int, block: () -> T): T {
+    TrafficStats.setThreadStatsTag(tag)
+    return try {
+        block()
+    } finally {
+        TrafficStats.clearThreadStatsTag()
+    }
+}

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
+    }
+}

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 <T> withTrafficTag(tag: Int, block: () -> T): T {
+    TrafficStats.setThreadStatsTag(tag)
+    return try {
+        block()
+    } finally {
+        TrafficStats.clearThreadStatsTag()
+    }
+}