feat: stream thinking with text responses

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/streaming/StreamingPromptRunner.kt

@@ -99,6 +99,17 @@ interface StreamingPromptRunner : PromptRunner {
          */
         fun generateStream(): Flux<String>
 
+        /**
+         * Generate a reactive stream containing text and thinking events.
+         *
+         * The default preserves compatibility for implementations that only
+         * provide text chunks.
+         *
+         * @return Flux emitting thinking and text events in response order
+         */
+        fun generateStreamWithThinking(): Flux<StreamingEvent<String>> =
+            generateStream().map { StreamingEvent.Object(it) }
+
     }
 }
 

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreaming.kt

@@ -38,6 +38,9 @@ internal data class DelegatingStreaming(
         return delegate.generateStream()
     }
 
+    override fun generateStreamWithThinking(): Flux<StreamingEvent<String>> =
+        delegate.generateStreamWithThinking()
+
     override fun <T> createObjectStream(itemClass: Class<T>): Flux<T> =
         delegate.createObjectStream(itemClass)
 

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt

@@ -341,6 +341,17 @@ internal data class OperationContextDelegate(
         )
     }
 
+    override fun generateStreamWithThinking(): Flux<StreamingEvent<String>> {
+        val streamingLlmOperations = streamingFactory().createStreamingOperations(llm)
+
+        return streamingLlmOperations.generateStreamWithThinking(
+            messages = messages,
+            interaction = streamingInteraction(),
+            agentProcess = context.processContext.agentProcess,
+            action = action,
+        )
+    }
+
     override fun <T> createObjectStream(itemClass: Class<T>): Flux<T> {
         val streamingLlmOperations = streamingFactory().createStreamingOperations(llm)
 

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt

@@ -139,6 +139,9 @@ internal interface PromptExecutionDelegate : LlmUse {
 
     fun generateStream(): Flux<String>
 
+    fun generateStreamWithThinking(): Flux<StreamingEvent<String>> =
+        generateStream().map { StreamingEvent.Object(it) }
+
     fun <T> createObjectStream(itemClass: Class<T>): Flux<T>
 
     fun <T> createObjectStreamWithThinking(itemClass: Class<T>): Flux<StreamingEvent<T>>

embabel-agent-api/src/main/kotlin/com/embabel/agent/core/internal/streaming/StreamingLlmOperations.kt

@@ -57,6 +57,21 @@ interface StreamingLlmOperations {
         action: Action?,
     ): Flux<String>
 
+    /**
+     * Generate text and thinking events from messages.
+     *
+     * The default wraps the existing text-only stream so third-party
+     * implementations remain source compatible.
+     */
+    fun generateStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        agentProcess: AgentProcess,
+        action: Action?,
+    ): Flux<StreamingEvent<String>> =
+        generateStream(messages, interaction, agentProcess, action)
+            .map { StreamingEvent.Object(it) }
+
     /**
      * Create a streaming list of objects from JSONL response in the context of an AgentProcess.
      * Each line in the LLM response should be a valid JSON object matching the output class.
@@ -142,6 +157,21 @@ interface StreamingLlmOperations {
         action: Action? = null,
     ): Flux<String>
 
+    /**
+     * Low-level text and thinking stream with optional platform context.
+     *
+     * The default wraps [doTransformStream] for source compatibility.
+     */
+    fun doTransformStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        llmRequestEvent: LlmRequestEvent<String>?,
+        agentProcess: AgentProcess? = null,
+        action: Action? = null,
+    ): Flux<StreamingEvent<String>> =
+        doTransformStream(messages, interaction, llmRequestEvent, agentProcess, action)
+            .map { StreamingEvent.Object(it) }
+
     /**
      * Low level object streaming transform with optional platform context.
      * Streams typed objects as they are parsed from JSONL response.

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/streaming/StreamingChatClientOperations.kt

@@ -31,6 +31,7 @@ import com.embabel.agent.spi.support.springai.ChatClientLlmOperations
 import com.embabel.agent.spi.support.springai.SpringAiLlmService
 import com.embabel.agent.spi.support.springai.toSpringAiMessage
 import com.embabel.agent.spi.support.springai.toSpringToolCallbacks
+import com.embabel.agent.spi.support.streaming.toTaggedThinkingEvents
 import com.embabel.chat.Message
 import com.embabel.common.ai.converters.streaming.StreamingJacksonOutputConverter
 import com.embabel.common.core.streaming.StreamingEvent
@@ -113,6 +114,14 @@ internal class StreamingChatClientOperations(
         return doTransformStream(messages, interaction, null, agentProcess, action)
     }
 
+    override fun generateStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        agentProcess: AgentProcess,
+        action: Action?,
+    ): Flux<StreamingEvent<String>> =
+        doTransformStreamWithThinking(messages, interaction, null, agentProcess, action)
+
     override fun <O> createObjectStream(
         messages: List<Message>,
         interaction: LlmInteraction,
@@ -190,6 +199,16 @@ internal class StreamingChatClientOperations(
         )
     }
 
+    override fun doTransformStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        llmRequestEvent: LlmRequestEvent<String>?,
+        agentProcess: AgentProcess?,
+        action: Action?,
+    ): Flux<StreamingEvent<String>> =
+        doTransformStream(messages, interaction, llmRequestEvent, agentProcess, action)
+            .toTaggedThinkingEvents()
+
     /**
      * Creates a stream of typed objects from LLM JSONL responses, with thinking content suppressed.
      *

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/StreamingLlmOperationsImpl.kt

@@ -75,6 +75,14 @@ internal class StreamingLlmOperationsImpl(
         return doTransformStream(messages, interaction, null, agentProcess, action)
     }
 
+    override fun generateStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        agentProcess: AgentProcess,
+        action: Action?,
+    ): Flux<StreamingEvent<String>> =
+        doTransformStreamWithThinking(messages, interaction, null, agentProcess, action)
+
     override fun <O> createObjectStream(
         messages: List<Message>,
         interaction: LlmInteraction,
@@ -137,6 +145,16 @@ internal class StreamingLlmOperationsImpl(
         return messageStreamer.stream(messagesWithContributions, tools, interaction.toolCallInspectors)
     }
 
+    override fun doTransformStreamWithThinking(
+        messages: List<Message>,
+        interaction: LlmInteraction,
+        llmRequestEvent: LlmRequestEvent<String>?,
+        agentProcess: AgentProcess?,
+        action: Action?,
+    ): Flux<StreamingEvent<String>> =
+        doTransformStream(messages, interaction, llmRequestEvent, agentProcess, action)
+            .toTaggedThinkingEvents()
+
     override fun <O> doTransformObjectStream(
         messages: List<Message>,
         interaction: LlmInteraction,

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/streaming/ThinkingStreamSupport.kt

@@ -0,0 +1,176 @@
+/*
+ * Copyright 2024-2026 Embabel Pty Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.embabel.agent.spi.support.streaming
+
+import com.embabel.common.core.streaming.StreamingEvent
+import com.embabel.common.core.thinking.ThinkingTags
+import reactor.core.publisher.Flux
+
+/**
+ * Extracts tagged thinking from text while preserving response order across
+ * arbitrary chunk boundaries. Parser state is scoped to each subscription.
+ */
+internal fun Flux<String>.toTaggedThinkingEvents(): Flux<StreamingEvent<String>> =
+    Flux.defer {
+        val parser = TaggedThinkingParser()
+        this@toTaggedThinkingEvents
+            .concatMap { Flux.fromIterable(parser.accept(it)) }
+            .concatWith(Flux.defer { Flux.fromIterable(parser.finish()) })
+    }
+
+private class TaggedThinkingParser {
+
+    private data class Tag(
+        val start: String,
+        val end: String,
+    )
+
+    private data class StartMatch(
+        val index: Int,
+        val tag: Tag?,
+        val legacy: Boolean,
+    )
+
+    private val tags = ThinkingTags.TAG_DEFINITIONS
+        .filterKeys { it != "legacy_prefix" && it != "no_prefix" }
+        .values
+        .map { Tag(it.first, it.second) }
+
+    private val legacyPrefix = ThinkingTags.TAG_DEFINITIONS["legacy_prefix"]?.first.orEmpty()
+    private val buffer = StringBuilder()
+    private var activeTag: Tag? = null
+    private var atLineStart = true
+
+    fun accept(text: String): List<StreamingEvent<String>> {
+        buffer.append(text)
+        val events = mutableListOf<StreamingEvent<String>>()
+
+        while (buffer.isNotEmpty()) {
+            val tag = activeTag
+            if (tag != null) {
+                val endIndex = buffer.indexOf(tag.end)
+                if (endIndex < 0) break
+                val thinking = buffer.substring(0, endIndex).trim()
+                if (thinking.isNotEmpty()) events += StreamingEvent.Thinking(thinking)
+                buffer.delete(0, endIndex + tag.end.length)
+                activeTag = null
+                atLineStart = false
+                continue
+            }
+
+            if (atLineStart && legacyPrefix.isNotEmpty() && buffer.startsWith(legacyPrefix)) {
+                val newlineIndex = buffer.indexOf("\n")
+                if (newlineIndex < 0) break
+                val thinking = buffer.substring(legacyPrefix.length, newlineIndex).trim()
+                if (thinking.isNotEmpty()) events += StreamingEvent.Thinking(thinking)
+                buffer.delete(0, newlineIndex + 1)
+                atLineStart = true
+                continue
+            }
+
+            val match = findStart()
+            if (match != null) {
+                if (match.index > 0) {
+                    emitText(events, buffer.substring(0, match.index))
+                    buffer.delete(0, match.index)
+                    continue
+                }
+                if (match.legacy) break
+                activeTag = match.tag
+                buffer.delete(0, match.tag!!.start.length)
+                atLineStart = false
+                continue
+            }
+
+            val heldSuffixLength = longestPartialStartSuffix()
+            val emitLength = buffer.length - heldSuffixLength
+            if (emitLength > 0) {
+                emitText(events, buffer.substring(0, emitLength))
+                buffer.delete(0, emitLength)
+            }
+            break
+        }
+
+        return events
+    }
+
+    fun finish(): List<StreamingEvent<String>> {
+        val events = mutableListOf<StreamingEvent<String>>()
+        val tag = activeTag
+        if (tag != null) {
+            emitText(events, tag.start + buffer.toString())
+        } else if (atLineStart && legacyPrefix.isNotEmpty() && buffer.startsWith(legacyPrefix)) {
+            val thinking = buffer.substring(legacyPrefix.length).trim()
+            if (thinking.isNotEmpty()) events += StreamingEvent.Thinking(thinking)
+        } else {
+            emitText(events, buffer.toString())
+        }
+        buffer.clear()
+        activeTag = null
+        atLineStart = true
+        return events
+    }
+
+    private fun findStart(): StartMatch? {
+        val tagMatch = tags
+            .mapNotNull { tag -> buffer.indexOf(tag.start).takeIf { it >= 0 }?.let { StartMatch(it, tag, false) } }
+            .minWithOrNull(compareBy<StartMatch> { it.index }.thenByDescending { it.tag?.start?.length ?: 0 })
+        val legacyMatch = if (legacyPrefix.isNotEmpty()) {
+            findLegacyStart()?.let { StartMatch(it, null, true) }
+        } else {
+            null
+        }
+        return listOfNotNull(tagMatch, legacyMatch).minByOrNull { it.index }
+    }
+
+    private fun findLegacyStart(): Int? {
+        var fromIndex = 0
+        while (fromIndex < buffer.length) {
+            val index = buffer.indexOf(legacyPrefix, fromIndex)
+            if (index < 0) return null
+            if ((index == 0 && atLineStart) || (index > 0 && buffer[index - 1] == '\n')) return index
+            fromIndex = index + 1
+        }
+        return null
+    }
+
+    private fun longestPartialStartSuffix(): Int {
+        val tagSuffixLength = tags.map { it.start }.maxOfOrNull { start ->
+            (1 until start.length)
+                .filter { length -> buffer.endsWith(start.substring(0, length)) }
+                .maxOrNull() ?: 0
+        } ?: 0
+        return maxOf(tagSuffixLength, longestPartialLegacySuffix())
+    }
+
+    private fun longestPartialLegacySuffix(): Int {
+        if (legacyPrefix.isEmpty()) return 0
+        return (1 until legacyPrefix.length)
+            .filter { length ->
+                if (!buffer.endsWith(legacyPrefix.substring(0, length))) return@filter false
+                val startIndex = buffer.length - length
+                (startIndex == 0 && atLineStart) || (startIndex > 0 && buffer[startIndex - 1] == '\n')
+            }
+            .maxOrNull() ?: 0
+    }
+
+    private fun emitText(events: MutableList<StreamingEvent<String>>, text: String) {
+        if (text.isNotEmpty()) {
+            events += StreamingEvent.Object(text)
+            atLineStart = text.endsWith("\n")
+        }
+    }
+}

embabel-agent-api/src/test/java/com/embabel/agent/api/streaming/StreamingThinkingJavaApiTest.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2024-2026 Embabel Pty Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.embabel.agent.api.streaming;
+
+import com.embabel.agent.api.common.streaming.StreamingPromptRunner;
+import com.embabel.common.core.streaming.StreamingEvent;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class StreamingThinkingJavaApiTest {
+
+    @Test
+    void exposesThinkingAwareTextStreamingToJava() throws NoSuchMethodException {
+        var method = StreamingPromptRunner.Streaming.class.getMethod("generateStreamWithThinking");
+
+        assertThat(method.getReturnType()).isEqualTo(Flux.class);
+        assertThat(method.isDefault()).isTrue();
+    }
+
+    Flux<StreamingEvent<String>> callFromJava(StreamingPromptRunner.Streaming streaming) {
+        return streaming.generateStreamWithThinking();
+    }
+}