try to fix typescript tracing

engine/baml-runtime/src/tracing/mod.rs

@@ -1,12 +1,9 @@
 pub mod api_wrapper;
-#[cfg(not(target_arch = "wasm32"))]
-mod threaded_tracer;
-#[cfg(target_arch = "wasm32")]
-mod wasm_tracer;
 
 use crate::on_log_event::LogEventCallbackSync;
 use anyhow::Result;
 use baml_types::{BamlMap, BamlMediaType, BamlValue};
+use cfg_if::cfg_if;
 use colored::Colorize;
 use internal_baml_jinja::RenderedPrompt;
 use std::collections::HashMap;
@@ -27,16 +24,17 @@ use self::api_wrapper::{
     },
     APIWrapper,
 };
-#[cfg(not(target_arch = "wasm32"))]
-use self::threaded_tracer::ThreadedTracer;
 
-#[cfg(target_arch = "wasm32")]
-use self::wasm_tracer::NonThreadedTracer;
+cfg_if! {
+    if #[cfg(target_arch = "wasm32")] {
+        mod wasm_tracer;
+        use self::wasm_tracer::NonThreadedTracer as TracerImpl;
+    } else {
+        mod threaded_tracer;
+        use self::threaded_tracer::ThreadedTracer as TracerImpl;
+    }
+}
 
-#[cfg(not(target_arch = "wasm32"))]
-type TracerImpl = ThreadedTracer;
-#[cfg(target_arch = "wasm32")]
-type TracerImpl = NonThreadedTracer;
 #[derive(Debug)]
 pub struct TracingSpan {
     span_id: Uuid,
@@ -95,6 +93,7 @@ impl BamlTracer {
         params: &BamlMap<String, BamlValue>,
     ) -> (Option<TracingSpan>, RuntimeContext) {
         let span_id = ctx.enter(function_name);
+        log::trace!("Entering span: {:#?}:::{:?}", span_id, function_name);
         if !self.enabled {
             return (None, ctx.create_ctx(tb));
         }
@@ -150,6 +149,12 @@ impl BamlTracer {
                 ctx
             );
         };
+        log::trace!(
+            "Finishing span: {:#?} {}\nevent chain {:?}",
+            span,
+            span_id,
+            event_chain
+        );
 
         if span.span_id != span_id {
             anyhow::bail!("Span ID mismatch: {} != {}", span.span_id, span_id);
@@ -211,6 +216,13 @@ impl BamlTracer {
             anyhow::bail!("Attempting to finish a span without first starting one");
         };
 
+        log::trace!(
+            "Finishing baml span: {:#?} {}\nevent chain {:?}",
+            span,
+            span_id,
+            event_chain
+        );
+
         if span.span_id != span_id {
             anyhow::bail!("Span ID mismatch: {} != {}", span.span_id, span_id);
         }

engine/language-client-codegen/src/typescript/generate_types.rs

@@ -7,14 +7,14 @@ use crate::GeneratorArgs;
 use super::ToTypeReferenceInClientDefinition;
 
 #[derive(askama::Template)]
-#[template(path = "type_builder.js.j2", escape = "none")]
+#[template(path = "type_builder.ts.j2", escape = "none")]
 pub(crate) struct TypeBuilder<'ir> {
     enums: Vec<TypescriptEnum<'ir>>,
     classes: Vec<TypescriptClass<'ir>>,
 }
 
 #[derive(askama::Template)]
-#[template(path = "types.js.j2", escape = "none")]
+#[template(path = "types.ts.j2", escape = "none")]
 pub(crate) struct TypescriptTypes<'ir> {
     enums: Vec<TypescriptEnum<'ir>>,
     classes: Vec<TypescriptClass<'ir>>,

engine/language-client-codegen/src/typescript/mod.rs

@@ -12,7 +12,7 @@ use self::typescript_language_features::{ToTypescript, TypescriptLanguageFeature
 use crate::dir_writer::FileCollector;
 
 #[derive(askama::Template)]
-#[template(path = "client.js.j2", escape = "none")]
+#[template(path = "client.ts.j2", escape = "none")]
 struct TypescriptClient {
     funcs: Vec<TypescriptFunction>,
     types: Vec<String>,

engine/language_client_typescript/async_context_vars.d.ts

@@ -1,12 +1,12 @@
 import { BamlSpan, RuntimeContextManager, BamlRuntime, BamlLogEvent } from './native';
-export declare class CtxManager {
+export declare class BamlCtxManager {
     private rt;
     private ctx;
     constructor(rt: BamlRuntime);
     upsertTags(tags: Record<string, string>): void;
     get(): RuntimeContextManager;
-    startTraceSync(name: string, args: Record<string, any>): BamlSpan;
-    startTraceAsync(name: string, args: Record<string, any>): BamlSpan;
+    startTraceSync(name: string, args: Record<string, any>): [RuntimeContextManager, BamlSpan];
+    startTraceAsync(name: string, args: Record<string, any>): [RuntimeContextManager, BamlSpan];
     endTrace(span: BamlSpan, response: any): void;
     flush(): void;
     onLogEvent(callback: (event: BamlLogEvent) => void): void;

engine/language_client_typescript/async_context_vars.js

@@ -1,9 +1,9 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.CtxManager = void 0;
+exports.BamlCtxManager = void 0;
 const native_1 = require("./native");
 const async_hooks_1 = require("async_hooks");
-class CtxManager {
+class BamlCtxManager {
     rt;
     ctx;
     constructor(rt) {
@@ -28,15 +28,11 @@ class CtxManager {
     }
     startTraceSync(name, args) {
         const mng = this.get();
-        // const clone = mng.deepClone()
-        // this.ctx.enterWith(clone)
-        return native_1.BamlSpan.new(this.rt, name, args, mng);
+        return [mng, native_1.BamlSpan.new(this.rt, name, args, mng)];
     }
     startTraceAsync(name, args) {
-        const mng = this.get();
-        const clone = mng.deepClone();
-        this.ctx.enterWith(clone);
-        return native_1.BamlSpan.new(this.rt, name, args, clone);
+        const mng = this.get().deepClone();
+        return [mng, native_1.BamlSpan.new(this.rt, name, args, mng)];
     }
     endTrace(span, response) {
         const manager = this.ctx.getStore();
@@ -62,16 +58,18 @@ class CtxManager {
                 ...acc,
                 [`arg${i}`]: arg, // generic way to label args
             }), {});
-            const span = this.startTraceSync(name, params);
-            try {
-                const response = func(...args);
-                this.endTrace(span, response);
-                return response;
-            }
-            catch (e) {
-                this.endTrace(span, e);
-                throw e;
-            }
+            const [mng, span] = this.startTraceSync(name, params);
+            this.ctx.run(mng, () => {
+                try {
+                    const response = func(...args);
+                    this.endTrace(span, response);
+                    return response;
+                }
+                catch (e) {
+                    this.endTrace(span, e);
+                    throw e;
+                }
+            });
         });
     }
     traceFnAync(name, func) {
@@ -81,17 +79,19 @@ class CtxManager {
                 ...acc,
                 [`arg${i}`]: arg, // generic way to label args
             }), {});
-            const span = this.startTraceAsync(funcName, params);
-            try {
-                const response = await func(...args);
-                this.endTrace(span, response);
-                return response;
-            }
-            catch (e) {
-                this.endTrace(span, e);
-                throw e;
-            }
+            const [mng, span] = this.startTraceAsync(name, params);
+            await this.ctx.run(mng, async () => {
+                try {
+                    const response = await func(...args);
+                    this.endTrace(span, response);
+                    return response;
+                }
+                catch (e) {
+                    this.endTrace(span, e);
+                    throw e;
+                }
+            });
         });
     }
 }
