1919import com .coze .loop .prompt .PromptProvider ;
2020import com .coze .loop .spec .tracespec .SpanKeys ;
2121import com .coze .loop .stream .StreamReader ;
22+ import com .coze .loop .trace .CozeLoopContext ;
2223import com .coze .loop .trace .CozeLoopSpan ;
2324import com .coze .loop .trace .CozeLoopTracerProvider ;
2425
@@ -84,10 +85,7 @@ public CozeLoopSpan startSpan(String name, String spanType, String scene) {
8485 public CozeLoopSpan startSpan (String name , String spanType , CozeLoopSpan parentSpan ) {
8586 checkNotClosed ();
8687
87- Context parentContext =
88- parentSpan != null && parentSpan .getSpan () != null
89- ? Context .current ().with (parentSpan .getSpan ())
90- : null ;
88+ CozeLoopContext parentContext = parentSpan != null ? parentSpan .getContext () : null ;
9189
9290 return startSpan (name , spanType , parentContext );
9391 }
@@ -97,23 +95,21 @@ public CozeLoopSpan startSpan(
9795 String name , String spanType , CozeLoopSpan parentSpan , String scene ) {
9896 checkNotClosed ();
9997
100- Context parentContext =
101- parentSpan != null && parentSpan .getSpan () != null
102- ? Context .current ().with (parentSpan .getSpan ())
103- : null ;
98+ CozeLoopContext parentContext = parentSpan != null ? parentSpan .getContext () : null ;
10499
105100 return startSpan (name , spanType , parentContext , scene );
106101 }
107102
108103 @ Override
109- public CozeLoopSpan startSpan (String name , String spanType , Context parentContext ) {
104+ public CozeLoopSpan startSpan (String name , String spanType , CozeLoopContext parentContext ) {
110105 checkNotClosed ();
111106
112107 return startSpan (name , spanType , parentContext , "" );
113108 }
114109
115110 @ Override
116- public CozeLoopSpan startSpan (String name , String spanType , Context parentContext , String scene ) {
111+ public CozeLoopSpan startSpan (
112+ String name , String spanType , CozeLoopContext parentContext , String scene ) {
117113 checkNotClosed ();
118114
119115 SpanBuilder spanBuilder = tracer .spanBuilder (name );
@@ -140,39 +136,35 @@ public CozeLoopSpan startSpan(String name, String spanType, Context parentContex
140136 Context finalContext = contextToUse .with (span );
141137 Scope scope = finalContext .makeCurrent ();
142138
143- return new CozeLoopSpan (span , scope , tracerProvider .getPropagators (), finalContext , scene );
139+ return new CozeLoopSpan (
140+ span , scope , tracerProvider .getPropagators (), new CozeLoopContext (finalContext ), scene );
144141 }
145142
146143 @ Override
147- public Context extractContext (Map <String , String > headers ) {
144+ public CozeLoopContext extractContext (Map <String , String > headers ) {
148145 checkNotClosed ();
149146 if (headers == null || headers .isEmpty ()) {
150- return Context .current ();
147+ return CozeLoopContext .current ();
151148 }
152149
153- return tracerProvider
154- .getPropagators ()
155- .getTextMapPropagator ()
156- .extract (
157- Context .current (),
158- headers ,
159- new TextMapGetter <Map <String , String >>() {
160- @ Override
161- public Iterable <String > keys (Map <String , String > carrier ) {
162- return carrier .keySet ();
163- }
164-
165- @ Override
166- public String get (Map <String , String > carrier , String key ) {
167- return carrier .get (key );
168- }
169- });
170- }
171-
172- @ Override
173- public Tracer getTracer () {
174- checkNotClosed ();
175- return tracer ;
150+ return new CozeLoopContext (
151+ tracerProvider
152+ .getPropagators ()
153+ .getTextMapPropagator ()
154+ .extract (
155+ Context .current (),
156+ headers ,
157+ new TextMapGetter <Map <String , String >>() {
158+ @ Override
159+ public Iterable <String > keys (Map <String , String > carrier ) {
160+ return carrier .keySet ();
161+ }
162+
163+ @ Override
164+ public String get (Map <String , String > carrier , String key ) {
165+ return carrier .get (key );
166+ }
167+ }));
176168 }
177169
178170 // ========== Prompt Operations ==========
0 commit comments