feat: Adds support for tracing streams with traceable (#117)

* Adds support for tracing streams with traceable

* Polish

* Make stream tracing opt-in

* Rework to use a passthrough instead of a proxy

* Record stream cancellations as errors

* Feedback

* Format and add to AGENTS.md

Author: jacoblee93

AGENTS.md

@@ -91,6 +91,25 @@ override fun toString(): String =
         "}"
 ```
 
+### Prefer extension functions over type casts
+
+When adding behavior to a type you don't own, use a `private` extension function instead of casting to an implementation type:
+
+```kotlin
+// Good — extension function, no cast needed
+private fun Stream<*>.withErrorTracking(
+    errorRef: AtomicReference<Throwable>,
+    exhaustedRef: AtomicBoolean,
+): Stream<Any?> { ... }
+
+val instrumented = result.withErrorTracking(iterationError, streamExhausted)
+
+// Bad — casting to implementation type
+val instrumented = wrapStreamWithErrorCapture(result, iterationError, streamExhausted)
+// or worse:
+(runs as? RunServiceImpl)?.flush()
+```
+
 ### Use `partition` instead of double `filter`
 
 ```kotlin
@@ -211,7 +230,7 @@ Tests skip gracefully via `assumeTrue` if keys are missing.
 ## Code style
 
 - `toString()` should be single-line, following the `ClassName{field=value, field=value}` convention used by the rest of the SDK.
-- Avoid `@Suppress("UNCHECKED_CAST")` — restructure code to use safe patterns (`as? String`, `is Map<*, *>` with `entries.associate`, etc).
+- Avoid `@Suppress("UNCHECKED_CAST")` — restructure code to use safe patterns (`as? String`, `is Map<*, *>` with `entries.associate`, etc). When unavoidable (e.g. generic type erasure after an `is` check), add a comment explaining why the cast is safe.
 - Use named arguments for constructor/function calls with 2+ parameters, especially when types could be confused:
   ```kotlin
   // Good
@@ -224,6 +243,7 @@ Tests skip gracefully via `assumeTrue` if keys are missing.
   // Bad — positional args are ambiguous
   PromptMessage(PromptMessage.Role.HUMAN, template, templateFormat = templateFormat)
   ```
+- Name functions from the caller's perspective — describe what the caller gets, not what the function does internally. Prefer `stream.withErrorTracking()` over `wrapStreamWithErrorCapture(stream)`.
 - When an `Optional` has a fallback default, use `orElse(default)` directly instead of `orElse(null) ?: default`:
   ```kotlin
   // Good — default goes straight into orElse

langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/TraceProcessIO.kt

@@ -13,10 +13,12 @@ import java.util.function.Function
  *     - **0-arg** ([traceSupplier]): `Map<String, Any?>` (empty map)
  *     - **2-arg** ([traceBiFunction]): `Pair<I1, I2>`
  *     - **3-arg** ([traceTriFunction]): `Triple<I1, I2, I3>`
- * - [PO] — the type passed to `processOutputs`. This is always the raw output type of the traced
- *   function.
+ * - [PO] — the type passed to `processOutputs`. For non-streaming functions, this is the raw return
+ *   type. For streaming functions (when [aggregator] is set), this is the **aggregated output
+ *   type** returned by the aggregator — not the stream type itself.
  *
- * When [inputs] is set, it replaces the default input serialization entirely. Same for [outputs].
+ * When `processInputs` is set, it replaces the default input serialization entirely. Same for
+ * `processOutputs`.
  *
  * ## Example (Kotlin)
  *
@@ -47,12 +49,32 @@ import java.util.function.Function
 class TraceProcessIO<PI, PO>(
     processInputs: Function<PI, Map<String, Any?>>? = null,
     processOutputs: Function<PO, Map<String, Any?>>? = null,
+    aggregator: Function<List<Any?>, Any?>? = null,
 ) {
     // Store erased wrappers so traceable overloads can call these without casts.
     // erase() converts Function<T, R> to Function<Any?, R> — safe because T erases to Object.
     internal val inputsFn: Function<Any?, Map<String, Any?>>? = processInputs?.erase()
     internal val outputsFn: Function<Any?, Map<String, Any?>>? = processOutputs?.erase()
 
+    /**
+     * Aggregator for streaming results. Setting this opts into stream tracing.
+     *
+     * When set and the traced function returns a [java.util.stream.Stream], `traceable` will:
+     * 1. Tee the stream via `peek()` to collect elements as the caller iterates
+     * 2. Register an `onClose()` handler that calls this aggregator with the collected chunks
+     * 3. Pass the aggregated result through [processOutputs] and record it as run output
+     *
+     * The caller gets back a real [java.util.stream.Stream] — no proxy, no type changes.
+     *
+     * Without an aggregator, `Stream` return values are treated normally — the run is completed
+     * immediately on return, same as any other value.
+     *
+     * **Lifecycle:** The run is finalized when the stream's `onClose` handler runs. Callers must
+     * close the stream (via `use {}` in Kotlin or try-with-resources in Java) to ensure the run is
+     * completed. Abandoned streams will leave runs open.
+     */
+    internal val aggregatorFn: Function<List<Any?>, Any?>? = aggregator
+
     companion object {
         /** Creates a new [Builder]. */
         @JvmStatic fun <PI, PO> builder() = Builder<PI, PO>()
@@ -82,6 +104,7 @@ class TraceProcessIO<PI, PO>(
     class Builder<PI, PO> internal constructor() {
         private var processInputs: Function<PI, Map<String, Any?>>? = null
         private var processOutputs: Function<PO, Map<String, Any?>>? = null
+        private var aggregator: Function<List<Any?>, Any?>? = null
 
         /** Callback to transform the input before it is recorded on the run. */
         fun processInputs(processInputs: Function<PI, Map<String, Any?>>) = apply {
@@ -93,9 +116,23 @@ class TraceProcessIO<PI, PO>(
             this.processOutputs = processOutputs
         }
 
+        /**
+         * Aggregator for streaming results. Reduces collected chunks into a single output that is
+         * then passed through [processOutputs].
+         */
+        fun aggregator(aggregator: Function<List<Any?>, Any?>) = apply {
+            this.aggregator = aggregator
+        }
+
         /** Builds the [TraceProcessIO]. */
-        fun build() = TraceProcessIO(processInputs = processInputs, processOutputs = processOutputs)
+        fun build() =
+            TraceProcessIO(
+                processInputs = processInputs,
+                processOutputs = processOutputs,
+                aggregator = aggregator,
+            )
     }
 
-    override fun toString(): String = "TraceProcessIO{inputs=$inputsFn, outputs=$outputsFn}"
+    override fun toString(): String =
+        "TraceProcessIO{inputs=$inputsFn, outputs=$outputsFn, aggregator=$aggregatorFn}"
 }

langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/Traceable.kt

@@ -686,6 +686,54 @@ private fun <T> executeTraced(config: TraceConfig, inputs: Map<String, Any?>?, b
     return CURRENT_RUN.runWith(run) {
         try {
             val result = block()
+
+            // If the result is a java.util.stream.Stream and an aggregator is configured,
+            // instrument it: collect chunks via peek(), finalize the run via onClose().
+            val aggregator = config.processTracedIO?.aggregatorFn
+            if (aggregator != null && result is java.util.stream.Stream<*>) {
+                val chunks = mutableListOf<Any?>()
+                val iterationError = java.util.concurrent.atomic.AtomicReference<Throwable>(null)
+                val streamExhausted = java.util.concurrent.atomic.AtomicBoolean(false)
+                val instrumented =
+                    result
+                        .withErrorTracking(iterationError, streamExhausted)
+                        .peek { chunks.add(it) }
+                        .onClose {
+                            try {
+                                // Always try to aggregate whatever chunks we have
+                                var aggregationError: Throwable? = null
+                                if (chunks.isNotEmpty()) {
+                                    try {
+                                        val output = aggregator.apply(chunks)
+                                        run.outputs =
+                                            applyProcessOutputs(config, output)
+                                                ?: toOutputMap(output)
+                                    } catch (e: Throwable) {
+                                        aggregationError = e
+                                    }
+                                }
+
+                                val error = iterationError.get()
+                                run.endTime = nowIso()
+                                if (error != null) {
+                                    run.error = error.stackTraceToString()
+                                } else if (!streamExhausted.get()) {
+                                    run.error = "Stream cancelled"
+                                } else if (aggregationError != null) {
+                                    run.error = aggregationError.stackTraceToString()
+                                }
+                                run.patchRun()
+                            } catch (e: Throwable) {
+                                run.endTime = nowIso()
+                                run.error = e.stackTraceToString()
+                                run.patchRun()
+                            }
+                        }
+                // Safe: T is Stream<something>, instrumented is Stream<Any?> — same erasure.
+                @Suppress("UNCHECKED_CAST")
+                return@runWith instrumented as T
+            }
+
             run.endTime = nowIso()
             run.outputs = applyProcessOutputs(config, result) ?: toOutputMap(result)
             run.patchRun()
@@ -698,3 +746,32 @@ private fun <T> executeTraced(config: TraceConfig, inputs: Map<String, Any?>?, b
         }
     }
 }
+
+/**
+ * Returns a new [java.util.stream.Stream] that captures exceptions during iteration into [errorRef]
+ * before rethrowing them. This lets the `onClose` handler record the real error instead of
+ * attempting aggregation on partial data.
+ */
+private fun java.util.stream.Stream<*>.withErrorTracking(
+    errorRef: java.util.concurrent.atomic.AtomicReference<Throwable>,
+    exhaustedRef: java.util.concurrent.atomic.AtomicBoolean,
+): java.util.stream.Stream<Any?> {
+    val delegate = this.spliterator()
+    val capturing =
+        object : java.util.Spliterator<Any?> {
+            override fun tryAdvance(action: java.util.function.Consumer<in Any?>): Boolean =
+                try {
+                    delegate.tryAdvance(action).also { if (!it) exhaustedRef.set(true) }
+                } catch (e: Throwable) {
+                    errorRef.compareAndSet(null, e)
+                    throw e
+                }
+
+            override fun trySplit(): java.util.Spliterator<Any?>? = null
+
+            override fun estimateSize(): Long = delegate.estimateSize()
+
+            override fun characteristics(): Int = delegate.characteristics()
+        }
+    return java.util.stream.StreamSupport.stream(capturing, false).onClose { this.close() }
+}