Skip to content

fix(prompt): preserve streamed tool-call identity and arguments #2272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kyboi wants to merge 2 commits into
base: develop
Choose a base branch
from
+442 −38
Open

fix(prompt): preserve streamed tool-call identity and arguments #2272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3105225
fix(openai): preserve Responses streaming tool identity
kyboi Sep 29, 2026
050d91f
fix(prompt): preserve interleaved streaming tool calls
kyboi Sep 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions ...ai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/OpenAILLMClient.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@ import kotlinx.coroutines.async
import kotlinx.coroutines.awaitAll
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.withContext
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonPrimitive
Expand Down Expand Up @@ -348,7 +350,7 @@ public open class OpenAILLMClient @JvmOverloads constructor(
model: LLModel,
tools: List<ToolDescriptor>,
params: OpenAIResponsesParams
): Flow<StreamFrame> {
): Flow<StreamFrame> = flow {
logger.debug { "Executing streaming prompt: $prompt with model: $model" }

val llmTools = tools.takeIf { it.isNotEmpty() }?.map {
Expand All @@ -368,8 +370,9 @@ public open class OpenAILLMClient @JvmOverloads constructor(
params = params,
stream = true
)
val toolCallsByItemId = mutableMapOf<String, Item.FunctionToolCall>()

return try {
val frames = try {
httpClient.sse(
path = settings.responsesAPIPath,
requestBody = request,
Expand All @@ -379,6 +382,14 @@ public open class OpenAILLMClient @JvmOverloads constructor(
},
processStreamingChunk = {
when (it) {
is OpenAIStreamEvent.ResponseOutputItemAdded -> {
val item = it.item
if (item is Item.FunctionToolCall) {
item.id?.let { itemId -> toolCallsByItemId[itemId] = item }
}
null
}

is OpenAIStreamEvent.ResponseOutputTextDelta -> {
StreamFrame.TextDelta(text = it.delta, index = it.outputIndex)
}
Expand All @@ -394,9 +405,10 @@ public open class OpenAILLMClient @JvmOverloads constructor(
}

is OpenAIStreamEvent.ResponseFunctionCallArgumentsDelta -> {
val toolCall = toolCallsByItemId[it.itemId]
StreamFrame.ToolCallDelta(
id = it.itemId,
name = null,
id = toolCall?.callId ?: it.itemId,
name = toolCall?.name,
content = it.delta,
index = it.outputIndex
)
Expand Down Expand Up @@ -425,6 +437,7 @@ public open class OpenAILLMClient @JvmOverloads constructor(
}

is Item.FunctionToolCall -> {
item.id?.let { itemId -> toolCallsByItemId.remove(itemId) }
StreamFrame.ToolCallComplete(
id = item.callId,
name = item.name,
Expand Down Expand Up @@ -464,6 +477,7 @@ public open class OpenAILLMClient @JvmOverloads constructor(
cause = e
)
}
emitAll(frames)
}

override suspend fun executeMultipleChoices(
Expand Down
244 changes: 244 additions & 0 deletions ...src/jvmTest/kotlin/ai/koog/prompt/executor/clients/openai/OpenAIPrimaryConstructorTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,23 @@ import ai.koog.test.utils.CapturingKoogHttpClient
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.take
import kotlinx.coroutines.flow.toList
import kotlinx.coroutines.test.runTest
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonNamingStrategy
import kotlin.reflect.KClass
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertIs

class OpenAIPrimaryConstructorTest {
private val streamingJson = Json {
encodeDefaults = true
explicitNulls = false
namingStrategy = JsonNamingStrategy.SnakeCase
}

private val responseJson = """
{
"id": "chatcmpl-123",
Expand Down Expand Up @@ -220,4 +229,239 @@ class OpenAIPrimaryConstructorTest {
assertEquals(inputTokens, end.metaInfo.inputTokensCount)
assertEquals(outputTokens, end.metaInfo.outputTokensCount)
}

@Test
fun testResponsesStreamingToolCallDeltasUseCanonicalIdentity() = runTest {
val firstItemId = "item_first"
val firstCallId = "call_first"
val firstName = "firstTool"
val secondItemId = "item_second"
val secondCallId = "call_second"
val secondName = "secondTool"
val events = listOf(
OpenAIStreamEvent.ResponseOutputItemAdded(
item = Item.FunctionToolCall(
arguments = "",
callId = firstCallId,
name = firstName,
id = firstItemId,
status = OpenAIInputStatus.IN_PROGRESS
),
outputIndex = 0,
sequenceNumber = 1
),
OpenAIStreamEvent.ResponseFunctionCallArgumentsDelta(
itemId = firstItemId,
outputIndex = 0,
delta = "{\"first\":",
sequenceNumber = 2
),
OpenAIStreamEvent.ResponseOutputItemAdded(
item = Item.FunctionToolCall(
arguments = "",
callId = secondCallId,
name = secondName,
id = secondItemId,
status = OpenAIInputStatus.IN_PROGRESS
),
outputIndex = 1,
sequenceNumber = 3
),
OpenAIStreamEvent.ResponseFunctionCallArgumentsDelta(
itemId = secondItemId,
outputIndex = 1,
delta = "{\"second\":true}",
sequenceNumber = 4
),
OpenAIStreamEvent.ResponseFunctionCallArgumentsDelta(
itemId = firstItemId,
outputIndex = 0,
delta = "true}",
sequenceNumber = 5
),
OpenAIStreamEvent.ResponseOutputItemDone(
item = Item.FunctionToolCall(
arguments = "{\"second\":true}",
callId = secondCallId,
name = secondName,
id = secondItemId,
status = OpenAIInputStatus.COMPLETED
),
outputIndex = 1,
sequenceNumber = 6
),
OpenAIStreamEvent.ResponseOutputItemDone(
item = Item.FunctionToolCall(
arguments = "{\"first\":true}",
callId = firstCallId,
name = firstName,
id = firstItemId,
status = OpenAIInputStatus.COMPLETED
),
outputIndex = 0,
sequenceNumber = 7
),
completedResponseEvent(sequenceNumber = 8)
)
val client = OpenAILLMClient(
settings = OpenAIClientSettings(baseUrl = "https://unused.test"),
httpClient = streamingTransport { events }
)

val frames = client.executeStreaming(
prompt = Prompt(
messages = listOf(Message.User("Use both tools", RequestMetaInfo.Empty)),
id = "test",
params = OpenAIResponsesParams()
),
model = OpenAIModels.Chat.GPT4o
).toList()

assertEquals(
listOf(
StreamFrame.ToolCallDelta(firstCallId, firstName, "{\"first\":", 0),
StreamFrame.ToolCallDelta(secondCallId, secondName, "{\"second\":true}", 1),
StreamFrame.ToolCallDelta(firstCallId, firstName, "true}", 0),
StreamFrame.ToolCallComplete(secondCallId, secondName, "{\"second\":true}", 1),
StreamFrame.ToolCallComplete(firstCallId, firstName, "{\"first\":true}", 0)
),
frames.dropLast(1)
)
assertIs<StreamFrame.End>(frames.last())
}

@Test
fun testResponsesStreamingFallsBackForUnknownAndCompletedItems() = runTest {
val events = listOf(
OpenAIStreamEvent.ResponseOutputItemAdded(Item.Text("ignored"), 0, 1),
OpenAIStreamEvent.ResponseOutputItemAdded(toolCall(id = null), 0, 2),
argumentDelta("unknown", 3),
OpenAIStreamEvent.ResponseOutputItemAdded(toolCall(), 0, 4),
argumentDelta("item", 5),
OpenAIStreamEvent.ResponseOutputItemDone(toolCall(), 0, 6),
argumentDelta("item", 7),
OpenAIStreamEvent.ResponseOutputItemDone(toolCall(id = null), 0, 8),
completedResponseEvent(9)
)
val frames = toolCallStream(streamingTransport { events }).toList()

assertEquals(
listOf(
StreamFrame.ToolCallDelta("unknown", null, "{}", 0),
StreamFrame.ToolCallDelta("call", "lookup", "{}", 0),
StreamFrame.ToolCallComplete("call", "lookup", "{}", 0),
StreamFrame.ToolCallDelta("item", null, "{}", 0),
StreamFrame.ToolCallComplete("call", "lookup", "{}", 0)
),
frames.dropLast(1)
)
assertIs<StreamFrame.End>(frames.last())
}

@Test
fun testResponsesStreamingDoesNotRetainIdentityAfterCancelledCollection() = runTest {
var collection = 0
val stream = toolCallStream(
streamingTransport {
collection++
if (collection == 1) {
listOf(
OpenAIStreamEvent.ResponseOutputItemAdded(toolCall(), 0, 1),
argumentDelta("item", 2),
completedResponseEvent(3)
)
} else {
listOf(argumentDelta("item", 1), completedResponseEvent(2))
}
}
)

assertEquals(listOf(StreamFrame.ToolCallDelta("call", "lookup", "{}", 0)), stream.take(1).toList())
val recollected = stream.toList()
assertEquals(StreamFrame.ToolCallDelta("item", null, "{}", 0), recollected.first())
assertIs<StreamFrame.End>(recollected.last())
}

private fun toolCall(id: String? = "item") = Item.FunctionToolCall(
arguments = "{}",
callId = "call",
name = "lookup",
id = id
)

private fun argumentDelta(itemId: String, sequenceNumber: Int) =
OpenAIStreamEvent.ResponseFunctionCallArgumentsDelta(itemId, 0, "{}", sequenceNumber)

private fun toolCallStream(transport: KoogHttpClient): Flow<StreamFrame> = OpenAILLMClient(
settings = OpenAIClientSettings(baseUrl = "https://unused.test"),
httpClient = transport
).executeStreaming(
prompt = Prompt(
messages = listOf(Message.User("Use lookup", RequestMetaInfo.Empty)),
id = "test",
params = OpenAIResponsesParams()
),
model = OpenAIModels.Chat.GPT4o
)

private fun streamingTransport(events: () -> List<OpenAIStreamEvent>): KoogHttpClient = object : KoogHttpClient {
override val clientName: String = "StreamingOpenAIClient"

override suspend fun <R : Any> get(
path: String,
responseType: KClass<R>,
parameters: Map<String, String>,
headers: Map<String, String>,
): R = error("GET is not expected in this test")

override suspend fun <T : Any, R : Any> post(
path: String,
requestBody: T,
requestBodyType: KClass<T>,
responseType: KClass<R>,
parameters: Map<String, String>,
headers: Map<String, String>,
): R = error("POST is not expected in this test")

override fun <T : Any, R : Any, O : Any> sse(
path: String,
requestBody: T,
requestBodyType: KClass<T>,
dataFilter: (String?) -> Boolean,
decodeStreamingResponse: (String) -> R,
processStreamingChunk: (R) -> O?,
parameters: Map<String, String>,
headers: Map<String, String>,
): Flow<O> = flow {
events().forEach { event ->
val data = streamingJson.encodeToString(OpenAIStreamEvent.serializer(), event)
processStreamingChunk(decodeStreamingResponse(data))?.let { emit(it) }
}
}

override fun <T : Any> lines(
path: String,
requestBody: T,
requestBodyType: KClass<T>,
parameters: Map<String, String>,
headers: Map<String, String>,
): Flow<String> = error("lines is not expected in this test")

override fun close(): Unit = Unit
}

private fun completedResponseEvent(sequenceNumber: Int): OpenAIStreamEvent.ResponseCompleted =
OpenAIStreamEvent.ResponseCompleted(
response = OpenAIResponsesAPIResponse(
created = 1716920005,
id = "response_id",
model = "gpt-5",
output = emptyList(),
parallelToolCalls = true,
status = OpenAIInputStatus.COMPLETED,
text = OpenAITextConfig(),
usage = null
),
sequenceNumber = sequenceNumber
)
}
Loading