support set span finish time

Author: JqRrt

CHANGLOG.md

@@ -2,6 +2,7 @@
 ### Added
 - add to_runnable decorator
 - lcc tool message add name and tool_call_id
+- support set finish time of span
 
 ## [0.1.24] - 2026-01-16
 ### Added

cozeloop/decorator/decorator.py

@@ -313,7 +313,6 @@ async def async_stream_wrapper(*args: Any, **kwargs: Any):
         else:
             return decorator(func)
 
-
     def to_runnable(
             self,
             func: Callable = None,
@@ -341,11 +340,15 @@ def sync_wrapper(*args: Any, **kwargs: Any):
                 config = _convert_config(config)
                 res = None
                 try:
+                    extra = {}
+                    if len(args) > 0 and is_class_func(func):
+                        extra = {"_inner_class_self": args[0]}
+                        args = args[1:]
                     inp = {
                         "args": args,
                         "kwargs": kwargs
                     }
-                    res = RunnableLambda(_param_wrapped_func).invoke(input=inp, config=config)
+                    res = RunnableLambda(_param_wrapped_func).invoke(input=inp, config=config, **extra)
                     if hasattr(res, "__iter__"):
                         return res
                 except StopIteration:
@@ -362,11 +365,15 @@ async def async_wrapper(*args: Any, **kwargs: Any):
                 config = _convert_config(config)
                 res = None
                 try:
+                    extra = {}
+                    if len(args) > 0 and is_class_func(func):
+                        extra = {"_inner_class_self": args[0]}
+                        args = args[1:]
                     inp = {
                         "args": args,
                         "kwargs": kwargs
                     }
-                    res = await RunnableLambda(_param_wrapped_func_async).ainvoke(input=inp, config=config)
+                    res = await RunnableLambda(_param_wrapped_func_async).ainvoke(input=inp, config=config, **extra)
                     if hasattr(res, "__aiter__"):
                         return res
                 except StopIteration:
@@ -387,11 +394,15 @@ def gen_wrapper(*args: Any, **kwargs: Any):
                 config = kwargs.pop("config", None)
                 config = _convert_config(config)
                 try:
+                    extra = {}
+                    if len(args) > 0 and is_class_func(func):
+                        extra = {"_inner_class_self": args[0]}
+                        args = args[1:]
                     inp = {
                         "args": args,
                         "kwargs": kwargs
                     }
-                    gen = RunnableLambda(_param_wrapped_func).invoke(input=inp, config=config)
+                    gen = RunnableLambda(_param_wrapped_func).invoke(input=inp, config=config, *extra)
                     try:
                         for item in gen:
                             yield item
@@ -405,11 +416,15 @@ async def async_gen_wrapper(*args: Any, **kwargs: Any):
                 config = kwargs.pop("config", None)
                 config = _convert_config(config)
                 try:
+                    extra = {}
+                    if len(args) > 0 and is_class_func(func):
+                        extra = {"_inner_class_self": args[0]}
+                        args = args[1:]
                     inp = {
                         "args": args,
                         "kwargs": kwargs
                     }
-                    gen = RunnableLambda(_param_wrapped_func_async).invoke(input=inp, config=config)
+                    gen = RunnableLambda(_param_wrapped_func_async).invoke(input=inp, config=config, **extra)
                     items = []
                     try:
                         async for item in gen:
@@ -428,15 +443,25 @@ async def async_gen_wrapper(*args: Any, **kwargs: Any):
                         raise e
 
             # for convert parameter
-            def _param_wrapped_func(input_dict: dict) -> Any:
-                args = input_dict.get("args", ())
-                kwargs = input_dict.get("kwargs", {})
-                return func(*args, **kwargs)
-
-            async def _param_wrapped_func_async(input_dict: dict) -> Any:
-                args = input_dict.get("args", ())
-                kwargs = input_dict.get("kwargs", {})
-                return await func(*args, **kwargs)
+            def _param_wrapped_func(input_dict: dict, **kwargs) -> Any:
+                real_args = input_dict.get("args", ())
+                real_kwargs = input_dict.get("kwargs", {})
+
+                inner_class_self = kwargs.get("_inner_class_self", None)
+                if inner_class_self is not None:
+                    real_args = (inner_class_self, *real_args)
+
+                return func(*real_args, **real_kwargs)
+
+            async def _param_wrapped_func_async(input_dict: dict, **kwargs) -> Any:
+                real_args = input_dict.get("args", ())
+                real_kwargs = input_dict.get("kwargs", {})
+
+                inner_class_self = kwargs.get("_inner_class_self", None)
+                if inner_class_self is not None:
+                    real_args = (inner_class_self, *real_args)
+
+                return await func(*real_args, **real_kwargs)
 
             def _convert_config(config: RunnableConfig = None) -> RunnableConfig | None:
                 if config is None:

cozeloop/internal/trace/noop_span.py

@@ -125,6 +125,9 @@ def set_system_tags(self, system_tags: Dict[str, Any]) -> None:
     def set_deployment_env(self, deployment_env: str) -> None:
         pass
 
+    def set_finish_time(self, finish_time: datetime) -> None:
+        pass
+
     def __enter__(self):
         return self
 

cozeloop/internal/trace/span.py

@@ -78,6 +78,7 @@ def __init__(self, span_type: str = '', name: str = '', space_id: str = '', trac
         self.space_id = space_id
         self.parent_span_id = parent_span_id
         self.start_time = start_time if start_time else datetime.now()
+        self.finish_time: datetime = None
         self.duration = duration
         self.tag_map = tag_map if tag_map else {}
         self.system_tag_map = system_tag_map if system_tag_map else {}
@@ -396,6 +397,10 @@ def set_system_tags(self, system_tags: Dict[str, Any]) -> None:
     def set_deployment_env(self, deployment_env: str) -> None:
         self.set_tags({DEPLOYMENT_ENV: deployment_env})
 
+    def set_finish_time(self, finish_time: datetime) -> None:
+        self.finish_time = finish_time
+
+
     def get_rectified_map(self, input_map: Dict[str, Any]) -> (Dict[str, Any], List[str], int):
         validate_map = {}
         cut_off_keys = []
@@ -541,7 +546,10 @@ def set_stat_info(self):
         if input_tokens > 0 or output_tokens > 0:
             self.set_tags({TOKENS: int(input_tokens) + int(output_tokens)})
 
-        duration = int((datetime.now().timestamp() - self.start_time.timestamp()) * 1000000)
+        finish_time_stamp = datetime.now().timestamp()
+        if self.finish_time is not None:
+            finish_time_stamp = self.finish_time.timestamp()
+        duration = int((finish_time_stamp - self.start_time.timestamp()) * 1000000)
         with self.lock:
             self.duration = duration
 

cozeloop/span.py

@@ -186,11 +186,18 @@ def set_system_tags(self, system_tags: Dict[str, Any]) -> None:
         Set system tags. DO NOT use this method unless you know what you are doing.
         """
 
+    @abstractmethod
     def set_deployment_env(self, deployment_env: str) -> None:
         """
         Set the deployment environment of the span, identify custom environments.
         """
 
+    @abstractmethod
+    def set_finish_time(self, finish_time: datetime) -> None:
+        """
+        Set the finish time of the span.
+        """
+
 
 class Span(CommonSpanSetter, SpanContext):
     """

examples/trace/simple.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import time
+from datetime import datetime, timedelta
 
 import cozeloop
 
@@ -100,6 +101,7 @@ def do_simple_demo():
         span.set_error(str(e))
 
     # 3. span finish
+    span.set_finish_time(datetime.now() + timedelta(seconds=3)) # set finish time as your need to change duration
     span.finish()
 
     # 4. (optional) flush or close