-exports.CtxManager = CtxManager;
+exports.BamlCtxManager = BamlCtxManager;

engine/language_client_typescript/index.d.ts

@@ -1,4 +1,4 @@
-export { BamlRuntime, FunctionResult, FunctionResultStream, BamlImage as Image, BamlAudio as Audio, invoke_runtime_cli } from './native';
+export { BamlRuntime, FunctionResult, FunctionResultStream, BamlImage as Image, BamlAudio as Audio, invoke_runtime_cli, } from './native';
 export { BamlStream } from './stream';
-export { CtxManager as BamlCtxManager } from './async_context_vars';
+export { BamlCtxManager } from './async_context_vars';
 //# sourceMappingURL=index.d.ts.map

engine/language_client_typescript/index.js

@@ -11,4 +11,4 @@ Object.defineProperty(exports, "invoke_runtime_cli", { enumerable: true, get: fu
 var stream_1 = require("./stream");
 Object.defineProperty(exports, "BamlStream", { enumerable: true, get: function () { return stream_1.BamlStream; } });
 var async_context_vars_1 = require("./async_context_vars");
-Object.defineProperty(exports, "BamlCtxManager", { enumerable: true, get: function () { return async_context_vars_1.CtxManager; } });
+Object.defineProperty(exports, "BamlCtxManager", { enumerable: true, get: function () { return async_context_vars_1.BamlCtxManager; } });