Skip to content

Commit 3adf4ac

Browse files
sdubovsdubov
committed
fix(agents): KG-808. SpanAdapter hooks run on fully prepared spans in OpenTelemetry feature
- `OpenTelemetry.kt` previously invoked `SpanAdapter.onBeforeSpanFinished` before `end<span_name>Span` had populated final attributes, so adapters such as Langfuse and Weave never observed them. - Each start and end span helpers now take a `SpanAdapter` and invokes the corresponding hook internally, after all attributes are set and right before the span is started or ended. - MCP enrichment for tool-call spans moved inside `startExecuteToolSpan` via a new optional `mcpToolMetadata` parameter; `mcpClientSpan.kt` is gone, its helper is now private inside `executeToolSpan.kt`. closes: [KG-808](https://youtrack.jetbrains.com/issue/KG-808)
1 parent cc3ebc5 commit 3adf4ac

11 files changed

Lines changed: 352 additions & 165 deletions

File tree

‎agents/agents-features/agents-features-opentelemetry/src/commonMain/kotlin/ai/koog/agents/features/opentelemetry/feature/OpenTelemetry.kt‎

Lines changed: 41 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ import ai.koog.agents.features.opentelemetry.span.endInvokeAgentSpan
3838
import ai.koog.agents.features.opentelemetry.span.endNodeExecuteSpan
3939
import ai.koog.agents.features.opentelemetry.span.endStrategySpan
4040
import ai.koog.agents.features.opentelemetry.span.endSubgraphExecuteSpan
41-
import ai.koog.agents.features.opentelemetry.span.enrichExecuteToolSpanWithMcpAttrs
4241
import ai.koog.agents.features.opentelemetry.span.startCreateAgentSpan
4342
import ai.koog.agents.features.opentelemetry.span.startExecuteToolSpan
4443
import ai.koog.agents.features.opentelemetry.span.startInferenceSpan
@@ -114,10 +113,10 @@ public class OpenTelemetry {
114113
id = eventContext.eventId,
115114
runId = eventContext.context.runId,
116115
nodeId = eventContext.node.id,
117-
nodeInput = nodeInput
116+
nodeInput = nodeInput,
117+
spanAdapter = spanAdapter,
118118
)
119119

120-
spanAdapter?.onBeforeSpanStarted(nodeExecuteSpan)
121120
spanCollector.collectSpan(
122121
span = nodeExecuteSpan,
123122
path = patchedExecutionInfo
@@ -137,11 +136,11 @@ public class OpenTelemetry {
137136

138137
val nodeOutput = nodeDataToString(eventContext.output, eventContext.outputType, pipeline.config.serializer)
139138

140-
spanAdapter?.onBeforeSpanFinished(nodeExecuteSpan)
141139
endNodeExecuteSpan(
142140
span = nodeExecuteSpan,
143141
nodeOutput = nodeOutput,
144-
verbose = config.isVerbose
142+
verbose = config.isVerbose,
143+
spanAdapter = spanAdapter,
145144
)
146145
spanCollector.removeSpan(
147146
span = nodeExecuteSpan,
@@ -160,12 +159,12 @@ public class OpenTelemetry {
160159
spanType = SpanType.NODE
161160
) ?: return@intercept
162161

163-
spanAdapter?.onBeforeSpanFinished(nodeExecuteSpan)
164162
endNodeExecuteSpan(
165163
span = nodeExecuteSpan,
166164
nodeOutput = null,
167165
error = eventContext.error,
168-
verbose = config.isVerbose
166+
verbose = config.isVerbose,
167+
spanAdapter = spanAdapter,
169168
)
170169
spanCollector.removeSpan(
171170
span = nodeExecuteSpan,
@@ -194,10 +193,10 @@ public class OpenTelemetry {
194193
id = eventContext.eventId,
195194
runId = eventContext.context.runId,
196195
subgraphId = eventContext.subgraph.id,
197-
subgraphInput = subgraphInput
196+
subgraphInput = subgraphInput,
197+
spanAdapter = spanAdapter,
198198
)
199199

200-
spanAdapter?.onBeforeSpanStarted(subgraphExecuteSpan)
201200
spanCollector.collectSpan(
202201
span = subgraphExecuteSpan,
203202
path = patchedExecutionInfo
@@ -217,11 +216,11 @@ public class OpenTelemetry {
217216

218217
val subgraphOutput = nodeDataToString(eventContext.output, eventContext.outputType, pipeline.config.serializer)
219218

220-
spanAdapter?.onBeforeSpanFinished(subgraphExecuteSpan)
221219
endSubgraphExecuteSpan(
222220
span = subgraphExecuteSpan,
223221
subgraphOutput = subgraphOutput,
224-
verbose = config.isVerbose
222+
verbose = config.isVerbose,
223+
spanAdapter = spanAdapter,
225224
)
226225
spanCollector.removeSpan(
227226
span = subgraphExecuteSpan,
@@ -240,12 +239,12 @@ public class OpenTelemetry {
240239
spanType = SpanType.SUBGRAPH
241240
) ?: return@intercept
242241

243-
spanAdapter?.onBeforeSpanFinished(subgraphExecuteSpan)
244242
endSubgraphExecuteSpan(
245243
span = subgraphExecuteSpan,
246244
subgraphOutput = null,
247245
error = eventContext.error,
248-
verbose = config.isVerbose
246+
verbose = config.isVerbose,
247+
spanAdapter = spanAdapter,
249248
)
250249
spanCollector.removeSpan(
251250
span = subgraphExecuteSpan,
@@ -311,10 +310,10 @@ public class OpenTelemetry {
311310
id = eventContext.eventId,
312311
model = eventContext.agent.agentConfig.model,
313312
agentId = eventContext.context.agentId,
314-
messages = messages
313+
messages = messages,
314+
spanAdapter = spanAdapter,
315315
)
316316

317-
spanAdapter?.onBeforeSpanStarted(createAgentSpan)
318317
spanCollector.collectSpan(
319318
span = createAgentSpan,
320319
path = eventContext.executionInfo
@@ -331,10 +330,10 @@ public class OpenTelemetry {
331330
runId = eventContext.runId,
332331
llmParams = eventContext.agent.agentConfig.prompt.params,
333332
messages = messages,
334-
tools = tools
333+
tools = tools,
334+
spanAdapter = spanAdapter,
335335
)
336336

337-
spanAdapter?.onBeforeSpanStarted(invokeAgentSpan)
338337
// Patch the agent execution info to include runId in the path.
339338
// This is required to create a path structure that matches the span structure in the OTel feature.
340339
spanCollector.collectSpan(
@@ -360,12 +359,12 @@ public class OpenTelemetry {
360359
)
361360
}
362361

363-
spanAdapter?.onBeforeSpanFinished(invokeAgentSpan)
364362
endInvokeAgentSpan(
365363
span = invokeAgentSpan,
366364
messages = eventContext.context.config.prompt.messages.toList(),
367365
model = eventContext.context.config.model,
368-
verbose = config.isVerbose
366+
verbose = config.isVerbose,
367+
spanAdapter = spanAdapter,
369368
)
370369
spanCollector.removeSpan(
371370
span = invokeAgentSpan,
@@ -395,13 +394,13 @@ public class OpenTelemetry {
395394
spanType = SpanType.INVOKE_AGENT
396395
) ?: return@intercept
397396

398-
spanAdapter?.onBeforeSpanFinished(invokeAgentSpan)
399397
endInvokeAgentSpan(
400398
span = invokeAgentSpan,
401399
messages = eventContext.context.config.prompt.messages.toList(),
402400
model = eventContext.context.config.model,
403401
error = eventContext.error,
404-
verbose = config.isVerbose
402+
verbose = config.isVerbose,
403+
spanAdapter = spanAdapter,
405404
)
406405
spanCollector.removeSpan(
407406
span = invokeAgentSpan,
@@ -428,10 +427,10 @@ public class OpenTelemetry {
428427
spanType = SpanType.CREATE_AGENT
429428
) ?: return@intercept
430429

431-
spanAdapter?.onBeforeSpanFinished(agentSpan)
432430
endCreateAgentSpan(
433431
span = agentSpan,
434-
verbose = config.isVerbose
432+
verbose = config.isVerbose,
433+
spanAdapter = spanAdapter,
435434
)
436435

437436
spanCollector.removeSpan(
@@ -468,10 +467,10 @@ public class OpenTelemetry {
468467
parentSpan = parentSpan,
469468
id = eventContext.eventId,
470469
runId = eventContext.context.runId,
471-
strategyName = eventContext.strategy.name
470+
strategyName = eventContext.strategy.name,
471+
spanAdapter = spanAdapter,
472472
)
473473

474-
spanAdapter?.onBeforeSpanStarted(strategySpan)
475474
spanCollector.collectSpan(
476475
span = strategySpan,
477476
path = patchedExecutionInfo
@@ -487,8 +486,11 @@ public class OpenTelemetry {
487486
spanType = SpanType.STRATEGY
488487
) ?: return@intercept
489488

490-
spanAdapter?.onBeforeSpanFinished(strategySpan)
491-
endStrategySpan(span = strategySpan, verbose = config.isVerbose)
489+
endStrategySpan(
490+
span = strategySpan,
491+
verbose = config.isVerbose,
492+
spanAdapter = spanAdapter,
493+
)
492494
spanCollector.removeSpan(
493495
span = strategySpan,
494496
path = patchedExecutionInfo
@@ -522,11 +524,10 @@ public class OpenTelemetry {
522524
model = eventContext.model,
523525
messages = messages,
524526
llmParams = eventContext.prompt.params,
525-
tools = eventContext.tools
527+
tools = eventContext.tools,
528+
spanAdapter = spanAdapter,
526529
)
527530

528-
// Start span
529-
spanAdapter?.onBeforeSpanStarted(inferenceSpan)
530531
spanCollector.collectSpan(
531532
span = inferenceSpan,
532533
path = patchedExecutionInfo
@@ -563,20 +564,13 @@ public class OpenTelemetry {
563564
)
564565
}
565566

566-
// Pre-populate gen_ai.output.messages so SpanAdapter implementations (Langfuse, Weave)
567-
// can reshape the response messages before the span ends.
568-
// endInferenceSpan re-adds the same attribute idempotently.
569-
if (eventContext.responses.isNotEmpty()) {
570-
inferenceSpan.addAttribute(GenAIAttributes.Output.Messages(eventContext.responses))
571-
}
572-
573567
// Stop InferenceSpan
574-
spanAdapter?.onBeforeSpanFinished(inferenceSpan)
575568
endInferenceSpan(
576569
span = inferenceSpan,
577570
messages = eventContext.responses,
578571
model = eventContext.model,
579-
verbose = config.isVerbose
572+
verbose = config.isVerbose,
573+
spanAdapter = spanAdapter,
580574
)
581575
spanCollector.removeSpan(
582576
span = inferenceSpan,
@@ -629,13 +623,13 @@ public class OpenTelemetry {
629623
spanType = SpanType.INFERENCE
630624
) ?: return@intercept
631625

632-
spanAdapter?.onBeforeSpanFinished(inferenceSpan)
633626
endInferenceSpan(
634627
span = inferenceSpan,
635628
messages = emptyList(),
636629
model = eventContext.model,
630+
error = eventContext.error,
637631
verbose = config.isVerbose,
638-
error = eventContext.error
632+
spanAdapter = spanAdapter,
639633
)
640634
spanCollector.removeSpan(
641635
span = inferenceSpan,
@@ -664,32 +658,11 @@ public class OpenTelemetry {
664658
toolName = eventContext.toolName,
665659
toolArgs = eventContext.toolArgs.toKotlinxJsonObject(),
666660
toolDescription = eventContext.toolDescription,
667-
toolCallId = eventContext.toolCallId
668-
)
669-
val mcpToolMetadata = eventContext.context.llm.toolRegistry.getMcpToolMeta(eventContext.toolName)
670-
if (mcpToolMetadata != null) {
671-
val mcpVersion = mcpToolMetadata[McpMetadataKeys.McpProtocolVersion]
672-
val mcpTransportType = mcpToolMetadata[McpMetadataKeys.McpTransportType]
673-
if (mcpVersion != null && mcpTransportType != null) {
674-
executeToolSpan.enrichExecuteToolSpanWithMcpAttrs(
675-
toolName = eventContext.toolName,
676-
sessionId = mcpToolMetadata[McpMetadataKeys.McpSessionId],
677-
methodName = "tools/call",
678-
serverPort = mcpToolMetadata[McpMetadataKeys.ServerPort]?.toIntOrNull(),
679-
serverAddress = mcpToolMetadata[McpMetadataKeys.ServerUrl],
680-
mcpProtocolVersion = mcpVersion,
681-
mcpTransportType = mcpTransportType,
682-
)
683-
} else {
684-
logger.error {
685-
"MCP protocol version ($mcpVersion) and transport type ($mcpTransportType) are required " +
686-
"for mcp tool call spans: tool=${eventContext.toolName}, " +
687-
"serverUrl=${mcpToolMetadata[McpMetadataKeys.ServerUrl]}"
688-
}
689-
}
690-
}
661+
toolCallId = eventContext.toolCallId,
662+
mcpToolMetadata = eventContext.context.llm.toolRegistry.getMcpToolMeta(eventContext.toolName),
663+
spanAdapter = spanAdapter,
664+
)
691665

692-
spanAdapter?.onBeforeSpanStarted(executeToolSpan)
693666
spanCollector.collectSpan(executeToolSpan, path)
694667

695668
// Metrics
@@ -918,12 +891,12 @@ public class OpenTelemetry {
918891
spanType = SpanType.EXECUTE_TOOL
919892
) ?: return
920893

921-
spanAdapter?.onBeforeSpanFinished(span = span)
922894
endExecuteToolSpan(
923895
span = span,
924896
toolResult = toolResult?.toKotlinxJsonElement(),
925897
error = error,
926-
verbose = config.isVerbose
898+
verbose = config.isVerbose,
899+
spanAdapter = spanAdapter,
927900
)
928901
spanCollector.removeSpan(
929902
span = span,

‎agents/agents-features/agents-features-opentelemetry/src/commonMain/kotlin/ai/koog/agents/features/opentelemetry/span/createAgentSpan.kt‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import ai.koog.agents.features.opentelemetry.attribute.KoogAttributes
55
import ai.koog.agents.features.opentelemetry.extension.addCommonErrorAttributes
66
import ai.koog.agents.features.opentelemetry.extension.systemMessages
77
import ai.koog.agents.features.opentelemetry.extension.toStatusData
8+
import ai.koog.agents.features.opentelemetry.integration.SpanAdapter
89
import ai.koog.prompt.llm.LLModel
910
import ai.koog.prompt.message.Message
1011
import io.opentelemetry.kotlin.factory.ContextFactory
@@ -30,6 +31,15 @@ import io.opentelemetry.kotlin.tracing.Tracer
3031
*
3132
* Custom attributes:
3233
* - koog.event.id
34+
*
35+
* @param tracer The tracer to use
36+
* @param contextFactory The context factory to use
37+
* @param parentSpan The parent span to use
38+
* @param id The id of the span
39+
* @param agentId The id of the agent
40+
* @param model The model to use
41+
* @param messages The messages to use
42+
* @param spanAdapter The span adapter to use
3343
*/
3444
internal fun startCreateAgentSpan(
3545
tracer: Tracer,
@@ -38,7 +48,8 @@ internal fun startCreateAgentSpan(
3848
id: String,
3949
agentId: String,
4050
model: LLModel,
41-
messages: List<Message>
51+
messages: List<Message>,
52+
spanAdapter: SpanAdapter? = null,
4253
): GenAIAgentSpan {
4354
val builder = GenAIAgentSpanBuilder(
4455
spanType = SpanType.CREATE_AGENT,
@@ -68,7 +79,9 @@ internal fun startCreateAgentSpan(
6879

6980
builder.addAttribute(KoogAttributes.Koog.Event.Id(id))
7081

71-
return builder.buildAndStart(tracer, contextFactory)
82+
val span = builder.buildAndStart(tracer, contextFactory)
83+
spanAdapter?.onBeforeSpanStarted(span)
84+
return span
7285
}
7386

7487
/**
@@ -79,11 +92,17 @@ internal fun startCreateAgentSpan(
7992
*
8093
* Span attribute:
8194
* - error.type (conditional)
95+
*
96+
* @param span The span to end
97+
* @param error The error to set as the span status
98+
* @param verbose Whether to log verbose information
99+
* @param spanAdapter The span adapter to use
82100
*/
83101
internal fun endCreateAgentSpan(
84102
span: GenAIAgentSpan,
85103
error: Throwable? = null,
86-
verbose: Boolean = false
104+
verbose: Boolean = false,
105+
spanAdapter: SpanAdapter? = null,
87106
) {
88107
check(span.type == SpanType.CREATE_AGENT) {
89108
"${span.logString} Expected to end span type of type: <${SpanType.CREATE_AGENT}>, but received span of type: <${span.type}>"
@@ -92,5 +111,6 @@ internal fun endCreateAgentSpan(
92111
// error.type
93112
span.addCommonErrorAttributes(error)
94113

114+
spanAdapter?.onBeforeSpanFinished(span)
95115
span.end(error.toStatusData(), verbose)
96116
}

0 commit comments

Comments
 (0)