diff --git a/crates/inference-core/src/session.rs b/crates/inference-core/src/session.rs index a0defd6..f6a9566 100644 --- a/crates/inference-core/src/session.rs +++ b/crates/inference-core/src/session.rs @@ -297,22 +297,30 @@ fn decode_token_piece_with(mut call: F) -> String where F: FnMut(*mut std::ffi::c_char, i32) -> i32, { - let mut capacity = 64usize; + let mut stack_buffer = [0u8; 256]; + let written = call( + stack_buffer.as_mut_ptr() as *mut std::ffi::c_char, + stack_buffer.len() as i32, + ); + + if written >= 0 { + return String::from_utf8_lossy(&stack_buffer[..written as usize]).into_owned(); + } + + let mut heap_buffer = vec![0u8; (-written) as usize]; loop { - let mut buffer = vec![0u8; capacity]; let written = call( - buffer.as_mut_ptr() as *mut std::ffi::c_char, - capacity as i32, + heap_buffer.as_mut_ptr() as *mut std::ffi::c_char, + heap_buffer.len() as i32, ); if written < 0 { - capacity = (-written) as usize; + heap_buffer.resize((-written) as usize, 0); continue; } - let written = written as usize; - return String::from_utf8_lossy(&buffer[..written]).into_owned(); + return String::from_utf8_lossy(&heap_buffer[..written as usize]).into_owned(); } } @@ -321,17 +329,14 @@ mod tests { use super::decode_token_piece_with; #[test] - fn retries_when_piece_buffer_is_too_small() { + fn uses_stack_buffer_when_piece_fits() { let mut calls = 0; let decoded = decode_token_piece_with(|buf, len| { calls += 1; - if len < 17 { - return -17; - } - let bytes = b" Quantum computer"; + assert!(len >= bytes.len() as i32); unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len()); } @@ -350,18 +355,21 @@ mod tests { calls += 1; if calls == 1 { - return -18; + return -300; } - assert!(len >= 18); - let bytes = b" Quantum computer"; + let bytes = b" Quantum computer rendered from a resized heap buffer"; + assert!(len >= bytes.len() as i32); unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len()); } bytes.len() as i32 }); - assert_eq!(decoded, " Quantum computer"); + assert_eq!( + decoded, + " Quantum computer rendered from a resized heap buffer" + ); assert_eq!(calls, 2); } }