From 0a8debb3676410f44a1c2fa42b8eaa0f6f90b217 Mon Sep 17 00:00:00 2001
From: Tinco Andringa <mail@tinco.nl>
Date: Fri, 24 Jan 2025 17:22:19 +0100
Subject: [PATCH] fix(tracing): Use `or_current()` to prevent orphaned tracing
 spans

---
 swiftide-agents/src/agent.rs         | 23 +++++++++++++++--------
 swiftide-indexing/src/pipeline.rs    | 14 +++++++-------
 swiftide-query/src/query/pipeline.rs |  8 ++++----
 3 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/swiftide-agents/src/agent.rs b/swiftide-agents/src/agent.rs
index d5d08b6e..6063bfd7 100644
--- a/swiftide-agents/src/agent.rs
+++ b/swiftide-agents/src/agent.rs
@@ -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?;
                 }
             }
         }
@@ -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?;
             }
         }
@@ -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(
@@ -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?;
             }
         }
 
@@ -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?;
                 }
             }
 
@@ -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));
         }
@@ -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?;
                 }
             }
@@ -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,
diff --git a/swiftide-indexing/src/pipeline.rs b/swiftide-indexing/src/pipeline.rs
index 46c8c0ab..9374536d 100644
--- a/swiftide-indexing/src/pipeline.rs
+++ b/swiftide-indexing/src/pipeline.rs
@@ -157,7 +157,7 @@ impl Pipeline {
                         Ok(Some(node))
                     }
                 }
-                .instrument(span)
+                .instrument(span.or_current())
             })
             .boxed()
             .into();
@@ -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>()
             })
@@ -243,7 +243,7 @@ impl Pipeline {
                         );
                         transformer.batch_transform(nodes).await
                     }
-                    .instrument(span),
+                    .instrument(span.or_current()),
                 )
                 .map_err(anyhow::Error::from)
             })
@@ -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)
             })
@@ -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)
 
@@ -346,7 +346,7 @@ impl Pipeline {
 
                             storage.store(node).await
                         }
-                        .instrument(span),
+                        .instrument(span.or_current()),
                     )
                     .err_into::<anyhow::Error>()
                 })
@@ -410,7 +410,7 @@ impl Pipeline {
                     })
                     .await;
             }
-            .instrument(span),
+            .instrument(span.or_current()),
         );
 
         let left_pipeline = Self {
diff --git a/swiftide-query/src/query/pipeline.rs b/swiftide-query/src/query/pipeline.rs
index 250bc58d..bf4571e7 100644
--- a/swiftide-query/src/query/pipeline.rs
+++ b/swiftide-query/src/query/pipeline.rs
@@ -123,7 +123,7 @@ where
 
                         Ok(transformed_query)
                     }
-                    .instrument(span),
+                    .instrument(span.or_current()),
                 )
                 .err_into::<anyhow::Error>()
             })
@@ -181,7 +181,7 @@ impl<'stream: 'static, STRATEGY: SearchStrategy + 'stream>
                             Ok(result)
                         }
                     }
-                    .instrument(span),
+                    .instrument(span.or_current()),
                 )
                 .err_into::<anyhow::Error>()
             })
@@ -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>()
             })
@@ -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>()
             })