diff --git a/‎engine/baml-runtime/src/cli/serve/mod.rs‎
Lines changed: 1 addition & 0 deletions b/‎engine/baml-runtime/src/cli/serve/mod.rs‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎engine/baml-runtime/src/internal/llm_client/orchestrator/stream.rs‎
Lines changed: 6 additions & 1 deletion b/‎engine/baml-runtime/src/internal/llm_client/orchestrator/stream.rs‎
Lines changed: 6 additions & 1 deletion
diff --git a/‎engine/baml-runtime/src/lib.rs‎
Lines changed: 8 additions & 1 deletion b/‎engine/baml-runtime/src/lib.rs‎
Lines changed: 8 additions & 1 deletion
diff --git a/‎engine/baml-runtime/src/tracingv2/storage/storage.rs‎
Lines changed: 2 additions & 2 deletions b/‎engine/baml-runtime/src/tracingv2/storage/storage.rs‎
Lines changed: 2 additions & 2 deletions
diff --git a/‎engine/baml-runtime/src/types/stream.rs‎
Lines changed: 8 additions & 3 deletions b/‎engine/baml-runtime/src/types/stream.rs‎
Lines changed: 8 additions & 3 deletions
diff --git a/‎engine/generators/data/sample/baml_src/main.baml‎
Lines changed: 15 additions & 1 deletion b/‎engine/generators/data/sample/baml_src/main.baml‎
Lines changed: 15 additions & 1 deletion
diff --git a/‎engine/generators/data/sample/go/go.mod‎
Lines changed: 9 additions & 2 deletions b/‎engine/generators/data/sample/go/go.mod‎
Lines changed: 9 additions & 2 deletions
diff --git a/‎engine/generators/data/sample/go/go.sum‎
Lines changed: 6 additions & 0 deletions b/‎engine/generators/data/sample/go/go.sum‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎engine/generators/data/sample/go/main_test.go‎
Lines changed: 98 additions & 0 deletions b/‎engine/generators/data/sample/go/main_test.go‎
Lines changed: 98 additions & 0 deletions
@@ -448,6 +448,7 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
                 Ok(mut result_stream) => {
                     let (result, _trace_id) = result_stream
                         .run(
+                            None::<fn()>,
                             Some(move |result| {
                                 // If the receiver is closed (either because it called close or it was dropped),
                                 // we can't really do anything
 
@@ -24,12 +24,13 @@ use crate::{
     FunctionResult, RuntimeContext,
 };
 
-pub async fn orchestrate_stream<F>(
+pub async fn orchestrate_stream<F, G>(
     iter: OrchestratorNodeIterator,
     ir: &IntermediateRepr,
     ctx: &RuntimeContext,
     prompt: &PromptRenderer,
     params: &BamlValue,
+    on_tick_fn: Option<G>,
     partial_parse_fn: impl Fn(&str) -> Result<ResponseBamlValue>,
     parse_fn: impl Fn(&str) -> Result<ResponseBamlValue>,
     on_event: Option<F>,
@@ -43,6 +44,7 @@ pub async fn orchestrate_stream<F>(
 )
 where
     F: Fn(FunctionResult),
+    G: Fn(),
 {
     let mut results = Vec::new();
     let mut total_sleep_duration = std::time::Duration::from_secs(0);
@@ -67,6 +69,9 @@ where
         let final_response = match stream_res {
             Ok(response) => response
                 .map(|stream_part| {
+                    if let Some(on_tick) = on_tick_fn.as_ref() {
+                        on_tick();
+                    }
                     if let Some(on_event) = on_event.as_ref() {
                         if let LLMResponse::Success(s) = &stream_part {
                             let response_value = partial_parse_fn(&s.content);
 
@@ -514,7 +514,14 @@ impl BamlRuntime {
                 vec![],
             )?;
             let (response_res, call_uuid) = stream
-                .run(on_event, ctx, type_builder.as_ref(), None, env_vars.clone())
+                .run(
+                    None::<fn()>,
+                    on_event,
+                    ctx,
+                    type_builder.as_ref(),
+                    None,
+                    env_vars.clone(),
+                )
                 .await;
             let res = response_res?;
             let (_, llm_resp, val) = res
 
@@ -248,7 +248,6 @@ fn build_function_log(
                 let entry = calls_map.entry(rid.clone()).or_default();
 
                 // find or insert the event
-
                 match &mut entry.http_response_stream {
                     Some(stream) => {
                         stream.lock().unwrap().push(http_res_stream.clone());
@@ -370,7 +369,8 @@ fn build_function_log(
     };
 
     let new_arc = Arc::new(Mutex::new(function_log_inner));
-    {
+    // only cache if we we've finished the function
+    if function_end_time.is_some() {
         // Insert into the cache
         let mut lock = storage.function_inners.lock().unwrap();
         lock.insert(function_id.clone(), new_arc.clone());
 
@@ -63,8 +63,9 @@ first.scope.clone();
 
 impl FunctionResultStream {
     #[cfg(not(target_arch = "wasm32"))]
-    pub fn run_sync<F>(
+    pub fn run_sync<F, G>(
         &mut self,
+        on_tick: Option<G>,
         on_event: Option<F>,
         ctx: &RuntimeContextManager,
         tb: Option<&TypeBuilder>,
@@ -73,14 +74,16 @@ impl FunctionResultStream {
     ) -> (Result<FunctionResult>, baml_ids::FunctionCallId)
     where
         F: Fn(FunctionResult),
+        G: Fn(),
     {
         let rt = self.tokio_runtime.clone();
-        let fut = self.run(on_event, ctx, tb, cb, env_vars);
+        let fut = self.run(on_tick, on_event, ctx, tb, cb, env_vars);
         rt.block_on(fut)
     }
 
-    pub async fn run<F>(
+    pub async fn run<F, G>(
         &mut self,
+        on_tick: Option<G>,
         on_event: Option<F>,
         ctx: &RuntimeContextManager,
         tb: Option<&TypeBuilder>,
@@ -89,6 +92,7 @@ impl FunctionResultStream {
     ) -> (Result<FunctionResult>, baml_ids::FunctionCallId)
     where
         F: Fn(FunctionResult),
+        G: Fn(),
     {
         let mut local_orchestrator = Vec::new();
         std::mem::swap(&mut local_orchestrator, &mut self.orchestrator);
@@ -114,6 +118,7 @@ impl FunctionResultStream {
                         &rctx,
                         &self.renderer,
                         &baml_types::BamlValue::Map(self.prepared_func.value.clone()),
+                        on_tick,
                         |content| self.renderer.parse(self.ir.as_ref(), &rctx, content, true),
                         |content| self.renderer.parse(self.ir.as_ref(), &rctx, content, false),
                         on_event,
 
@@ -1,4 +1,5 @@
 class Example {
+    type "example_1" @stream.not_null
     a int
     b string
 
@@ -7,13 +8,26 @@ class Example {
 }
 
 class Example2 {
+    type "example_2" @stream.not_null
     item Example
     element string
     element2 string
 }
 
+client<llm> Thinking {
+    provider "anthropic"
+    options {
+        model "claude-sonnet-4-20250514"
+        thinking {
+            type enabled
+            budget_tokens 1024
+        }
+        max_tokens 2048
+    }
+}
+
 function Foo(x: int) -> Example2 | Example {
-    client "openai/gpt-4o-mini"
+    client Thinking
     prompt #"
         Fill out this data model with some examples.
 
 
@@ -2,8 +2,15 @@ module sample
 
 go 1.24.0
 
-require github.com/boundaryml/baml v0.90.0
+require (
+	github.com/boundaryml/baml v0.90.0
+	github.com/tidwall/gjson v1.18.0
+)
 
-require google.golang.org/protobuf v1.36.6 // indirect
+require (
+	github.com/tidwall/match v1.1.1 // indirect
+	github.com/tidwall/pretty v1.2.0 // indirect
+	google.golang.org/protobuf v1.36.6 // indirect
+)
 
 replace github.com/boundaryml/baml => ../../../../../../
@@ -6,6 +6,12 @@ github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY=
+github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
+github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
+github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
+github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
 
@@ -10,8 +10,106 @@ import (
 	"testing"
 
 	baml "github.com/boundaryml/baml/engine/language_client_go/pkg"
+	"github.com/tidwall/gjson"
 )
 
+func getOnTick() (baml.TickCallback, *string, *int) {
+	var lastThinking string
+	var tickCount int
+	onTick := func(ctx context.Context, reason baml.TickReason, log baml.FunctionLog) baml.FunctionSignal {
+		tickCount++
+		calls, err := log.Calls()
+		if err != nil {
+			fmt.Println("Error getting selected call: ", err)
+			return nil
+		}
+
+		if len(calls) == 0 {
+			fmt.Println("No calls found")
+			return nil
+		}
+
+		selectedCall := calls[len(calls)-1]
+
+		if as_stream, ok := selectedCall.(baml.LLMStreamCall); ok {
+			stream, err := as_stream.SSEChunks()
+			if err != nil {
+				fmt.Println("Error getting stream: ", err)
+				return nil
+			}
+			fmt.Println("Tick: ", len(stream), " chunks")
+			accumulated := ""
+			for _, chunk := range stream {
+				text, err := chunk.Text()
+				if err != nil {
+					// this is an openai stream
+				} else {
+					if content := gjson.Get(text, "delta.thinking"); content.Exists() {
+						accumulated += content.String()
+					}
+				}
+			}
+			fmt.Println("Accumulated: ", accumulated)
+			lastThinking = accumulated
+			fmt.Println("--------------------------------")
+		} else {
+			fmt.Println("Response is not a stream: ", selectedCall)
+		}
+
+		return nil
+	}
+
+	return onTick, &lastThinking, &tickCount
+}
+
+func TestOnTickRequest(t *testing.T) {
+	ctx := context.Background()
+
+	onTick, lastThinking, tickCount := getOnTick()
+
+	_, err := b.Foo(ctx, 8192, b.WithExperimentalOnTick(onTick))
+	if err != nil {
+		t.Fatalf("Error in Foo: %v", err)
+	}
+
+	if *tickCount < 10 {
+		t.Errorf("Expected more than 10 ticks, got %d", *tickCount)
+	}
+
+	if *lastThinking == "" {
+		t.Errorf("Expected thinking, got %s", *lastThinking)
+	}
+}
+
+func TestOnTickStream(t *testing.T) {
+	ctx := context.Background()
+
+	onTick, lastThinking, tickCount := getOnTick()
+
+	result, err := b.Stream.Foo(ctx, 8192, b.WithExperimentalOnTick(onTick))
+	if err != nil {
+		t.Fatalf("Error in Foo: %v", err)
+	}
+	for result := range result {
+		if result.IsError {
+			t.Fatalf("Error in Foo: %v", result.Error)
+		} else if result.IsFinal {
+			final := result.Final()
+			fmt.Println("final", final)
+		} else {
+			fmt.Printf("Stream: %+v\n", result.Stream())
+		}
+	}
+
+	if *tickCount < 10 {
+		t.Errorf("Expected more than 10 ticks, got %d", *tickCount)
+	}
+
+	if *lastThinking == "" {
+		t.Errorf("Expected thinking, got %s", *lastThinking)
+	}
+}
+
 func TestFoo(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()