Add trackInference block overload and update docs/examples

README.md

@@ -223,11 +223,7 @@ val handle = wildEdge.registerModel("my-model", ModelInfo(
 ))
 
 handle.trackLoad(durationMs = loadMs, accelerator = Accelerator.CPU, coldStart = true)
-
-val start = System.currentTimeMillis()
-val output = model.run(input)
-handle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt())
-
+val output = handle.trackInference { model.run(input) }
 handle.trackUnload()
 ```
 
@@ -260,19 +256,8 @@ Group related inferences so the server can reconstruct the full pipeline:
 
 ```kotlin
 wildEdge.trace("user-query") { trace ->
-    trace.span("embed") {
-        val start = System.currentTimeMillis()
-        val embedding = embedModel.run(input)
-        embedHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt())
-        embedding
-    }
-
-    trace.span("classify") {
-        val start = System.currentTimeMillis()
-        val label = classifyModel.run(embedding)
-        classifyHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt())
-        label
-    }
+    val embedding = trace.span("embed") { embedHandle.trackInference { embedModel.run(input) } }
+    trace.span("classify") { classifyHandle.trackInference { classifyModel.run(embedding) } }
 }
 ```
 
@@ -318,12 +303,7 @@ Declare your field against `WildEdgeClient` and inject `WildEdgeClient.noop()` i
 class InferenceService(private val wildEdge: WildEdgeClient) {
     private val handle = wildEdge.registerModel("my-model", ...)
 
-    fun run(input: ByteArray): Result {
-        val start = System.currentTimeMillis()
-        val result = model.run(input)
-        handle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt())
-        return result
-    }
+    fun run(input: ByteArray): Result = handle.trackInference { model.run(input) }
 }
 
 // In tests:

examples/tracing/TracingExample.kt

@@ -2,6 +2,7 @@ package examples.tracing
 
 import dev.wildedge.sdk.ModelInfo
 import dev.wildedge.sdk.WildEdge
+import dev.wildedge.sdk.trackInference
 
 // Assumes WildEdge.init() has already run (manifest meta-data or Application.onCreate()).
 class TracingExample {
@@ -15,19 +16,8 @@ class TracingExample {
     // Nested span() calls set parent_span_id so the server can reconstruct the tree.
     fun runPipeline(input: ByteArray) {
         wildEdge.trace("pipeline") { trace ->
-            val embedding = trace.span("embed") {
-                val start = System.currentTimeMillis()
-                val result = runEmbedding(input)
-                embedHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt())
-                result
-            }
-
-            trace.span("classify") {
-                val start = System.currentTimeMillis()
-                val result = runClassification(embedding)
-                classifyHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt())
-                result
-            }
+            val embedding = trace.span("embed") { embedHandle.trackInference { runEmbedding(input) } }
+            trace.span("classify") { classifyHandle.trackInference { runClassification(embedding) } }
         }
     }
 

wildedge/src/main/kotlin/dev/wildedge/sdk/ModelHandle+Coroutines.kt

@@ -6,6 +6,53 @@ import dev.wildedge.sdk.events.TextInputMeta
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flow
 
+/**
+ * Times a blocking inference block and emits a tracking event on completion or error.
+ *
+ * Usage:
+ * ```
+ * val result = handle.trackInference { model.run(input) }
+ * ```
+ */
+@Suppress("TooGenericExceptionCaught")
+fun <T> ModelHandle.trackInference(
+    inputModality: InputModality? = null,
+    outputModality: OutputModality? = null,
+    inputMeta: Map<String, Any?>? = null,
+    outputMeta: Map<String, Any?>? = null,
+    outputMetaExtractor: ((T) -> Map<String, Any?>?)? = null,
+    traceId: String? = null,
+    parentSpanId: String? = null,
+    block: () -> T,
+): T {
+    val start = System.currentTimeMillis()
+    return try {
+        val result = block()
+        trackInference(
+            durationMs = (System.currentTimeMillis() - start).toInt(),
+            inputModality = inputModality,
+            outputModality = outputModality,
+            inputMeta = inputMeta,
+            outputMeta = outputMetaExtractor?.invoke(result) ?: outputMeta,
+            traceId = traceId,
+            parentSpanId = parentSpanId,
+        )
+        result
+    } catch (e: Exception) {
+        trackInference(
+            durationMs = (System.currentTimeMillis() - start).toInt(),
+            inputModality = inputModality,
+            outputModality = outputModality,
+            inputMeta = inputMeta,
+            success = false,
+            errorCode = e.javaClass.simpleName,
+            traceId = traceId,
+            parentSpanId = parentSpanId,
+        )
+        throw e
+    }
+}
+
 /**
  * Times a suspend inference block and emits a tracking event on completion or error.
  *

wildedge/src/test/kotlin/dev/wildedge/sdk/CoroutinesTrackingTest.kt

@@ -18,6 +18,58 @@ class CoroutinesTrackingTest {
         return handle to events
     }
 
+    // --- trackInference (block) ---
+
+    @Test fun blockTrackingEmitsInferenceEvent() {
+        val (handle, events) = captureHandle()
+        handle.trackInference { "result" }
+        assertEquals(1, events.count { it["event_type"] == "inference" })
+    }
+
+    @Test fun blockTrackingReturnsBlockResult() {
+        val (handle, _) = captureHandle()
+        val result = handle.trackInference { 42 }
+        assertEquals(42, result)
+    }
+
+    @Test fun blockTrackingRecordsSuccessTrue() {
+        val (handle, events) = captureHandle()
+        handle.trackInference { Unit }
+        val inference = events.first { it["event_type"] == "inference" }
+        @Suppress("UNCHECKED_CAST")
+        assertEquals(true, (inference["inference"] as Map<String, Any?>)["success"])
+    }
+
+    @Test fun blockTrackingRecordsFailureOnException() {
+        val (handle, events) = captureHandle()
+        runCatching { handle.trackInference<Unit> { throw IllegalStateException("boom") } }
+        val inference = events.first { it["event_type"] == "inference" }
+
+        @Suppress("UNCHECKED_CAST")
+        val inf = inference["inference"] as Map<String, Any?>
+        assertEquals(false, inf["success"])
+        assertEquals("IllegalStateException", inf["error_code"])
+    }
+
+    @Test fun blockTrackingRethrowsException() {
+        val (handle, _) = captureHandle()
+        val ex = runCatching {
+            handle.trackInference<Unit> { throw RuntimeException("rethrow me") }
+        }.exceptionOrNull()
+        assertNotNull(ex)
+        assertEquals("rethrow me", ex!!.message)
+    }
+
+    @Test fun blockTrackingOutputMetaExtractorReceivesResult() {
+        val (handle, events) = captureHandle()
+        handle.trackInference(outputMetaExtractor = { r: String -> mapOf("extracted" to r) }) { "hello" }
+        val inference = events.first { it["event_type"] == "inference" }
+
+        @Suppress("UNCHECKED_CAST")
+        val outputMeta = (inference["inference"] as Map<String, Any?>)["output_meta"] as Map<String, Any?>?
+        assertEquals("hello", outputMeta?.get("extracted"))
+    }
+
     // --- trackSuspendInference ---
 
     @Test fun suspendTrackingEmitsInferenceEvent() = runBlocking {