Skip to content

Commit

Permalink
fix(tracing): Use or_current() to prevent orphaned tracing spans
Browse files Browse the repository at this point in the history
  • Loading branch information
@tinco
tinco committed Jan 24, 2025
1 parent c5408a9 commit 0a8debb
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 19 deletions.
23 changes: 15 additions & 8 deletions swiftide-agents/src/agent.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -260,7 +260,7 @@ impl Agent {
"otel.name" = format!("hook.{}", HookTypes::AfterTool)
);
tracing::info!("Calling {} hook", HookTypes::AfterTool);
hook(&*self.context).instrument(span).await?;
hook(&*self.context).instrument(span.or_current()).await?;
}
}
}
Expand Down Expand Up @@ -316,7 +316,7 @@ impl Agent {
);
tracing::info!("Calling {} hook", HookTypes::BeforeCompletion);
hook(&*self.context, &mut chat_completion_request)
.instrument(span)
.instrument(span.or_current())
.await?;
}
}
Expand All @@ -341,7 +341,9 @@ impl Agent {
"otel.name" = format!("hook.{}", HookTypes::AfterCompletion)
);
tracing::info!("Calling {} hook", HookTypes::AfterCompletion);
hook(&*self.context, &mut response).instrument(span).await?;
hook(&*self.context, &mut response)
.instrument(span.or_current())
.await?;
}
}
self.add_message(ChatMessage::Assistant(
Expand All @@ -361,7 +363,7 @@ impl Agent {
"otel.name" = format!("hook.{}", HookTypes::AfterEach)
);
tracing::info!("Calling {} hook", HookTypes::AfterEach);
hook(&*self.context).instrument(span).await?;
hook(&*self.context).instrument(span.or_current()).await?;
}
}

Expand Down Expand Up @@ -389,7 +391,9 @@ impl Agent {
"otel.name" = format!("hook.{}", HookTypes::BeforeTool)
);
tracing::info!("Calling {} hook", HookTypes::BeforeTool);
hook(&*self.context, &tool_call).instrument(span).await?;
hook(&*self.context, &tool_call)
.instrument(span.or_current())
.await?;
}
}

Expand All @@ -403,7 +407,7 @@ impl Agent {
tracing::debug!(output = output.to_string(), args = ?tool_args, tool_name = tool.name(), "Completed tool call");

Ok(output)
}.instrument(tool_span));
}.instrument(tool_span.or_current()));

handles.push((handle, tool_call));
}
Expand All @@ -420,7 +424,7 @@ impl Agent {
);
tracing::info!("Calling {} hook", HookTypes::AfterTool);
hook(&*self.context, &tool_call, &mut output)
.instrument(span)
.instrument(span.or_current())
.await?;
}
}
Expand Down Expand Up @@ -464,7 +468,10 @@ impl Agent {
"hook",
"otel.name" = format!("hook.{}", HookTypes::OnNewMessage)
);
if let Err(err) = hook(&*self.context, &mut message).instrument(span).await {
if let Err(err) = hook(&*self.context, &mut message)
.instrument(span.or_current())
.await
{
tracing::error!(
"Error in {hooktype} hook: {err}",
hooktype = HookTypes::OnNewMessage,
Expand Down
14 changes: 7 additions & 7 deletions swiftide-indexing/src/pipeline.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ impl Pipeline {
Ok(Some(node))
}
}
.instrument(span)
.instrument(span.or_current())
})
.boxed()
.into();
Expand Down Expand Up @@ -194,7 +194,7 @@ impl Pipeline {
task::spawn(async move {
tracing::debug!(node = ?node, transformer = transformer.name(), "Transforming node");
transformer.transform_node(node).await
}.instrument(span)
}.instrument(span.or_current())
)
.err_into::<anyhow::Error>()
})
Expand Down Expand Up @@ -243,7 +243,7 @@ impl Pipeline {
);
transformer.batch_transform(nodes).await
}
.instrument(span),
.instrument(span.or_current()),
)
.map_err(anyhow::Error::from)
})
Expand Down Expand Up @@ -279,7 +279,7 @@ impl Pipeline {
tracing::debug!(chunker = chunker.name(), "Chunking node");
chunker.transform_node(node).await
}
.instrument(span),
.instrument(span.or_current()),
)
.map_err(anyhow::Error::from)
})
Expand Down Expand Up @@ -323,7 +323,7 @@ impl Pipeline {
tracing::debug!(storage = storage.name(), num_nodes = nodes.len(), "Batch Storing nodes");
storage.batch_store(nodes).await
}
.instrument(span)
.instrument(span.or_current())
)
.map_err(anyhow::Error::from)

Expand All @@ -346,7 +346,7 @@ impl Pipeline {

storage.store(node).await
}
.instrument(span),
.instrument(span.or_current()),
)
.err_into::<anyhow::Error>()
})
Expand Down Expand Up @@ -410,7 +410,7 @@ impl Pipeline {
})
.await;
}
.instrument(span),
.instrument(span.or_current()),
);

let left_pipeline = Self {
Expand Down
8 changes: 4 additions & 4 deletions swiftide-query/src/query/pipeline.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ where

Ok(transformed_query)
}
.instrument(span),
.instrument(span.or_current()),
)
.err_into::<anyhow::Error>()
})
Expand Down Expand Up @@ -181,7 +181,7 @@ impl<'stream: 'static, STRATEGY: SearchStrategy + 'stream>
Ok(result)
}
}
.instrument(span),
.instrument(span.or_current()),
)
.err_into::<anyhow::Error>()
})
Expand Down Expand Up @@ -229,7 +229,7 @@ impl<'stream: 'static, STRATEGY: SearchStrategy> Pipeline<'stream, STRATEGY, sta

Ok(transformed_query)
}
.instrument(span),
.instrument(span.or_current()),
)
.err_into::<anyhow::Error>()
})
Expand Down Expand Up @@ -281,7 +281,7 @@ impl<'stream: 'static, STRATEGY: SearchStrategy> Pipeline<'stream, STRATEGY, sta
Ok(result)
}
}
.instrument(span),
.instrument(span.or_current()),
)
.err_into::<anyhow::Error>()
})
Expand Down

0 comments on commit 0a8debb

Please sign in to comment.