legion-ref/mesh-llm
1
0
Fork 0
mirror of https://github.com/Mesh-LLM/mesh-llm.git synced 2026-08-08 22:23:19 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(skippy): preserve native tool-call semantics (#1144)

Browse source 
* accept role-only assistant messages
* preserve sampled grammar history
* ci: use Q8 model for skippy correctness
This commit is contained in:
James Dumay 2026-08-04 18:28:34 +10:00 • committed by GitHub
parent 735dc004ed
commit aba53315ab
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 3830 additions and 2760 deletions
Download patch file Download diff file Expand all files Collapse all files

3
.agents/skills/manage-ci/references/current-inventory.md
View file

@ -425,6 +425,9 @@ bulk Cargo target caches retain cross-run reuse. Other producer and grouped-test
jobs remain remote-enabled. An explicitly authorized Depot call selects
`disk,webdav` before that GHA opt-out. Swift restores a
mode-independent Rust dependency cache that only trusted main pushes save.
The main and PR `rust_crate_tests` shard containing `skippy-runtime` downloads
the public Qwen3 correctness fixture and exposes `SKIPPY_CORRECTNESS_MODEL` to
the crate tests; this fixture does not require `HF_TOKEN`.
Persistent Cargo target and ABI reuse remains owned by
`Swatinem/rust-cache` and `actions/cache`. Current PR jobs use the normal
`mesh-llm` key namespace; native `actions/cache` writes remain merge-ref scoped,

13
.github/workflows/ci.yml vendored
View file

@ -451,8 +451,21 @@ jobs:
path: ${{ runner.temp }}/static-abi-input
- name: Restore immutable static ABI input
run: scripts/restore-static-abi-input.sh "$RUNNER_TEMP/static-abi-input" "$LLAMA_STAGE_BUILD_DIR" x86_64-unknown-linux-gnu cpu
- name: Download Skippy correctness model
if: ${{ contains(matrix.batch.crates, 'skippy-runtime') }}
env:
SKIPPY_CORRECTNESS_MODEL: ${{ runner.temp }}/skippy-correctness-model/Qwen3-0.6B-Q8_0.gguf
run: |
set -euo pipefail
mkdir -p "$RUNNER_TEMP/skippy-correctness-model"
hf download \
Qwen/Qwen3-0.6B-GGUF \
Qwen3-0.6B-Q8_0.gguf \
--local-dir "$RUNNER_TEMP/skippy-correctness-model"
test -s "$SKIPPY_CORRECTNESS_MODEL"
- name: Run crate tests
env:
SKIPPY_CORRECTNESS_MODEL: ${{ runner.temp }}/skippy-correctness-model/Qwen3-0.6B-Q8_0.gguf
TEST_CRATES: ${{ toJson(matrix.batch.crates) }}
run: |
mapfile -t crates < <(jq -r '.[]' <<<"$TEST_CRATES")

13
.github/workflows/pr_builds.yml vendored
View file

@ -767,8 +767,21 @@ jobs:
path: ${{ runner.temp }}/static-abi-input
- name: Restore immutable static ABI input
run: scripts/restore-static-abi-input.sh "$RUNNER_TEMP/static-abi-input" "$LLAMA_STAGE_BUILD_DIR" x86_64-unknown-linux-gnu cpu
- name: Download Skippy correctness model
if: ${{ contains(matrix.batch.crates, 'skippy-runtime') }}
env:
SKIPPY_CORRECTNESS_MODEL: ${{ runner.temp }}/skippy-correctness-model/Qwen3-0.6B-Q8_0.gguf
run: |
set -euo pipefail
mkdir -p "$RUNNER_TEMP/skippy-correctness-model"
hf download \
Qwen/Qwen3-0.6B-GGUF \
Qwen3-0.6B-Q8_0.gguf \
--local-dir "$RUNNER_TEMP/skippy-correctness-model"
test -s "$SKIPPY_CORRECTNESS_MODEL"
- name: Run crate tests
env:
SKIPPY_CORRECTNESS_MODEL: ${{ runner.temp }}/skippy-correctness-model/Qwen3-0.6B-Q8_0.gguf
TEST_CRATES: ${{ toJson(matrix.batch.crates) }}
run: |
mapfile -t crates < <(jq -r '.[]' <<<"$TEST_CRATES")

5
ci/ci.md
View file

@ -267,7 +267,10 @@ flowchart TD
that downstream smoke jobs consume before long validation groups finish.
Every affected Rust workspace crate is assigned to a generated
`rust_crate_tests` matrix and runs its complete `cargo test -p <crate>` suite;
protocol compatibility and Skippy smoke remain separate integration rows.
the shard containing `skippy-runtime` downloads the public Qwen3 correctness
fixture and sets `SKIPPY_CORRECTNESS_MODEL`, so the model-backed grammar
equivalence test runs instead of being skipped; protocol compatibility and
Skippy smoke remain separate integration rows.
Linux host/CPU-runtime and macOS host/Metal-runtime producers run
independently, and their product composers never compile. Linux backend rows
are split into one independent CUDA, ROCm, or Vulkan runtime producer plus

30
crates/openai-frontend/src/chat.rs
View file

@ -1,6 +1,6 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde::{Deserialize, Deserializer, Serialize, de::Error as _};
use serde_json::Value;
use crate::{
@ -118,11 +118,28 @@ fn invalid_tools_value(value: &Value) -> bool {
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct ChatMessage {
pub role: String,
#[serde(
default,
deserialize_with = "deserialize_present_message_content",
skip_serializing_if = "Option::is_none"
)]
pub content: Option<MessageContent>,
#[serde(flatten)]
pub extra: BTreeMap<String, Value>,
}
fn deserialize_present_message_content<'de, D>(
deserializer: D,
) -> Result<Option<MessageContent>, D::Error>
where
D: Deserializer<'de>,
{
let value = Value::deserialize(deserializer)?;
serde_json::from_value(value)
.map(Some)
.map_err(D::Error::custom)
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(untagged)]
pub enum MessageContent {
@ -387,4 +404,15 @@ mod tests {
assert_eq!(value, json!({ "reasoning_content": "Still thinking." }));
}
#[test]
fn chat_message_round_trip_preserves_missing_and_null_content() {
for expected in [
json!({"role": "assistant", "tool_calls": []}),
json!({"role": "assistant", "content": null, "tool_calls": []}),
] {
let message: ChatMessage = serde_json::from_value(expected.clone()).unwrap();
assert_eq!(serde_json::to_value(message).unwrap(), expected);
}
}
}

13
crates/skippy-runtime/src/activation.rs
View file

@ -625,6 +625,19 @@ impl StageSession {
))
}
pub(crate) fn verify_tokens_sampled_without_mtp(
&mut self,
token_ids: &[i32],
sampling: Option<&SamplingConfig>,
) -> Result<Vec<i32>> {
if token_ids.is_empty() {
return Ok(Vec::new());
}
Ok(self
.verify_tokens_frame_raw(token_ids, sampling, None, 0, 0)?
.predicted_tokens)
}
fn verify_tokens_frame_raw(
&mut self,
token_ids: &[i32],

8
crates/skippy-runtime/src/session.rs
View file

@ -361,6 +361,14 @@ impl StageSession {
Ok(predicted)
}
pub fn verify_tokens_sampled(
&mut self,
token_ids: &[i32],
sampling: Option<&SamplingConfig>,
) -> Result<Vec<i32>> {
self.verify_tokens_sampled_without_mtp(token_ids, sampling)
}
/// Runs batched verification and trims the speculative suffix.
pub fn verify_tokens_rewound(&mut self, token_ids: &[i32]) -> Result<Vec<i32>> {
if token_ids.is_empty() {

84
crates/skippy-runtime/src/tests.rs
View file

@ -216,6 +216,90 @@ mod tests {
Ok(())
}
#[test]
fn batched_sampled_verification_matches_serial_across_lazy_grammar_trigger()
-> anyhow::Result<()> {
let Some(model_path) = correctness_model() else {
eprintln!("skipping: SKIPPY_CORRECTNESS_MODEL is not set");
return Ok(());
};
let model = open_correctness_model(&model_path)?;
let rendered = model.apply_chat_template_json(
r#"[{"role":"user","content":"Call execute_bash."}]"#,
ChatTemplateJsonOptions {
tools_json: Some(
r#"[{"type":"function","function":{"name":"execute_bash","description":"Run a command.","parameters":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}}]"#
.to_string(),
),
..ChatTemplateJsonOptions::default()
},
)?;
let metadata: Value = serde_json::from_str(&rendered.metadata_json)?;
assert!(
metadata
.get("grammar")
.and_then(Value::as_str)
.is_some_and(|grammar| !grammar.is_empty()),
"tool-capable template must produce a grammar"
);
assert_eq!(
metadata.get("grammar_lazy").and_then(Value::as_bool),
Some(true),
"tool grammar must wait for its trigger"
);
let prompt_tokens = model.tokenize(&rendered.prompt, true)?;
assert!(prompt_tokens.len() > 1);
let mut verify_inputs = vec![*prompt_tokens.last().expect("checked nonempty prompt")];
verify_inputs.extend(model.tokenize("<tool_call>", false)?);
verify_inputs.extend(model.tokenize(
"execute_bash<arg_key>command</arg_key><arg_value>pwd</arg_value></tool_call>",
false,
)?);
let sampling = SamplingConfig {
enabled: true,
temperature: 0.0,
top_p: 0.95,
top_k: 40,
min_p: 0.05,
..SamplingConfig::default()
};
let prompt_prefix = &prompt_tokens[..prompt_tokens.len() - 1];
let prompt_token_count = u64::try_from(prompt_tokens.len())?;
let mut serial = model.create_session()?;
serial.prefill_chunked(prompt_prefix)?;
serial.configure_chat_sampling(
&rendered.metadata_json,
prompt_token_count,
Some(&sampling),
)?;
let serial_predictions = verify_inputs
.iter()
.copied()
.map(|token| serial.decode_step_sampled(token, Some(&sampling)))
.collect::<Result<Vec<_>>>()?;
let serial_token_count = serial.token_count();
let serial_native_position = serial.native_position()?;
drop(serial);
let mut batched = model.create_session()?;
batched.prefill_chunked(prompt_prefix)?;
batched.configure_chat_sampling(
&rendered.metadata_json,
prompt_token_count,
Some(&sampling),
)?;
let batched_predictions =
batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?;
assert_eq!(batched_predictions, serial_predictions);
assert_eq!(batched.token_count(), serial_token_count);
assert_eq!(batched.native_position()?, serial_native_position);
Ok(())
}
#[test]
fn stage_session_exposes_non_frame_native_mtp_decode_api() {
type DecodeStepSampledMtp = fn(

21
crates/skippy-server/src/frontend/generation/parsing.rs
View file

@ -7,6 +7,7 @@ use openai_frontend::ChatCompletionResponse;
use openai_frontend::ChatHookAction;
use openai_frontend::ChatHookOutcome;
use openai_frontend::FinishReason;
use openai_frontend::MessageContent;
use openai_frontend::OpenAiError;
use openai_frontend::OpenAiResult;
use serde_json::Value;
@ -396,20 +397,12 @@ pub(in crate::frontend) fn chat_message_generation_value(
) -> OpenAiResult<Value> {
let mut value = serde_json::to_value(message)
.map_err(|error| OpenAiError::invalid_request(format!("serialize message: {error}")))?;
let content = message
.content
.as_ref()
.map(|content| message_content_to_generation_text(content, marker, media))
.transpose()?;
if let Some(object) = value.as_object_mut() {
match content {
Some(content) => {
object.insert("content".to_string(), Value::String(content));
}
None => {
object.insert("content".to_string(), Value::Null);
}
}
let content = match message.content.as_ref() {
Some(MessageContent::Other(Value::Null)) | None => None,
Some(content) => Some(message_content_to_generation_text(content, marker, media)?),
};
if let (Some(object), Some(content)) = (value.as_object_mut(), content) {
object.insert("content".to_string(), Value::String(content));
}
Ok(value)
}

367
crates/skippy-server/src/frontend/linear_proposal.rs
View file

@ -1,3 +1,7 @@
mod execution;
pub(crate) use execution::{LinearProposalExecutionParams, elapsed_us};
use std::{
collections::BTreeMap,
sync::Arc,
@ -9,10 +13,7 @@ use openai_frontend::{OpenAiError, OpenAiResult};
use serde_json::json;
use skippy_runtime::SamplingConfig;
use crate::frontend::{
NativeMtpVerifyWindowDecision, StageOpenAiBackend, TokenControl,
classify_native_mtp_verify_window, openai_backend_error,
};
use crate::frontend::openai_backend_error;
const MAX_OPAQUE_DECISION_ID_BYTES: usize = 64;
const MAX_LINEAR_PROPOSAL_TOKENS: usize = 256;
@ -446,320 +447,16 @@ pub(crate) fn greedy_linear_proposal_admitted(
}
match chat_sampling_metadata {
None => true,
Some(metadata) => serde_json::from_str::<serde_json::Value>(metadata)
.ok()
.is_some_and(|value| {
value
.get("grammar")
.and_then(serde_json::Value::as_str)
.is_none_or(str::is_empty)
}),
Some(metadata) => serde_json::from_str::<serde_json::Value>(metadata).is_ok(),
}
}
struct LinearProposalExecution {
decision: NativeMtpVerifyWindowDecision,
predictions: Vec<i32>,
committed_tokens: Vec<i32>,
reached_stop: bool,
position_after_verification: u64,
canonical_position: u64,
verification_elapsed_us: u64,
repair_elapsed_us: u64,
runtime_lock_wait_us: u64,
runtime_lock_hold_us: u64,
runtime_lock_acquires: usize,
}
#[derive(Clone, Copy)]
pub(crate) struct LinearProposalExecutionParams<'a> {
pub(crate) session_id: &'a str,
pub(crate) current: i32,
pub(crate) base_position: u64,
pub(crate) generated_len: usize,
pub(crate) max_new_tokens: usize,
}
#[derive(Default)]
struct LinearProposalRepairTiming {
elapsed_us: u64,
runtime_lock_wait_us: u64,
runtime_lock_hold_us: u64,
runtime_lock_acquires: usize,
}
impl StageOpenAiBackend {
pub(crate) fn execute_local_linear_proposal(
&self,
params: LinearProposalExecutionParams<'_>,
queried: QueriedLinearProposal,
on_token: &mut impl FnMut(i32) -> OpenAiResult<TokenControl>,
) -> OpenAiResult<Option<LinearProposalReceipt>> {
let proposal_token_count = queried.proposal.token_ids.len();
let mut verify_inputs = Vec::with_capacity(proposal_token_count.saturating_add(1));
verify_inputs.push(params.current);
verify_inputs.extend_from_slice(&queried.proposal.token_ids);
let Some(execution) = self.execute_local_linear_proposal_inner(
params,
&queried.proposal.token_ids,
&verify_inputs,
on_token,
)?
else {
return Ok(None);
};
let accepted_proposal_tokens = execution
.decision
.accepted_proposal_tokens
.min(execution.committed_tokens.len());
let disposition = linear_proposal_disposition(
execution.decision,
proposal_token_count,
execution.committed_tokens.len(),
execution.reached_stop,
);
if execution.committed_tokens.is_empty() {
return Err(OpenAiError::backend(
"linear proposal committed no target token",
));
}
let correction_or_boundary_token = (disposition != LinearProposalDisposition::Stopped)
.then(|| {
execution
.committed_tokens
.last()
.copied()
.expect("checked non-empty committed tokens")
});
let total_elapsed_us = elapsed_us(queried.operation_started);
Ok(Some(LinearProposalReceipt {
decision_id: queried.proposal.decision_id,
disposition,
proposal_token_count,
verification_rows: verify_inputs.len(),
accepted_proposal_tokens,
canonical_prediction_count: execution.committed_tokens.len(),
committed_tokens: execution.committed_tokens.into_boxed_slice(),
verification_row_predictions: execution.predictions.into_boxed_slice(),
correction_or_boundary_token,
base_position: params.base_position,
position_after_verification: execution.position_after_verification,
canonical_position: execution.canonical_position,
trimmed_rows: usize::try_from(
execution
.position_after_verification
.saturating_sub(execution.canonical_position),
)
.map_err(|_| OpenAiError::backend("trimmed row count exceeds usize"))?,
proposal_elapsed_us: queried.proposal_elapsed_us,
verification_elapsed_us: execution.verification_elapsed_us,
repair_elapsed_us: execution.repair_elapsed_us,
total_elapsed_us,
runtime_lock_wait_us: execution.runtime_lock_wait_us,
runtime_lock_hold_us: execution.runtime_lock_hold_us,
runtime_lock_acquires: execution.runtime_lock_acquires,
}))
}
fn execute_local_linear_proposal_inner(
&self,
params: LinearProposalExecutionParams<'_>,
proposal_tokens: &[i32],
verify_inputs: &[i32],
on_token: &mut impl FnMut(i32) -> OpenAiResult<TokenControl>,
) -> OpenAiResult<Option<LinearProposalExecution>> {
let verify_timer = Instant::now();
let verify_lock_timer = Instant::now();
let mut runtime = self
.runtime
.lock()
.map_err(|_| OpenAiError::backend("runtime lock poisoned"))?;
let verify_lock_wait_us = elapsed_us(verify_lock_timer);
let verify_hold_timer = Instant::now();
let observed_position = runtime
.session_token_count(params.session_id)
.ok_or_else(|| OpenAiError::backend("linear proposal session is not active"))?;
if observed_position != params.base_position {
return Ok(None);
}
let predictions = runtime
.verify_tokens(params.session_id, verify_inputs)
.map_err(openai_backend_error)?;
let decision = classify_native_mtp_verify_window(
proposal_tokens,
&predictions,
params.generated_len,
params.max_new_tokens,
|token| {
runtime
.model
.token_is_eog(token)
.map_err(openai_backend_error)
},
)?;
let position_after_verification = runtime
.session_token_count(params.session_id)
.ok_or_else(|| OpenAiError::backend("linear proposal session disappeared"))?;
let expected_position_after_verification = params
.base_position
.checked_add(
u64::try_from(verify_inputs.len())
.map_err(|_| OpenAiError::backend("verification row count exceeds u64"))?,
)
.ok_or_else(|| OpenAiError::backend("linear proposal position overflow"))?;
if position_after_verification != expected_position_after_verification {
return Err(OpenAiError::backend(format!(
"linear proposal verification position mismatch: observed {position_after_verification}, expected {expected_position_after_verification}"
)));
}
let verify_lock_hold_us = elapsed_us(verify_hold_timer);
drop(runtime);
let verification_elapsed_us = elapsed_us(verify_timer);
let mut committed_tokens = Vec::with_capacity(decision.commit_count);
let mut reached_stop = false;
let mut callback_error = None;
for token in predictions.iter().copied().take(decision.commit_count) {
committed_tokens.push(token);
match on_token(token) {
Ok(TokenControl::Continue) => {}
Ok(TokenControl::Stop) => {
reached_stop = true;
break;
}
Err(error) => {
callback_error = Some(error);
break;
}
}
}
if committed_tokens.is_empty() {
return Err(OpenAiError::backend(
"linear proposal classifier committed no target prediction",
));
}
let canonical_position = params
.base_position
.checked_add(
u64::try_from(committed_tokens.len())
.map_err(|_| OpenAiError::backend("committed token count exceeds u64"))?,
)
.ok_or_else(|| OpenAiError::backend("linear proposal canonical position overflow"))?;
let repair = finish_linear_proposal_after_repair(callback_error, || {
self.trim_branch_suffix_or_retire(
params.session_id,
params.base_position,
verify_inputs.len(),
canonical_position,
position_after_verification,
)
})?;
Ok(Some(LinearProposalExecution {
decision,
predictions,
committed_tokens,
reached_stop,
position_after_verification,
canonical_position,
verification_elapsed_us,
repair_elapsed_us: repair.elapsed_us,
runtime_lock_wait_us: verify_lock_wait_us.saturating_add(repair.runtime_lock_wait_us),
runtime_lock_hold_us: verify_lock_hold_us.saturating_add(repair.runtime_lock_hold_us),
runtime_lock_acquires: 1usize.saturating_add(repair.runtime_lock_acquires),
}))
}
fn trim_branch_suffix_or_retire(
&self,
session_id: &str,
checkpoint_start: u64,
checkpoint_count: usize,
canonical_position: u64,
position_after_verification: u64,
) -> OpenAiResult<LinearProposalRepairTiming> {
if canonical_position >= position_after_verification {
let mut runtime = self.runtime.lock().map_err(|_| {
OpenAiError::backend("runtime lock poisoned during verify retirement")
})?;
runtime
.retire_verify_checkpoint(session_id, checkpoint_start, checkpoint_count as u64)
.map_err(openai_backend_error)?;
return Ok(LinearProposalRepairTiming::default());
}
let repair_timer = Instant::now();
let repair_lock_timer = Instant::now();
let mut runtime = self
.runtime
.lock()
.map_err(|_| OpenAiError::backend("runtime lock poisoned during proposal repair"))?;
let runtime_lock_wait_us = elapsed_us(repair_lock_timer);
let repair_hold_timer = Instant::now();
let trim_result = runtime.trim_session(session_id, canonical_position);
let runtime_lock_hold_us = elapsed_us(repair_hold_timer);
if let Err(error) = trim_result {
let _ = runtime.drop_session_timed(session_id);
return Err(OpenAiError::backend(format!(
"linear proposal repair failed and the session was retired: {error:#}"
)));
}
let repaired_position = runtime
.session_token_count(session_id)
.ok_or_else(|| OpenAiError::backend("repaired linear proposal session disappeared"))?;
if repaired_position != canonical_position {
let _ = runtime.drop_session_timed(session_id);
return Err(OpenAiError::backend(format!(
"linear proposal repair position mismatch: observed {repaired_position}, expected {canonical_position}"
)));
}
Ok(LinearProposalRepairTiming {
elapsed_us: elapsed_us(repair_timer),
runtime_lock_wait_us,
runtime_lock_hold_us,
runtime_lock_acquires: 1,
})
}
}
fn elapsed_us(started: Instant) -> u64 {
u64::try_from(started.elapsed().as_micros()).unwrap_or(u64::MAX)
}
fn linear_proposal_disposition(
decision: NativeMtpVerifyWindowDecision,
proposal_token_count: usize,
committed_token_count: usize,
reached_stop: bool,
) -> LinearProposalDisposition {
if reached_stop
|| (!decision.rejected
&& (decision.accepted_proposal_tokens != proposal_token_count
|| committed_token_count != proposal_token_count.saturating_add(1)))
{
LinearProposalDisposition::Stopped
} else if decision.rejected {
LinearProposalDisposition::FirstMismatch
} else {
LinearProposalDisposition::FullAccept
}
}
fn finish_linear_proposal_after_repair<T>(
callback_error: Option<OpenAiError>,
repair: impl FnOnce() -> OpenAiResult<T>,
) -> OpenAiResult<T> {
let repaired = repair()?;
callback_error.map_or(Ok(repaired), Err)
}
#[cfg(test)]
mod tests {
use std::{cell::Cell, sync::Mutex, thread};
use std::{sync::Mutex, thread};
use super::*;
use crate::frontend::{NativeMtpVerifyWindowDecision, classify_native_mtp_verify_window};
#[derive(Debug, PartialEq, Eq)]
struct RecordedQuery {
@ -865,50 +562,6 @@ mod tests {
}
}
#[test]
fn disposition_distinguishes_full_mismatch_and_early_stop() {
let full = decision(&[11, 12], &[11, 12, 13]);
assert_eq!(
linear_proposal_disposition(full, 2, 3, false),
LinearProposalDisposition::FullAccept
);
let mismatch = decision(&[11, 12], &[11, 99, 13]);
assert_eq!(
linear_proposal_disposition(mismatch, 2, 2, false),
LinearProposalDisposition::FirstMismatch
);
assert_eq!(
linear_proposal_disposition(full, 2, 1, true),
LinearProposalDisposition::Stopped
);
assert_eq!(
linear_proposal_disposition(full, 2, 1, false),
LinearProposalDisposition::Stopped
);
}
#[test]
fn callback_error_is_returned_only_after_repair_runs() {
let repair_ran = Cell::new(false);
let result = finish_linear_proposal_after_repair(
Some(OpenAiError::backend("synthetic callback failure")),
|| {
repair_ran.set(true);
Ok(())
},
);
assert!(repair_ran.get());
assert!(
result
.unwrap_err()
.to_string()
.contains("synthetic callback failure")
);
}
#[test]
fn execution_error_discards_exactly_once_without_masking_primary_error() {
let source = Arc::new(FakeIngress::default());
@ -988,7 +641,7 @@ mod tests {
}
#[test]
fn greedy_admission_accepts_zero_temperature_and_rejects_logit_modifiers() {
fn greedy_admission_rejects_stochastic_sampling_but_accepts_valid_grammar_metadata() {
let disabled = SamplingConfig::default();
let temperature_zero = SamplingConfig {
enabled: true,
@ -1022,7 +675,7 @@ mod tests {
assert!(greedy_linear_proposal_admitted(&temperature_zero, None));
assert!(!greedy_linear_proposal_admitted(&stochastic, None));
assert!(!greedy_linear_proposal_admitted(&biased_greedy, None));
assert!(!greedy_linear_proposal_admitted(
assert!(greedy_linear_proposal_admitted(
&disabled,
Some(r#"{"grammar":"root ::= value"}"#)
));

409
crates/skippy-server/src/frontend/linear_proposal/execution.rs Normal file
View file

@ -0,0 +1,409 @@
use std::time::Instant;
use openai_frontend::{OpenAiError, OpenAiResult};
use skippy_runtime::SamplingConfig;
use crate::frontend::{
LinearProposalDisposition, NativeMtpVerifyWindowDecision, StageOpenAiBackend, TokenControl,
classify_native_mtp_verify_window, openai_backend_error,
};
use super::{LinearProposalReceipt, QueriedLinearProposal};
struct LinearProposalExecution {
decision: NativeMtpVerifyWindowDecision,
predictions: Vec<i32>,
committed_tokens: Vec<i32>,
reached_stop: bool,
position_after_verification: u64,
canonical_position: u64,
verification_elapsed_us: u64,
repair_elapsed_us: u64,
runtime_lock_wait_us: u64,
runtime_lock_hold_us: u64,
runtime_lock_acquires: usize,
}
#[derive(Clone, Copy)]
pub(crate) struct LinearProposalExecutionParams<'a> {
pub(crate) session_id: &'a str,
pub(crate) current: i32,
pub(crate) base_position: u64,
pub(crate) generated_len: usize,
pub(crate) max_new_tokens: usize,
pub(crate) sampling: &'a SamplingConfig,
pub(crate) chat_sampling_metadata: Option<&'a str>,
pub(crate) prompt_token_count: usize,
}
#[derive(Default)]
struct LinearProposalRepairTiming {
elapsed_us: u64,
runtime_lock_wait_us: u64,
runtime_lock_hold_us: u64,
runtime_lock_acquires: usize,
}
struct LinearProposalRepairParams<'a> {
session_id: &'a str,
checkpoint_start: u64,
checkpoint_count: usize,
canonical_position: u64,
position_after_verification: u64,
sampling: &'a SamplingConfig,
chat_sampling_metadata: Option<&'a str>,
prompt_token_count: usize,
}
impl StageOpenAiBackend {
pub(crate) fn execute_local_linear_proposal(
&self,
params: LinearProposalExecutionParams<'_>,
queried: QueriedLinearProposal,
on_token: &mut impl FnMut(i32) -> OpenAiResult<TokenControl>,
) -> OpenAiResult<Option<LinearProposalReceipt>> {
let proposal_token_count = queried.proposal.token_ids.len();
let mut verify_inputs = Vec::with_capacity(proposal_token_count.saturating_add(1));
verify_inputs.push(params.current);
verify_inputs.extend_from_slice(&queried.proposal.token_ids);
let Some(execution) = self.execute_local_linear_proposal_inner(
params,
&queried.proposal.token_ids,
&verify_inputs,
on_token,
)?
else {
return Ok(None);
};
let accepted_proposal_tokens = execution
.decision
.accepted_proposal_tokens
.min(execution.committed_tokens.len());
let disposition = linear_proposal_disposition(
execution.decision,
proposal_token_count,
execution.committed_tokens.len(),
execution.reached_stop,
);
if execution.committed_tokens.is_empty() {
return Err(OpenAiError::backend(
"linear proposal committed no target token",
));
}
let correction_or_boundary_token = (disposition != LinearProposalDisposition::Stopped)
.then(|| {
execution
.committed_tokens
.last()
.copied()
.expect("checked non-empty committed tokens")
});
let total_elapsed_us = elapsed_us(queried.operation_started);
Ok(Some(LinearProposalReceipt {
decision_id: queried.proposal.decision_id,
disposition,
proposal_token_count,
verification_rows: verify_inputs.len(),
accepted_proposal_tokens,
canonical_prediction_count: execution.committed_tokens.len(),
committed_tokens: execution.committed_tokens.into_boxed_slice(),
verification_row_predictions: execution.predictions.into_boxed_slice(),
correction_or_boundary_token,
base_position: params.base_position,
position_after_verification: execution.position_after_verification,
canonical_position: execution.canonical_position,
trimmed_rows: usize::try_from(
execution
.position_after_verification
.saturating_sub(execution.canonical_position),
)
.map_err(|_| OpenAiError::backend("trimmed row count exceeds usize"))?,
proposal_elapsed_us: queried.proposal_elapsed_us,
verification_elapsed_us: execution.verification_elapsed_us,
repair_elapsed_us: execution.repair_elapsed_us,
total_elapsed_us,
runtime_lock_wait_us: execution.runtime_lock_wait_us,
runtime_lock_hold_us: execution.runtime_lock_hold_us,
runtime_lock_acquires: execution.runtime_lock_acquires,
}))
}
fn execute_local_linear_proposal_inner(
&self,
params: LinearProposalExecutionParams<'_>,
proposal_tokens: &[i32],
verify_inputs: &[i32],
on_token: &mut impl FnMut(i32) -> OpenAiResult<TokenControl>,
) -> OpenAiResult<Option<LinearProposalExecution>> {
let verify_timer = Instant::now();
let verify_lock_timer = Instant::now();
let mut runtime = self
.runtime
.lock()
.map_err(|_| OpenAiError::backend("runtime lock poisoned"))?;
let verify_lock_wait_us = elapsed_us(verify_lock_timer);
let verify_hold_timer = Instant::now();
let observed_position = runtime
.session_token_count(params.session_id)
.ok_or_else(|| OpenAiError::backend("linear proposal session is not active"))?;
if observed_position != params.base_position {
return Ok(None);
}
let predictions = runtime
.verify_tokens_sampled(
params.session_id,
verify_inputs,
params.sampling.enabled.then_some(params.sampling),
)
.map_err(openai_backend_error)?;
let decision = classify_native_mtp_verify_window(
proposal_tokens,
&predictions,
params.generated_len,
params.max_new_tokens,
|token| {
runtime
.model
.token_is_eog(token)
.map_err(openai_backend_error)
},
)?;
let position_after_verification = runtime
.session_token_count(params.session_id)
.ok_or_else(|| OpenAiError::backend("linear proposal session disappeared"))?;
let expected_position_after_verification = params
.base_position
.checked_add(
u64::try_from(verify_inputs.len())
.map_err(|_| OpenAiError::backend("verification row count exceeds u64"))?,
)
.ok_or_else(|| OpenAiError::backend("linear proposal position overflow"))?;
if position_after_verification != expected_position_after_verification {
return Err(OpenAiError::backend(format!(
"linear proposal verification position mismatch: observed {position_after_verification}, expected {expected_position_after_verification}"
)));
}
let verify_lock_hold_us = elapsed_us(verify_hold_timer);
drop(runtime);
let verification_elapsed_us = elapsed_us(verify_timer);
let mut committed_tokens = Vec::with_capacity(decision.commit_count);
let mut reached_stop = false;
let mut callback_error = None;
for token in predictions.iter().copied().take(decision.commit_count) {
committed_tokens.push(token);
match on_token(token) {
Ok(TokenControl::Continue) => {}
Ok(TokenControl::Stop) => {
reached_stop = true;
break;
}
Err(error) => {
callback_error = Some(error);
break;
}
}
}
if committed_tokens.is_empty() {
return Err(OpenAiError::backend(
"linear proposal classifier committed no target prediction",
));
}
let canonical_position = params
.base_position
.checked_add(
u64::try_from(committed_tokens.len())
.map_err(|_| OpenAiError::backend("committed token count exceeds u64"))?,
)
.ok_or_else(|| OpenAiError::backend("linear proposal canonical position overflow"))?;
let repair = finish_linear_proposal_after_repair(callback_error, || {
self.trim_branch_suffix_or_retire(LinearProposalRepairParams {
session_id: params.session_id,
checkpoint_start: params.base_position,
checkpoint_count: verify_inputs.len(),
canonical_position,
position_after_verification,
sampling: params.sampling,
chat_sampling_metadata: params.chat_sampling_metadata,
prompt_token_count: params.prompt_token_count,
})
})?;
Ok(Some(LinearProposalExecution {
decision,
predictions,
committed_tokens,
reached_stop,
position_after_verification,
canonical_position,
verification_elapsed_us,
repair_elapsed_us: repair.elapsed_us,
runtime_lock_wait_us: verify_lock_wait_us.saturating_add(repair.runtime_lock_wait_us),
runtime_lock_hold_us: verify_lock_hold_us.saturating_add(repair.runtime_lock_hold_us),
runtime_lock_acquires: 1usize.saturating_add(repair.runtime_lock_acquires),
}))
}
fn trim_branch_suffix_or_retire(
&self,
params: LinearProposalRepairParams<'_>,
) -> OpenAiResult<LinearProposalRepairTiming> {
let LinearProposalRepairParams {
session_id,
checkpoint_start,
checkpoint_count,
canonical_position,
position_after_verification,
sampling,
chat_sampling_metadata,
prompt_token_count,
} = params;
if canonical_position >= position_after_verification {
let retire_timer = Instant::now();
let retire_lock_timer = Instant::now();
let mut runtime = self.runtime.lock().map_err(|_| {
OpenAiError::backend("runtime lock poisoned during verify retirement")
})?;
let runtime_lock_wait_us = elapsed_us(retire_lock_timer);
let retire_hold_timer = Instant::now();
runtime
.retire_verify_checkpoint(session_id, checkpoint_start, checkpoint_count as u64)
.map_err(openai_backend_error)?;
return Ok(LinearProposalRepairTiming {
elapsed_us: elapsed_us(retire_timer),
runtime_lock_wait_us,
runtime_lock_hold_us: elapsed_us(retire_hold_timer),
runtime_lock_acquires: 1,
});
}
let repair_timer = Instant::now();
let repair_lock_timer = Instant::now();
let mut runtime = self
.runtime
.lock()
.map_err(|_| OpenAiError::backend("runtime lock poisoned during proposal repair"))?;
let runtime_lock_wait_us = elapsed_us(repair_lock_timer);
let repair_hold_timer = Instant::now();
let trim_result = runtime.trim_session(session_id, canonical_position);
let runtime_lock_hold_us = elapsed_us(repair_hold_timer);
if let Err(error) = trim_result {
let _ = runtime.drop_session_timed(session_id);
return Err(OpenAiError::backend(format!(
"linear proposal repair failed and the session was retired: {error:#}"
)));
}
if let Some(metadata) = chat_sampling_metadata {
runtime
.configure_chat_sampling(
session_id,
metadata,
u64::try_from(prompt_token_count).unwrap_or(u64::MAX),
sampling.enabled.then_some(sampling),
)
.map_err(openai_backend_error)?;
}
let repaired_position = runtime
.session_token_count(session_id)
.ok_or_else(|| OpenAiError::backend("repaired linear proposal session disappeared"))?;
if repaired_position != canonical_position {
let _ = runtime.drop_session_timed(session_id);
return Err(OpenAiError::backend(format!(
"linear proposal repair position mismatch: observed {repaired_position}, expected {canonical_position}"
)));
}
Ok(LinearProposalRepairTiming {
elapsed_us: elapsed_us(repair_timer),
runtime_lock_wait_us,
runtime_lock_hold_us,
runtime_lock_acquires: 1,
})
}
}
pub(crate) fn elapsed_us(started: Instant) -> u64 {
u64::try_from(started.elapsed().as_micros()).unwrap_or(u64::MAX)
}
fn linear_proposal_disposition(
decision: NativeMtpVerifyWindowDecision,
proposal_token_count: usize,
committed_token_count: usize,
reached_stop: bool,
) -> LinearProposalDisposition {
if reached_stop
|| (!decision.rejected
&& (decision.accepted_proposal_tokens != proposal_token_count
|| committed_token_count != proposal_token_count.saturating_add(1)))
{
LinearProposalDisposition::Stopped
} else if decision.rejected {
LinearProposalDisposition::FirstMismatch
} else {
LinearProposalDisposition::FullAccept
}
}
fn finish_linear_proposal_after_repair<T>(
callback_error: Option<OpenAiError>,
repair: impl FnOnce() -> OpenAiResult<T>,
) -> OpenAiResult<T> {
let repaired = repair()?;
callback_error.map_or(Ok(repaired), Err)
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
fn decision(proposal: &[i32], predictions: &[i32]) -> NativeMtpVerifyWindowDecision {
classify_native_mtp_verify_window(proposal, predictions, 0, 64, |_| Ok(false)).unwrap()
}
#[test]
fn disposition_distinguishes_full_mismatch_and_early_stop() {
let full = decision(&[11, 12], &[11, 12, 13]);
assert_eq!(
linear_proposal_disposition(full, 2, 3, false),
LinearProposalDisposition::FullAccept
);
let mismatch = decision(&[11, 12], &[11, 99, 13]);
assert_eq!(
linear_proposal_disposition(mismatch, 2, 2, false),
LinearProposalDisposition::FirstMismatch
);
assert_eq!(
linear_proposal_disposition(full, 2, 1, true),
LinearProposalDisposition::Stopped
);
assert_eq!(
linear_proposal_disposition(full, 2, 1, false),
LinearProposalDisposition::Stopped
);
}
#[test]
fn callback_error_is_returned_only_after_repair_runs() {
let repair_ran = Cell::new(false);
let result = finish_linear_proposal_after_repair(
Some(OpenAiError::backend("synthetic callback failure")),
|| {
repair_ran.set(true);
Ok(())
},
);
assert!(repair_ran.get());
assert!(
result
.unwrap_err()
.to_string()
.contains("synthetic callback failure")
);
}
}

1200
crates/skippy-server/src/frontend/local_generation.rs
View file

File diff suppressed because it is too large Load diff

220
crates/skippy-server/src/frontend/local_generation/tests.rs Normal file
View file

@ -0,0 +1,220 @@
use std::sync::{
Arc, Mutex,
atomic::{AtomicBool, AtomicUsize, Ordering},
};
use anyhow::{Result, bail};
use skippy_protocol::{LoadMode, StageConfig};
use skippy_runtime::SamplingConfig;
use tokio::sync::Semaphore;
use crate::binary_transport::DecodeFrameBatcher;
use crate::frontend::SpeculativeDecodeConfig;
use crate::frontend::admission::GenerationTokenBudget;
use crate::frontend::decode_batcher::DecodeBatcher;
use crate::frontend::generation::{
LocalGeneration, OpenAiBackendMode, OpenAiCacheHints, OpenAiGenerationIds, StageOpenAiBackend,
TokenControl,
};
use crate::frontend::local_generation::{
native_mtp_dispatch_counts_for_test, prompt_fits_single_prefill_sample,
};
use crate::frontend::{
EmbeddedOpenAiRequestDefaults, GenerationReceipt, GenerationReceiptConfig,
GenerationReceiptSink, GenerationTermination,
};
use crate::runtime_state::load_runtime;
use crate::telemetry::{Telemetry, TelemetryLevel};
#[derive(Default)]
struct RecordingReceiptSink {
receipts: Mutex<Vec<GenerationReceipt>>,
fail: AtomicBool,
}
impl GenerationReceiptSink for RecordingReceiptSink {
fn record(&self, receipt: &GenerationReceipt) -> Result<()> {
self.receipts.lock().unwrap().push(receipt.clone());
if self.fail.load(Ordering::Relaxed) {
bail!("synthetic generation receipt sink failure");
}
Ok(())
}
}
#[test]
fn single_prefill_sample_requires_prompt_to_fit_session_batch() {
assert!(!prompt_fits_single_prefill_sample(0, 2048));
assert!(!prompt_fits_single_prefill_sample(1, 2048));
assert!(prompt_fits_single_prefill_sample(2048, 2048));
assert!(!prompt_fits_single_prefill_sample(2049, 2048));
}
#[test]
fn local_native_mtp_decode_uses_non_frame_runtime_api() {
let (sampled_calls, frame_calls) = native_mtp_dispatch_counts_for_test();
assert_eq!(sampled_calls, 1);
assert_eq!(frame_calls, 0);
}
#[test]
fn local_generation_delivers_receipt_before_cleanup_and_propagates_sink_errors() -> Result<()> {
let Some(model_path) = std::env::var_os("SKIPPY_GENERATION_RECEIPT_MODEL") else {
eprintln!("skipping: SKIPPY_GENERATION_RECEIPT_MODEL is not set");
return Ok(());
};
let Some(layer_count) = std::env::var_os("SKIPPY_GENERATION_RECEIPT_MODEL_LAYERS") else {
eprintln!("skipping: SKIPPY_GENERATION_RECEIPT_MODEL_LAYERS is not set");
return Ok(());
};
let layer_count = layer_count
.to_string_lossy()
.parse::<u32>()
.map_err(|error| anyhow::anyhow!("invalid receipt test layer count: {error}"))?;
let config = StageConfig {
run_id: "generation-receipt-test".to_string(),
topology_id: "generation-receipt-test".to_string(),
model_id: "generation-receipt-test".to_string(),
package_ref: None,
manifest_sha256: None,
source_model_path: None,
source_model_sha256: None,
source_model_bytes: None,
materialized_path: None,
materialized_pinned: false,
model_path: Some(model_path.to_string_lossy().into_owned()),
projector_path: None,
stage_id: "stage-0".to_string(),
stage_index: 0,
layer_start: 0,
layer_end: layer_count,
ctx_size: 128,
lane_count: 1,
n_batch: Some(32),
n_ubatch: Some(32),
n_gpu_layers: 0,
mmap: Some(true),
mlock: false,
cache_type_k: "f16".to_string(),
cache_type_v: "f16".to_string(),
flash_attn_type: Default::default(),
filter_tensors_on_load: false,
selected_device: None,
kv_cache: None,
native_mtp_enabled: false,
load_mode: LoadMode::RuntimeSlice,
bind_addr: "127.0.0.1:0".to_string(),
upstream: None,
downstream: None,
};
let runtime = load_runtime(&config)?
.ok_or_else(|| anyhow::anyhow!("receipt test runtime was not loaded"))?;
let sink = Arc::new(RecordingReceiptSink::default());
let telemetry = Telemetry::new(None, 1, config.clone(), TelemetryLevel::Off);
let speculative = SpeculativeDecodeConfig::default();
let decode_batcher = DecodeBatcher::new(runtime.clone(), 1);
let decode_frame_batcher = DecodeFrameBatcher::new(runtime.clone(), 1);
let backend = StageOpenAiBackend {
runtime: runtime.clone(),
config,
telemetry,
model_id: "generation-receipt-test".to_string(),
default_max_tokens: 1,
request_defaults: EmbeddedOpenAiRequestDefaults::default(),
ctx_size: 128,
mode: OpenAiBackendMode::LocalRuntime,
draft: None,
speculative_window: 0,
adaptive_speculative_window: false,
ngram_max: 0,
speculative: speculative.clone(),
generation_limit: Arc::new(Semaphore::new(1)),
generation_queue_depth: Arc::new(AtomicUsize::new(0)),
generation_queue_limit: 1,
generation_token_budget: Arc::new(GenerationTokenBudget::new(128)),
hook_policy: None,
generation_receipt: Some(GenerationReceiptConfig::new(sink.clone())),
linear_proposal_ingress: None,
kv: None,
decode_batcher,
decode_frame_batcher,
};
let sampling = SamplingConfig::default();
let prompt_token_ids = [1];
let ids = OpenAiGenerationIds::new(OpenAiCacheHints::default());
let mut emitted = Vec::new();
backend.generate_local_tokens(
LocalGeneration {
prompt_token_ids: &prompt_token_ids,
max_tokens: 1,
sampling: &sampling,
chat_sampling_metadata: None,
speculative: &speculative,
native_mtp_enabled: false,
hook_request: None,
hook_runtime: None,
cancellation: None,
ids: &ids,
},
|token_id| {
emitted.push(token_id);
Ok(TokenControl::Continue)
},
)?;
let receipts = sink.receipts.lock().unwrap();
assert_eq!(receipts.len(), 1);
assert_eq!(receipts[0].request_id, ids.request_id);
assert_eq!(receipts[0].session_id, ids.session_id);
assert_eq!(receipts[0].prompt_token_count, prompt_token_ids.len());
assert_eq!(receipts[0].generated_token_ids.as_ref(), emitted.as_slice());
assert_eq!(receipts[0].termination, GenerationTermination::MaxTokens);
assert!(receipts[0].final_session_position >= prompt_token_ids.len() as u64);
drop(receipts);
assert!(
runtime
.lock()
.unwrap()
.session_stats()
.lanes
.iter()
.all(|lane| lane.session_id.as_deref() != Some(&ids.session_label))
);
sink.fail.store(true, Ordering::Relaxed);
let failing_ids = OpenAiGenerationIds::new(OpenAiCacheHints::default());
let error = match backend.generate_local_tokens(
LocalGeneration {
prompt_token_ids: &prompt_token_ids,
max_tokens: 1,
sampling: &sampling,
chat_sampling_metadata: None,
speculative: &speculative,
native_mtp_enabled: false,
hook_request: None,
hook_runtime: None,
cancellation: None,
ids: &failing_ids,
},
|_| Ok(TokenControl::Continue),
) {
Ok(_) => panic!("sink failure should fail local generation"),
Err(error) => error,
};
assert!(
error
.to_string()
.contains("synthetic generation receipt sink failure")
);
assert!(
runtime
.lock()
.unwrap()
.session_stats()
.lanes
.iter()
.all(|lane| lane.session_id.as_deref() != Some(&failing_ids.session_label))
);
Ok(())
}

1315
crates/skippy-server/src/frontend/local_generation/token_generation.rs Normal file
View file

File diff suppressed because it is too large Load diff

19
crates/skippy-server/src/frontend/tests/prompting.rs
View file

@ -552,3 +552,22 @@ fn chat_message_generation_value_preserves_tool_history() {
assert_eq!(value["tool_calls"][0]["id"], "call_123");
assert_eq!(value["tool_calls"][0]["function"]["name"], "lookup");
}
#[test]
fn chat_message_generation_value_preserves_omitted_content() {
let message: openai_frontend::ChatMessage = serde_json::from_value(json!({
"role": "assistant",
"tool_calls": [{
"id": "call_123",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"}
}]
}))
.unwrap();
let mut media = Vec::new();
let value = chat_message_generation_value(&message, "<__media__>", &mut media).unwrap();
assert!(!value.as_object().unwrap().contains_key("content"));
assert_eq!(value["tool_calls"][0]["id"], "call_123");
}

1203
crates/skippy-server/src/runtime_state.rs
View file

File diff suppressed because it is too large Load diff

556
crates/skippy-server/src/runtime_state/frame_operations.rs Normal file
View file

@ -0,0 +1,556 @@
use super::*;
impl RuntimeState {
pub fn prefill(&mut self, session_id: &str, token_ids: &[i32]) -> Result<()> {
let session = self.session(session_id)?;
session.prefill_chunked(token_ids)?;
self.add_session_tokens(session_id, token_ids.len() as u64);
Ok(())
}
pub fn media_marker(&self) -> String {
self.model.media_marker()
}
pub fn has_media_projector(&self) -> bool {
self.model.has_media_projector()
}
pub fn prefill_media(
&mut self,
session_id: &str,
prompt: &str,
media: &[MediaInput],
sampling: Option<&SamplingConfig>,
) -> Result<MediaPrefill> {
let model = &self.model as *const StageModel;
let session = self.session(session_id)?;
// `session()` mutably borrows the session map, while the projector lives
// on the same RuntimeState. RuntimeState serializes access behind one
// outer mutex, so this split borrow only aliases immutable model state.
let prefill = unsafe { (&*model).prefill_media(session, prompt, media, sampling) }?;
self.session_token_counts
.insert(session_id.to_string(), prefill.position);
Ok(prefill)
}
pub fn prefill_media_frame(
&mut self,
session_id: &str,
prompt: &str,
media: &[MediaInput],
) -> Result<MediaPrefillFrame> {
let model = &self.model as *const StageModel;
let session = self.session(session_id)?;
// `session()` mutably borrows the session map, while the projector lives
// on the same RuntimeState. RuntimeState serializes access behind one
// outer mutex, so this split borrow only aliases immutable model state.
let prefill = unsafe { (&*model).prefill_media_frame(session, prompt, media) }?;
self.session_token_counts
.insert(session_id.to_string(), prefill.position);
Ok(prefill)
}
pub fn decode(&mut self, session_id: &str, token_id: i32) -> Result<i32> {
self.decode_sampled(session_id, token_id, None)
}
pub fn decode_sampled(
&mut self,
session_id: &str,
token_id: i32,
sampling: Option<&SamplingConfig>,
) -> Result<i32> {
let session = self.session(session_id)?;
let token = session.decode_step_sampled(token_id, sampling)?;
self.add_session_tokens(session_id, 1);
Ok(token)
}
pub fn decode_batch_sampled(
&mut self,
requests: &[RuntimeDecodeBatchRequest<'_>],
) -> Result<Vec<i32>> {
if requests.is_empty() {
return Ok(Vec::new());
}
Self::ensure_unique_batch_sessions(requests)?;
for request in requests {
self.session(request.session_id)?;
}
let mut lane_sessions = Vec::with_capacity(requests.len());
for request in requests {
let lane_session = self.sessions.remove(request.session_id).ok_or_else(|| {
anyhow::anyhow!(
"session {} was not active after admission",
request.session_id
)
})?;
lane_sessions.push((request.session_id.to_string(), lane_session));
}
let result = {
let mut decode_requests = lane_sessions
.iter_mut()
.zip(requests.iter())
.map(|((_, lane_session), request)| DecodeBatchRequest {
session: &mut lane_session.session,
token_id: request.token_id,
sampling: request.sampling,
})
.collect::<Vec<_>>();
StageSession::decode_batch_sampled(&mut decode_requests)
};
for (session_id, lane_session) in lane_sessions {
self.sessions.insert(session_id, lane_session);
}
if result.is_ok() {
for request in requests {
self.add_session_tokens(request.session_id, 1);
}
}
result
}
pub fn session_batch_size(&mut self, session_id: &str) -> Result<usize> {
self.active_session(session_id)?.batch_size()
}
pub fn ensure_session_active(&mut self, session_id: &str) -> Result<()> {
self.session(session_id).map(|_| ())
}
pub fn configure_chat_sampling(
&mut self,
session_id: &str,
metadata_json: &str,
prompt_token_count: u64,
sampling: Option<&SamplingConfig>,
) -> Result<()> {
self.session(session_id)?.configure_chat_sampling(
metadata_json,
prompt_token_count,
sampling,
)
}
pub fn last_token_signal(&mut self, session_id: &str) -> Result<TokenSignal> {
self.session(session_id)?.last_token_signal()
}
pub fn signal_window(
&mut self,
session_id: &str,
window_tokens: u32,
) -> Result<GenerationSignalWindow> {
self.session(session_id)?.signal_window(window_tokens)
}
pub fn prefill_frame(
&mut self,
session_id: &str,
token_ids: &[i32],
input: Option<&ActivationFrame>,
) -> Result<ActivationFrame> {
self.prefill_frame_with_positions(session_id, token_ids, &[], input)
}
pub fn prefill_frame_with_positions(
&mut self,
session_id: &str,
token_ids: &[i32],
positions: &[i32],
input: Option<&ActivationFrame>,
) -> Result<ActivationFrame> {
let session = self.session(session_id)?;
let frame = session.prefill_chunk_frame_with_positions(token_ids, positions, input, 0)?;
self.add_session_tokens(session_id, token_ids.len() as u64);
Ok(frame)
}
pub fn prefill_final_frame_sampled(
&mut self,
session_id: &str,
token_ids: &[i32],
positions: &[i32],
sampling: Option<&SamplingConfig>,
input: Option<&ActivationFrame>,
) -> Result<(i32, ActivationFrame)> {
let session = self.session(session_id)?;
let (predicted, frame) = session
.prefill_chunk_frame_sampled_with_positions(token_ids, positions, sampling, input, 0)?;
self.add_session_tokens(session_id, token_ids.len() as u64);
Ok((predicted, frame))
}
#[allow(dead_code)]
pub fn decode_frame(
&mut self,
session_id: &str,
token_id: i32,
input: Option<&ActivationFrame>,
) -> Result<(i32, ActivationFrame)> {
self.decode_frame_sampled(session_id, token_id, None, input, 0)
}
pub fn decode_frame_sampled(
&mut self,
session_id: &str,
token_id: i32,
sampling: Option<&SamplingConfig>,
input: Option<&ActivationFrame>,
output_capacity: usize,
) -> Result<(i32, ActivationFrame)> {
let session = self.session(session_id)?;
let output =
session.decode_step_frame_sampled(token_id, sampling, input, output_capacity)?;
self.add_session_tokens(session_id, 1);
Ok(output)
}
pub fn decode_frame_sampled_mtp(
&mut self,
session_id: &str,
token_id: i32,
sampling: Option<&SamplingConfig>,
input: Option<&ActivationFrame>,
output_capacity: usize,
max_draft_tokens: usize,
) -> Result<(i32, Option<NativeMtpDraft>, ActivationFrame)> {
let session = self.session(session_id)?;
let output = session.decode_step_frame_sampled_mtp(
token_id,
sampling,
input,
output_capacity,
max_draft_tokens,
)?;
self.add_session_tokens(session_id, 1);
Ok(output)
}
pub fn decode_sampled_mtp(
&mut self,
session_id: &str,
token_id: i32,
sampling: Option<&SamplingConfig>,
max_draft_tokens: usize,
) -> Result<(i32, Option<NativeMtpDraft>)> {
let session = self.session(session_id)?;
let output = session.decode_step_sampled_mtp(token_id, sampling, max_draft_tokens)?;
self.add_session_tokens(session_id, 1);
Ok(output)
}
pub fn decode_frame_batch_sampled(
&mut self,
requests: &[RuntimeDecodeFrameBatchRequest<'_>],
) -> Result<Vec<DecodeFrameBatchOutput>> {
if requests.is_empty() {
return Ok(Vec::new());
}
Self::ensure_unique_frame_batch_sessions(requests)?;
for request in requests {
self.session(request.session_id)?;
}
let mut lane_sessions = Vec::with_capacity(requests.len());
for request in requests {
let lane_session = self.sessions.remove(request.session_id).ok_or_else(|| {
anyhow::anyhow!(
"session {} was not active after admission",
request.session_id
)
})?;
lane_sessions.push((request.session_id.to_string(), lane_session));
}
let result = {
let mut decode_requests = lane_sessions
.iter_mut()
.zip(requests.iter())
.map(|((_, lane_session), request)| DecodeFrameBatchRequest {
session: &mut lane_session.session,
token_id: request.token_id,
sampling: request.sampling,
input: request.input,
})
.collect::<Vec<_>>();
StageSession::decode_step_frame_batch_sampled(&mut decode_requests)
};
for (session_id, lane_session) in lane_sessions {
self.sessions.insert(session_id, lane_session);
}
if result.is_ok() {
for request in requests {
self.add_session_tokens(request.session_id, 1);
}
}
result
}
pub fn verify_frame(
&mut self,
session_id: &str,
token_ids: &[i32],
input: Option<&ActivationFrame>,
output_capacity: usize,
) -> Result<(Vec<i32>, Option<NativeMtpDraft>, ActivationFrame)> {
self.verify_frame_sampled(session_id, token_ids, None, input, output_capacity, 0)
}
pub(crate) fn canonical_session_position(&self, session_id: &str) -> Result<u64> {
let tracked_position = self
.session_token_counts
.get(session_id)
.copied()
.with_context(|| format!("session {session_id} has no tracked position"))?;
let session = self
.sessions
.get(session_id)
.with_context(|| format!("session {session_id} is not active"))?;
let rust_position = session.session.token_count();
let native_position = session.session.native_position()?;
if tracked_position != rust_position || tracked_position != native_position {
bail!(
"session {session_id} position mismatch: tracked={tracked_position}, rust={rust_position}, native={native_position}"
);
}
Ok(native_position)
}
pub(crate) fn verify_tokens_sampled(
&mut self,
session_id: &str,
token_ids: &[i32],
sampling: Option<&SamplingConfig>,
) -> Result<Vec<i32>> {
let token_count = u64::try_from(token_ids.len())
.context("linear verification token count exceeds u64")?;
let session = self.session(session_id)?;
let predicted = session.verify_tokens_sampled(token_ids, sampling)?;
self.add_session_tokens(session_id, token_count);
Ok(predicted)
}
pub(crate) fn session_token_count(&self, session_id: &str) -> Option<u64> {
self.session_token_counts.get(session_id).copied()
}
pub fn verify_frame_sampled(
&mut self,
session_id: &str,
token_ids: &[i32],
sampling: Option<&SamplingConfig>,
input: Option<&ActivationFrame>,
output_capacity: usize,
max_draft_tokens: usize,
) -> Result<(Vec<i32>, Option<NativeMtpDraft>, ActivationFrame)> {
let session = self.session(session_id)?;
let output = session.verify_tokens_frame_sampled(
token_ids,
sampling,
input,
output_capacity,
max_draft_tokens,
)?;
self.add_session_tokens(session_id, token_ids.len() as u64);
Ok(output)
}
pub fn verify_frame_sampled_serial(
&mut self,
session_id: &str,
token_ids: &[i32],
sampling: Option<&SamplingConfig>,
input: Option<&ActivationFrame>,
output_capacity: usize,
) -> Result<(Vec<i32>, Option<NativeMtpDraft>, ActivationFrame)> {
if token_ids.is_empty() {
bail!("serial verify_frame requires at least one token");
}
let input_frames = split_activation_frame(input, token_ids.len())?;
let mut predicted_tokens = Vec::with_capacity(token_ids.len());
let mut output_frames = Vec::with_capacity(token_ids.len());
let mut last_draft = None;
for (index, token_id) in token_ids.iter().copied().enumerate() {
let input_frame = input_frames.as_ref().map(|frames| &frames[index]);
let (predicted, native_mtp, output) = self.decode_frame_sampled_mtp(
session_id,
token_id,
sampling,
input_frame,
output_capacity,
1,
)?;
if predicted >= 0 {
predicted_tokens.push(predicted);
}
last_draft = native_mtp;
output_frames.push(output);
}
Ok((
predicted_tokens,
last_draft,
combine_activation_frames(&output_frames)?,
))
}
pub fn retire_verify_checkpoint(
&mut self,
session_id: &str,
token_start: u64,
token_count: u64,
) -> Result<()> {
self.active_session(session_id)?
.retire_verify_checkpoint(token_start, token_count)
}
pub fn trim_session(&mut self, session_id: &str, token_count: u64) -> Result<()> {
let session = self.session(session_id)?;
session.trim_session(token_count)?;
self.session_token_counts
.insert(session_id.to_string(), token_count);
Ok(())
}
pub fn align_session_to_token_count_if_ahead(
&mut self,
session_id: &str,
token_count: u64,
) -> Result<Option<RuntimeSessionAlignStats>> {
let Some(current) = self.session_token_counts.get(session_id).copied() else {
return Ok(None);
};
if current <= token_count {
return Ok(None);
}
self.trim_session(session_id, token_count)?;
Ok(Some(RuntimeSessionAlignStats {
before_token_count: current,
after_token_count: token_count,
}))
}
pub(super) fn session(&mut self, session_id: &str) -> Result<&mut StageSession> {
if !self.sessions.contains_key(session_id) {
let lane_session = self.take_idle_session().map(Ok).unwrap_or_else(|| {
if self.sessions.len() >= self.lane_count as usize {
bail!("all execution lanes are busy");
}
self.create_lane_session()
})?;
self.sessions.insert(session_id.to_string(), lane_session);
}
Ok(&mut self
.sessions
.get_mut(session_id)
.expect("session inserted above")
.session)
}
fn ensure_unique_batch_sessions(requests: &[RuntimeDecodeBatchRequest<'_>]) -> Result<()> {
let mut seen = BTreeSet::new();
for request in requests {
if !seen.insert(request.session_id) {
bail!("duplicate session {} in decode batch", request.session_id);
}
}
Ok(())
}
fn ensure_unique_frame_batch_sessions(
requests: &[RuntimeDecodeFrameBatchRequest<'_>],
) -> Result<()> {
let mut seen = BTreeSet::new();
for request in requests {
if !seen.insert(request.session_id) {
bail!(
"duplicate session {} in decode frame batch",
request.session_id
);
}
}
Ok(())
}
pub(super) fn active_session(&mut self, session_id: &str) -> Result<&mut StageSession> {
self.sessions
.get_mut(session_id)
.map(|lane_session| &mut lane_session.session)
.ok_or_else(|| anyhow::anyhow!("session {session_id} is not active"))
}
}
fn split_activation_frame(
input: Option<&ActivationFrame>,
token_count: usize,
) -> Result<Option<Vec<ActivationFrame>>> {
let Some(input) = input else {
return Ok(None);
};
if token_count == 0 {
bail!("cannot split activation frame for zero tokens");
}
if input.desc.token_count as usize != token_count {
bail!(
"activation token count mismatch: frame={} tokens={}",
input.desc.token_count,
token_count
);
}
if input.payload.len() % token_count != 0 {
bail!(
"activation payload is not divisible by token count: payload={} tokens={}",
input.payload.len(),
token_count
);
}
let row_bytes = input.payload.len() / token_count;
let frames = input
.payload
.chunks(row_bytes)
.map(|row| {
let mut desc = input.desc;
desc.token_count = 1;
desc.sequence_count = 1;
desc.payload_bytes = row.len() as u64;
ActivationFrame {
desc,
payload: row.to_vec(),
}
})
.collect();
Ok(Some(frames))
}
fn combine_activation_frames(frames: &[ActivationFrame]) -> Result<ActivationFrame> {
let Some(first) = frames.first() else {
bail!("cannot combine empty activation frames");
};
let mut desc = first.desc;
let mut payload = Vec::new();
let mut token_count = 0u32;
for frame in frames {
if frame.desc.dtype != desc.dtype
|| frame.desc.layout != desc.layout
|| frame.desc.producer_stage_index != desc.producer_stage_index
|| frame.desc.layer_start != desc.layer_start
|| frame.desc.layer_end != desc.layer_end
|| frame.desc.sequence_count != desc.sequence_count
|| frame.desc.flags != desc.flags
{
bail!("cannot combine incompatible activation frames");
}
token_count = token_count
.checked_add(frame.desc.token_count)
.context("combined activation token count overflow")?;
payload.extend_from_slice(&frame.payload);
}
desc.token_count = token_count;
desc.payload_bytes = payload.len() as u64;
Ok(ActivationFrame { desc, payload })
}

653
crates/skippy-server/src/runtime_state/lane_lifecycle.rs Normal file
View file

@ -0,0 +1,653 @@
use super::*;
impl RuntimeState {
pub fn prewarm_idle_sessions(
&mut self,
target_idle_sessions: usize,
) -> Result<RuntimeSessionStats> {
while self.idle_sessions.len() < target_idle_sessions {
if self.sessions.len() + self.idle_sessions.len() >= self.lane_count as usize {
break;
}
let lane_session = self.create_lane_session()?;
self.idle_sessions.push(lane_session);
}
Ok(self.session_stats())
}
/// Release the session slot identified by `session_id`.
///
/// This is the cleanup path called at the end of every chat
/// completion (success, cancellation, or backend error). It must
/// leave [`Self`] in a self-consistent state regardless of whether
/// the underlying StageSession can be reset cleanly:
///
/// - The lane is either returned to `idle_sessions` (reset OK) or
/// dropped entirely (reset failed). Dropping the lane triggers
/// `StageSession::drop`, which calls `skippy_session_free` on
/// the C side — the authoritative path for releasing native KV
/// cells held by that sequence id.
/// - `session_token_counts` and `session_resident_prefixes` for
/// `session_id` are always removed.
/// - The function always returns `Ok` so per-request cleanup at
/// callsites never propagates a reset failure as a request
/// error. The outcome is reported via [`RuntimeSessionDropStats`]
/// fields (`lane_discarded`, `lane_discard_reason`) for
/// telemetry.
///
/// Previously a reset error propagated `?` through this function,
/// which left `session_token_counts` holding stale entries and dropped the lane on the floor without
/// any record. That accumulated bookkeeping drift over time and
/// could leave the native KV cache reporting "all slots in use"
/// long after the owning sessions were gone, producing
/// `failed to find a memory slot` errors on subsequent admissions.
pub fn drop_session_timed(&mut self, session_id: &str) -> Result<RuntimeSessionDropStats> {
let reset_started = Instant::now();
let mut reset_session = false;
let preserved_resident_prefix = false;
let mut lane_discarded = false;
let mut lane_discard_reason: Option<String> = None;
if let Some(mut lane_session) = self.sessions.remove(session_id) {
let lane_index = lane_session.index;
// Always release the lane's native KV cells back to the
// unified pool. The trim+preserve path kept the lane's cells
// pinned to a specific (`page_id`, `token_count`) pair so a
// future request whose content prefix hashed to the *exact*
// same `page_id` AND same `token_count` could acquire the
// warm lane via `acquire_resident_prefix_lane`. Real chat /
// agent workloads vary the conversation tail every turn, so
// both the hash and the length change request-to-request and
// that exact-match acquisition almost never fires. Meanwhile
// the pinned cells remain claimed in the unified pool, in
// parallel with the cells the cache layer itself pins, and
// the pool runs out of contiguous space — producing
// `decode: failed to find a memory slot` under repeated
// tool-using agent traffic (#652). Cross-request prefix
// reuse is still done by the cache layer (by `page_id`); we
// just stop double-claiming cells on the lane side.
self.session_resident_prefixes.remove(session_id);
reset_session = true;
match lane_session.session.reset() {
Ok(()) => {
lane_session.resident_prefix = None;
self.idle_sessions.push(lane_session);
}
Err(reset_err) => {
lane_discarded = true;
let reason = format!("reset() failed ({reset_err:#})");
eprintln!(
"skippy::runtime_state: drop_session_timed: discarding lane {lane_index} for session {session_id}: {reason}"
);
lane_discard_reason = Some(reason);
drop(lane_session);
self.free_lane_indices.push(lane_index);
}
}
}
// Always clear per-session bookkeeping. The previous version
// skipped these when reset returned Err, which leaked entries.
//
// session_resident_prefixes is also cleared here defensively:
// it's already removed above on the active-session path, but
// calling drop_session_timed for an id that's no longer in
// `sessions` (idempotent cleanup, stale callers) must still
// clear any stray resident-prefix entry under that id.
self.session_token_counts.remove(session_id);
self.session_resident_prefixes.remove(session_id);
Ok(RuntimeSessionDropStats {
reset_session,
reset_ms: reset_started.elapsed().as_secs_f64() * 1000.0,
preserved_resident_prefix,
lane_discarded,
lane_discard_reason,
stats_after: self.session_stats(),
})
}
pub fn session_stats(&self) -> RuntimeSessionStats {
let mut max_session_tokens = 0u64;
let mut total_session_tokens = 0u64;
let mut lanes = (0..self.lane_count as usize)
.map(|index| RuntimeSessionLaneStats {
index,
active: false,
session_id: None,
token_count: None,
})
.collect::<Vec<_>>();
for (session_id, lane_session) in &self.sessions {
if let Some(token_count) = self.session_token_counts.get(session_id).copied() {
max_session_tokens = max_session_tokens.max(token_count);
total_session_tokens = total_session_tokens.saturating_add(token_count);
}
if let Some(lane) = lanes.get_mut(lane_session.index) {
lane.active = true;
lane.session_id = Some(session_id.clone());
lane.token_count = self.session_token_counts.get(session_id).copied();
}
}
RuntimeSessionStats {
lane_count: self.lane_count as usize,
active_sessions: self.sessions.len(),
idle_sessions: self.idle_sessions.len(),
idle_resident_prefixes: self
.idle_sessions
.iter()
.filter(|idle| idle.resident_prefix.is_some())
.count(),
tracked_token_counts: self.session_token_counts.len(),
max_session_tokens,
total_session_tokens,
lanes,
}
}
pub(super) fn take_idle_session(&mut self) -> Option<RuntimeLaneSession> {
if let Some(index) = self
.idle_sessions
.iter()
.position(|idle| idle.resident_prefix.is_none())
{
return Some(self.idle_sessions.swap_remove(index));
}
self.idle_sessions.pop()
}
pub fn retain_resident_prefix_on_drop(
&mut self,
session_id: &str,
page_id: String,
token_count: u64,
) -> Result<()> {
if !self.sessions.contains_key(session_id) {
bail!("session {session_id} does not exist");
}
if self
.session_resident_prefixes
.get(session_id)
.is_some_and(|current| current.token_count >= token_count)
{
return Ok(());
}
self.session_resident_prefixes.insert(
session_id.to_string(),
ResidentLanePrefix {
page_id,
token_count,
},
);
Ok(())
}
pub fn acquire_resident_prefix_lane(
&mut self,
session_id: &str,
page_id: &str,
token_count: u64,
) -> Result<bool> {
if self.sessions.contains_key(session_id) {
bail!("session {session_id} already exists");
}
let Some(index) = self.idle_sessions.iter().position(|idle| {
idle.resident_prefix.as_ref().is_some_and(|prefix| {
prefix.page_id == page_id && prefix.token_count == token_count
})
}) else {
return Ok(false);
};
let mut idle = self.idle_sessions.swap_remove(index);
idle.resident_prefix = None;
self.sessions.insert(session_id.to_string(), idle);
self.session_token_counts
.insert(session_id.to_string(), token_count);
self.session_resident_prefixes.insert(
session_id.to_string(),
ResidentLanePrefix {
page_id: page_id.to_string(),
token_count,
},
);
Ok(true)
}
pub fn has_session_range(&self, session_id: &str, token_start: u64, token_count: u64) -> bool {
let Some(token_end) = token_start.checked_add(token_count) else {
return false;
};
self.session_token_counts
.get(session_id)
.copied()
.is_some_and(|known_tokens| token_end <= known_tokens)
}
#[allow(dead_code)]
pub fn export_kv_page(
&mut self,
session_id: &str,
token_start: u64,
token_count: u64,
) -> Result<RuntimeKvPage> {
self.validate_export_range(session_id, token_start, token_count)?;
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.export_kv_page(layer_start, layer_end, token_start, token_count)
}
#[allow(dead_code)]
pub fn probe_kv_page(
&mut self,
session_id: &str,
token_start: u64,
token_count: u64,
) -> Result<RuntimeKvPageDesc> {
self.validate_export_range(session_id, token_start, token_count)?;
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
let page = session.export_kv_page(layer_start, layer_end, token_start, token_count)?;
Ok(page.desc)
}
pub fn import_kv_page(
&mut self,
session_id: &str,
desc: &RuntimeKvPageDesc,
bytes: &[u8],
) -> Result<()> {
let session = self.session(session_id)?;
session.import_kv_page(desc, bytes)?;
let token_end = desc
.token_start
.checked_add(desc.token_count)
.ok_or_else(|| anyhow::anyhow!("KV page token range overflows"))?;
self.session_token_counts
.entry(session_id.to_string())
.and_modify(|current| *current = (*current).max(token_end))
.or_insert(token_end);
Ok(())
}
#[allow(dead_code)]
pub fn export_state(&mut self, session_id: &str) -> Result<Vec<u8>> {
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.export_state(layer_start, layer_end)
}
pub fn import_state(&mut self, session_id: &str, bytes: &[u8]) -> Result<()> {
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.import_state(layer_start, layer_end, bytes)
}
pub fn import_state_for_token_count(
&mut self,
session_id: &str,
bytes: &[u8],
token_count: u64,
) -> Result<()> {
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.import_state_for_token_count(layer_start, layer_end, bytes, token_count)?;
record_restored_session_token_count(
&mut self.session_token_counts,
session_id,
token_count,
);
Ok(())
}
pub fn export_full_state(&mut self, session_id: &str) -> Result<Vec<u8>> {
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.export_full_state(layer_start, layer_end)
}
pub fn import_full_state(&mut self, session_id: &str, bytes: &[u8]) -> Result<()> {
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.import_full_state(layer_start, layer_end, bytes)
}
pub fn import_full_state_for_token_count(
&mut self,
session_id: &str,
bytes: &[u8],
token_count: u64,
) -> Result<()> {
let layer_start = i32::try_from(self.model_layer_start())?;
let layer_end = i32::try_from(self.model_layer_end())?;
let session = self.session(session_id)?;
session.import_full_state_for_token_count(layer_start, layer_end, bytes, token_count)?;
record_restored_session_token_count(
&mut self.session_token_counts,
session_id,
token_count,
);
Ok(())
}
pub fn export_recurrent_state(&mut self, session_id: &str) -> Result<Vec<u8>> {
self.session(session_id)?.export_recurrent_state()
}
pub fn import_recurrent_state_for_token_count(
&mut self,
session_id: &str,
bytes: &[u8],
token_count: u64,
) -> Result<()> {
self.session(session_id)?
.import_recurrent_state_for_token_count(bytes, token_count)?;
record_restored_session_token_count(
&mut self.session_token_counts,
session_id,
token_count,
);
Ok(())
}
pub fn save_resident_prefix(
&mut self,
session_id: &str,
cache_seq_id: i32,
token_count: u64,
) -> Result<()> {
self.session(session_id)?
.save_prefix(cache_seq_id, token_count)
}
pub fn restore_resident_prefix(
&mut self,
session_id: &str,
cache_seq_id: i32,
token_ids: &[i32],
) -> Result<()> {
let session = self.session(session_id)?;
session.restore_prefix(cache_seq_id, token_ids)?;
self.session_token_counts
.insert(session_id.to_string(), token_ids.len() as u64);
Ok(())
}
pub fn borrow_resident_prefix_session(
&mut self,
session_id: &str,
cache_seq_id: i32,
token_ids: &[i32],
) -> Result<()> {
if self.sessions.contains_key(session_id) {
bail!("session {session_id} already exists");
}
let model = &self.model;
let (index, session) = create_indexed_lane_resource(
&mut self.next_lane_index,
&mut self.free_lane_indices,
self.lane_count,
|| model.create_session_from_resident_prefix(cache_seq_id, token_ids),
)?;
let lane_session = RuntimeLaneSession {
index,
session,
resident_prefix: None,
};
self.sessions.insert(session_id.to_string(), lane_session);
self.session_token_counts
.insert(session_id.to_string(), token_ids.len() as u64);
Ok(())
}
pub fn drop_resident_prefix_sequence(
&mut self,
session_id: &str,
cache_seq_id: i32,
) -> Result<()> {
self.active_session(session_id)?.drop_sequence(cache_seq_id)
}
pub(super) fn add_session_tokens(&mut self, session_id: &str, count: u64) {
self.session_token_counts
.entry(session_id.to_string())
.and_modify(|current| *current = current.saturating_add(count))
.or_insert(count);
}
fn validate_export_range(
&self,
session_id: &str,
token_start: u64,
token_count: u64,
) -> Result<()> {
let token_end = token_start
.checked_add(token_count)
.ok_or_else(|| anyhow::anyhow!("KV page token range overflows"))?;
let known_tokens = self
.session_token_counts
.get(session_id)
.copied()
.unwrap_or_default();
if token_end > known_tokens {
bail!(
"cannot export KV page [{token_start}, {token_end}) from session with {known_tokens} known tokens"
);
}
Ok(())
}
fn model_layer_start(&self) -> u32 {
self.layer_start
}
fn model_layer_end(&self) -> u32 {
self.layer_end
}
pub(super) fn create_lane_session(&mut self) -> Result<RuntimeLaneSession> {
let model = &self.model;
let (index, session) = create_indexed_lane_resource(
&mut self.next_lane_index,
&mut self.free_lane_indices,
self.lane_count,
|| model.create_session(),
)?;
Ok(RuntimeLaneSession {
index,
session,
resident_prefix: None,
})
}
}
fn record_restored_session_token_count(
session_token_counts: &mut BTreeMap<String, u64>,
session_id: &str,
token_count: u64,
) {
// A prefix restore can move an existing lane backwards to a shorter
// common prefix. The tracked position must follow the imported native
// state exactly; retaining the previous high-water mark submits the next
// divergent token at the wrong position and makes llama_decode fail.
session_token_counts.insert(session_id.to_string(), token_count);
}
/// Allocate the next lane slot.
///
/// Prefers indices in `free_lane_indices` (lanes previously discarded
/// via [`RuntimeState::drop_session_timed`]) so they can be reused
/// without growing `next_lane_index` past `lane_count`. If the free
/// list is empty, falls through to bumping `next_lane_index`. If both
/// are exhausted, returns "all execution lanes are busy".
///
/// If `create()` fails after popping from the free list, the index is
/// pushed back so a retry can reuse it. The high-water counter is only
/// bumped on success, matching the prior behavior.
fn create_indexed_lane_resource<T>(
next_lane_index: &mut usize,
free_lane_indices: &mut Vec<usize>,
lane_count: u32,
create: impl FnOnce() -> Result<T>,
) -> Result<(usize, T)> {
if let Some(index) = free_lane_indices.pop() {
let resource = match create() {
Ok(resource) => resource,
Err(err) => {
// Return the freed index so the next allocation can
// still reuse it.
free_lane_indices.push(index);
return Err(err);
}
};
return Ok((index, resource));
}
if *next_lane_index >= lane_count as usize {
bail!("all execution lanes are busy");
}
let index = *next_lane_index;
let resource = create()?;
*next_lane_index = index + 1;
Ok((index, resource))
}
#[cfg(test)]
mod tests {
use super::*;
use anyhow::{Result, bail};
#[test]
fn prefix_restore_moves_tracked_position_backwards() {
let mut token_counts = std::collections::BTreeMap::from([("lane-a".to_string(), 3_535)]);
record_restored_session_token_count(&mut token_counts, "lane-a", 3_530);
assert_eq!(token_counts.get("lane-a"), Some(&3_530));
}
#[test]
fn create_indexed_lane_resource_keeps_index_available_when_creation_fails() {
let mut next_lane_index = 0;
let mut free_lane_indices: Vec<usize> = Vec::new();
let error = create_indexed_lane_resource(
&mut next_lane_index,
&mut free_lane_indices,
2,
|| -> Result<()> { bail!("transient session creation failure") },
)
.expect_err("failed creation should propagate the original error");
assert_eq!(error.to_string(), "transient session creation failure");
assert_eq!(next_lane_index, 0);
assert!(free_lane_indices.is_empty());
let (index, resource) =
create_indexed_lane_resource(&mut next_lane_index, &mut free_lane_indices, 2, || {
Ok("lane")
})
.expect("successful retry should reuse the unconsumed lane index");
assert_eq!(index, 0);
assert_eq!(resource, "lane");
assert_eq!(next_lane_index, 1);
}
#[test]
fn create_indexed_lane_resource_reuses_freed_indices_before_growing() {
// Simulate the wedge scenario: all lanes allocated, one lane
// freed via the discard path, next allocation must reuse the
// freed index rather than bailing with "all execution lanes
// are busy".
let mut next_lane_index = 0;
let mut free_lane_indices: Vec<usize> = Vec::new();
let lane_count = 2;
// Allocate both lanes.
let (a_idx, _) = create_indexed_lane_resource(
&mut next_lane_index,
&mut free_lane_indices,
lane_count,
|| Ok("a"),
)
.expect("first allocation should succeed");
let (b_idx, _) = create_indexed_lane_resource(
&mut next_lane_index,
&mut free_lane_indices,
lane_count,
|| Ok("b"),
)
.expect("second allocation should succeed");
assert_eq!(a_idx, 0);
assert_eq!(b_idx, 1);
assert_eq!(next_lane_index, 2);
// Pool is full at the high-water mark. A third allocation must
// fail.
let error = create_indexed_lane_resource(
&mut next_lane_index,
&mut free_lane_indices,
lane_count,
|| Ok("c"),
)
.expect_err("allocating past lane_count should fail when no slots are free");
assert!(error.to_string().contains("all execution lanes are busy"));
// Discard one lane: the caller pushes its freed index onto the
// free list (this is what drop_session_timed does on the
// discard branch).
free_lane_indices.push(a_idx);
// The next allocation MUST reuse the freed index instead of
// bailing. This is the wedge regression: previously
// next_lane_index stayed at lane_count and every allocation
// failed forever.
let (reused_idx, _) = create_indexed_lane_resource(
&mut next_lane_index,
&mut free_lane_indices,
lane_count,
|| Ok("c"),
)
.expect("allocation must reuse a freed index, not stay wedged");
assert_eq!(reused_idx, 0);
assert_eq!(next_lane_index, 2);
assert!(free_lane_indices.is_empty());
}
#[test]
fn create_indexed_lane_resource_returns_freed_index_on_create_failure() {
// If create() fails while consuming a freed index, the index
// must go back onto the free list so a retry can use it.
let mut next_lane_index = 1;
let mut free_lane_indices: Vec<usize> = vec![0];
let error = create_indexed_lane_resource(
&mut next_lane_index,
&mut free_lane_indices,
2,
|| -> Result<()> { bail!("create failed mid-reuse") },
)
.expect_err("failed creation should propagate");
assert_eq!(error.to_string(), "create failed mid-reuse");
assert_eq!(next_lane_index, 1);
assert_eq!(free_lane_indices, vec![0]);
// A retry should now succeed using the same freed index.
let (idx, _) =
create_indexed_lane_resource(&mut next_lane_index, &mut free_lane_indices, 2, || {
Ok("retry")
})
.expect("retry should succeed");
assert_eq!(idx, 0);
assert_eq!(next_lane_index, 1);
assert!(free_lane_indices.is_empty());
}
}

37
third_party/llama.cpp/patches/0048-Preserve-grammar-state-during-sampled-verification.patch vendored Normal file
View file

@ -0,0 +1,37 @@
From caee7b5ada3a05963938f2903093c41cc79f3c96 Mon Sep 17 00:00:00 2001
From: Mesh-LLM CI <ci@mesh-llm.local>
Date: Sun, 2 Aug 2026 14:41:51 +1000
Subject: [PATCH] Preserve grammar state during sampled verification
---
src/skippy.cpp | 10 +++++++---
1 file changed, 7 insertions(+), 3 deletions(-)
diff --git a/src/skippy.cpp b/src/skippy.cpp
index 21db2b77d..38ea48649 100644
--- a/src/skippy.cpp
+++ b/src/skippy.cpp
@@ -5966,9 +5966,11 @@ enum skippy_status skippy_verify_tokens_frame_sampled(
}
if (session->stage_model->config.include_output) {
+ session->token_history.resize(session->token_history.size() - token_count);
- const int32_t n_tokens = static_cast<int32_t>(token_count);
- for (int32_t i = 0; i < n_tokens; ++i) {
- output_tokens[i] = skippy_sample_token_ith(session, sampling, i);
+ for (size_t i = 0; i < token_count; ++i) {
+ skippy_record_tokens(session, &token_ids[i], 1);
+ const int32_t logits_index = static_cast<int32_t>(i);
+ output_tokens[i] = skippy_sample_token_ith(session, sampling, logits_index);
}
if (out_mtp_draft != nullptr && token_count > 0) {
status = skippy_mtp_propose_next(
@@ -5981,6 +5983,8 @@ enum skippy_status skippy_verify_tokens_frame_sampled(
return status;
}
}
+ } else {
+ skippy_record_tokens(session, token_ids, token_count);
}
return skippy_success(out_error);

280
third_party/llama.cpp/patches/0049-Accept-tagged-tool-arguments-in-object-order.patch vendored Normal file
View file

@ -0,0 +1,280 @@
From ad70b08a7c192e09f766694a2379da69c37d7e81 Mon Sep 17 00:00:00 2001
From: Mesh-LLM CI <ci@mesh-llm.local>
Date: Sun, 2 Aug 2026 14:42:14 +1000
Subject: [PATCH] Accept tagged tool arguments in object order
---
common/chat-auto-parser-generator.cpp | 37 +++++---------
common/chat-peg-parser.cpp | 19 +++++++-
common/chat-peg-parser.h | 11 +++++
tests/test-chat.cpp | 70 ++++++++++++++++++++++++++-
4 files changed, 111 insertions(+), 26 deletions(-)
diff --git a/common/chat-auto-parser-generator.cpp b/common/chat-auto-parser-generator.cpp
index a68e2c916..98db861e3 100644
--- a/common/chat-auto-parser-generator.cpp
+++ b/common/chat-auto-parser-generator.cpp
@@ -409,12 +409,11 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
auto schema_info = common_schema_info();
schema_info.resolve_refs(params);
- // Build parser for each argument, separating required and optional
- std::vector<common_peg_parser> required_parsers;
- std::vector<common_peg_parser> optional_parsers;
+ // Object member order is semantically irrelevant. Build one choice of
+ // every declared argument and retain required-field metadata for the
+ // mapper to validate when the complete tool call closes.
+ std::vector<common_peg_parser> arg_parsers;
for (const auto & [param_name, param_schema] : properties.items()) {
- bool is_required = required.find(param_name) != required.end();
-
auto arg =
p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param_name)) +
arguments.name_suffix) +
@@ -427,29 +426,19 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
p.tool_arg_close(p.literal(arguments.value_suffix)))));
auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg);
- if (is_required) {
- required_parsers.push_back(named_arg);
- } else {
- optional_parsers.push_back(named_arg);
- }
+ arg_parsers.push_back(named_arg);
}
- // Build required arg sequence in definition order
common_peg_parser args_seq = p.eps();
- for (size_t i = 0; i < required_parsers.size(); i++) {
- if (i > 0) {
- args_seq = args_seq + p.space();
- }
- args_seq = args_seq + required_parsers[i];
+ for (const auto & required_arg : required) {
+ args_seq = args_seq + p.tool_required_arg(required_arg);
}
-
- // Build optional args with flexible ordering
- if (!optional_parsers.empty()) {
- common_peg_parser any_opt = p.choice();
- for (const auto & opt : optional_parsers) {
- any_opt |= opt;
+ if (!arg_parsers.empty()) {
+ common_peg_parser any_arg = p.choice();
+ for (const auto & arg : arg_parsers) {
+ any_arg |= arg;
}
- args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1);
+ args_seq = args_seq + p.zero_or_more(any_arg + p.space());
}
if (!arguments.start.empty()) {
@@ -474,7 +463,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
// Only peek for an arg tag when there are required args that must follow.
// When all args are optional, the model may emit no arg tags at all (#20650).
- auto atomic_peek = (!arguments.name_prefix.empty() && !required_parsers.empty()) ?
+ auto atomic_peek = (!arguments.name_prefix.empty() && !required.empty()) ?
std::optional(p.peek(p.literal(arguments.name_prefix))) : std::nullopt;
auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_seq, atomic_peek);
tool_choice |= p.rule("tool-" + name, func_parser);
diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp
index a309f0276..f2c42c417 100644
--- a/common/chat-peg-parser.cpp
+++ b/common/chat-peg-parser.cpp
@@ -300,6 +300,7 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
bool is_arg_name = node.tag == common_chat_peg_builder::TOOL_ARG_NAME;
bool is_arg_value = node.tag == common_chat_peg_builder::TOOL_ARG_VALUE;
bool is_arg_string_value = node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE;
+ bool is_required_arg = node.tag.rfind(common_chat_peg_builder::TOOL_REQUIRED_ARG_PREFIX, 0) == 0;
if (is_tool_open) {
pending_tool_call = common_chat_tool_call();
@@ -307,6 +308,13 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
arg_count = 0;
args_buffer.clear();
closing_quote_pending = false;
+ required_args.clear();
+ seen_args.clear();
+ }
+
+ if (is_required_arg && current_tool) {
+ required_args.insert(node.tag.substr(std::char_traits<char>::length(
+ common_chat_peg_builder::TOOL_REQUIRED_ARG_PREFIX)));
}
if (is_tool_id && current_tool) {
@@ -348,11 +356,15 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
}
if (is_arg_name && current_tool) {
+ const std::string arg_name(trim(node.text));
+ if (!seen_args.insert(arg_name).second) {
+ throw std::runtime_error("Duplicate tool argument: " + arg_name);
+ }
std::string arg_entry;
if (arg_count > 0) {
arg_entry = ",";
}
- arg_entry += ordered_json(trim(node.text)).dump() + ":";
+ arg_entry += ordered_json(arg_name).dump() + ":";
++arg_count;
auto & target = args_target();
@@ -393,6 +405,11 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
}
if (is_tool_close && current_tool) {
+ for (const auto & required_arg : required_args) {
+ if (seen_args.find(required_arg) == seen_args.end()) {
+ throw std::runtime_error("Missing required tool argument: " + required_arg);
+ }
+ }
// Flush buffer to arguments if tool name was never seen
if (current_tool->name.empty() && !args_buffer.empty()) {
current_tool->arguments = args_buffer;
diff --git a/common/chat-peg-parser.h b/common/chat-peg-parser.h
index b3ffd7de2..203ac0107 100644
--- a/common/chat-peg-parser.h
+++ b/common/chat-peg-parser.h
@@ -5,6 +5,7 @@
#include <map>
#include <optional>
+#include <set>
#include <vector>
class common_chat_peg_mapper {
@@ -26,6 +27,8 @@ class common_chat_peg_mapper {
int arg_count = 0;
bool closing_quote_pending = false;
std::string args_buffer; // Buffer to delay arguments until tool name is known
+ std::set<std::string> required_args;
+ std::set<std::string> seen_args;
// Returns a reference to the active argument destination string.
// Before tool_name is known, writes go to args_buffer; after, to current_tool->arguments.
@@ -63,6 +66,7 @@ class common_chat_peg_builder : public common_peg_parser_builder {
static constexpr const char * TOOL_ARG_NAME = "tool-arg-name";
static constexpr const char * TOOL_ARG_VALUE = "tool-arg-value";
static constexpr const char * TOOL_ARG_STRING_VALUE = "tool-arg-string-value"; // For schema-declared string types
+ static constexpr const char * TOOL_REQUIRED_ARG_PREFIX = "tool-required-arg:";
// Low-level tag methods (from former common_chat_peg_base_builder)
common_peg_parser reasoning_block(const common_peg_parser & p) { return tag(REASONING_BLOCK, p); }
@@ -92,6 +96,13 @@ class common_chat_peg_builder : public common_peg_parser_builder {
common_peg_parser tool_arg_string_value(const common_peg_parser & p) { return tag(TOOL_ARG_STRING_VALUE, p); }
common_peg_parser tool_arg_json_value(const common_peg_parser & p) { return tag(TOOL_ARG_VALUE, p); }
+ // Attach schema metadata to the AST without consuming model output. The
+ // mapper uses this to enforce required tagged arguments after accepting
+ // them in arbitrary object-key order.
+ common_peg_parser tool_required_arg(const std::string & name) {
+ return tag(std::string(TOOL_REQUIRED_ARG_PREFIX) + name, eps());
+ }
+
// Return a parser that parses the prefix of a string, up to a given delimiter.
common_peg_parser prefix(const std::string & s, const std::string & delimiter = {});
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index 4dd00efdd..431d80dc7 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -1388,6 +1388,8 @@ class peg_tester {
const std::string & template_path() const { return template_path_; }
+ common_chat_templates * templates() const { return tmpls_.get(); }
+
peg_test_builder test(const std::string & input);
};
@@ -4181,6 +4183,15 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
{
auto tst = peg_tester("models/templates/GLM-4.7-Flash.jinja", detailed_debug);
+ static const common_chat_tool terminal_tool{
+ "terminal", "Run or interact with a terminal command",
+ R"({"type":"object","properties":{"security_risk":{"type":"string"},"summary":{"type":"string"},"command":{"type":"string"},"is_input":{"type":"boolean"},"timeout":{"type":"integer"},"reset":{"type":"boolean"}},"required":["command","security_risk"]})",
+ };
+ static const common_chat_tool think_tool{
+ "think", "Record reasoning",
+ R"({"type":"object","properties":{"summary":{"type":"string"},"thought":{"type":"string"}},"required":["thought"]})",
+ };
+
// Pure content (no reasoning)
tst.test("Hello, world!\nWhat's up?")
.enable_thinking(false)
@@ -4207,6 +4218,63 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
.expect_reconstruction()
.run();
+ // OpenAI tool arguments are object members: required and optional
+ // fields may be interleaved in any order. These are exact ordering
+ // shapes observed in a multi-tool agent regression case.
+ tst.test(
+ "<tool_call>terminal"
+ "<arg_key>command</arg_key><arg_value>C-c</arg_value>"
+ "<arg_key>is_input</arg_key><arg_value>true</arg_value>"
+ "<arg_key>security_risk</arg_key><arg_value>LOW</arg_value>"
+ "</tool_call>")
+ .enable_thinking(false)
+ .tools({ terminal_tool, think_tool })
+ .expect_tool_calls({
+ { "terminal", R"({"command":"C-c","is_input":true,"security_risk":"LOW"})", {} },
+ })
+ .run();
+
+ tst.test(
+ "<tool_call>think"
+ "<arg_key>summary</arg_key><arg_value>Inspect the failure</arg_value>"
+ "<arg_key>thought</arg_key><arg_value>Check the process state.</arg_value>"
+ "</tool_call>")
+ .enable_thinking(false)
+ .tools({ terminal_tool, think_tool })
+ .expect_tool_calls({
+ { "think", R"({"summary":"Inspect the failure","thought":"Check the process state."})", {} },
+ })
+ .run();
+
+ // Flexible ordering must not weaken the schema's presence and
+ // uniqueness rules.
+ common_chat_templates_inputs validation_inputs;
+ validation_inputs.messages = { message_user };
+ validation_inputs.tools = { terminal_tool, think_tool };
+ validation_inputs.enable_thinking = false;
+ auto validation_parser = make_peg_parser(tst.templates(), validation_inputs, detailed_debug);
+
+ try {
+ validation_parser.parse(
+ "<tool_call>terminal"
+ "<arg_key>command</arg_key><arg_value>pwd</arg_value>"
+ "</tool_call>", false);
+ throw std::runtime_error("Expected missing required tagged argument to fail");
+ } catch (const std::exception & e) {
+ assert_contains(e.what(), "Missing required tool argument: security_risk");
+ }
+
+ try {
+ validation_parser.parse(
+ "<tool_call>think"
+ "<arg_key>thought</arg_key><arg_value>one</arg_value>"
+ "<arg_key>thought</arg_key><arg_value>two</arg_value>"
+ "</tool_call>", false);
+ throw std::runtime_error("Expected duplicate tagged argument to fail");
+ } catch (const std::exception & e) {
+ assert_contains(e.what(), "Duplicate tool argument: thought");
+ }
+
// Tool call with reasoning (forced-open mode)
tst.test(
"I'm\nthinking</think>"
@@ -4318,7 +4386,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
"Thinking.\n"
"</think>"
"<tool_call>get_weather"
- "<arg_key>city</arg_key><arg_value>Tokyo</arg_value>"
+ "<arg_key>country</arg_key><arg_value>Japan</arg_value>"
"</tool_call>\n";
bool got_runtime_error = false;

102
third_party/llama.cpp/patches/0050-Preserve-OpenAI-message-content-presence.patch vendored Normal file
View file

@ -0,0 +1,102 @@
From 551993765b392414535d7359d841b7b1ab614779 Mon Sep 17 00:00:00 2001
From: Mesh-LLM CI <ci@mesh-llm.local>
Date: Sun, 2 Aug 2026 14:42:39 +1000
Subject: [PATCH] Preserve OpenAI message content presence
---
common/chat.cpp | 10 +++++++++-
common/chat.h | 5 ++++-
tests/test-chat.cpp | 22 ++++++++++++++++++++++
3 files changed, 35 insertions(+), 2 deletions(-)
diff --git a/common/chat.cpp b/common/chat.cpp
index 91658e6df..5498e676f 100644
--- a/common/chat.cpp
+++ b/common/chat.cpp
@@ -226,6 +226,11 @@ json common_chat_msg::to_json_oaicompat(bool concat_typed_text) const {
});
}
}
+ } else if (!content_present) {
+ // Preserve an omitted OpenAI-compatible content field. Some templates
+ // distinguish this from explicit null and an empty string.
+ } else if (content_is_null) {
+ jmsg["content"] = nullptr;
} else {
jmsg["content"] = "";
}
@@ -390,6 +395,7 @@ std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const json & messa
auto has_content = message.contains("content");
auto has_tool_calls = message.contains("tool_calls");
+ msg.content_present = has_content;
if (has_content) {
const auto & content = message.at("content");
if (content.is_string()) {
@@ -408,7 +414,9 @@ std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const json & messa
msg_part.text = part.at("text");
msg.content_parts.push_back(msg_part);
}
- } else if (!content.is_null()) {
+ } else if (content.is_null()) {
+ msg.content_is_null = true;
+ } else {
throw std::invalid_argument("Invalid 'content' type: expected string or array, got " +
content.dump() +
" (ref: https://github.com/ggml-org/llama.cpp/issues/8367)");
diff --git a/common/chat.h b/common/chat.h
index d79f4ecd7..b3c846aff 100644
--- a/common/chat.h
+++ b/common/chat.h
@@ -86,6 +86,8 @@ struct common_chat_msg {
std::string reasoning_content;
std::string tool_name;
std::string tool_call_id;
+ bool content_present = true;
+ bool content_is_null = false;
nlohmann::ordered_json to_json_oaicompat(bool concat_typed_text = false) const;
@@ -122,7 +124,8 @@ struct common_chat_msg {
bool operator==(const common_chat_msg & other) const {
return role == other.role && content == other.content && content_parts == other.content_parts &&
tool_calls == other.tool_calls && reasoning_content == other.reasoning_content &&
- tool_name == other.tool_name && tool_call_id == other.tool_call_id;
+ tool_name == other.tool_name && tool_call_id == other.tool_call_id &&
+ content_present == other.content_present && content_is_null == other.content_is_null;
}
bool operator!=(const common_chat_msg & other) const { return !(*this == other); }
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index 431d80dc7..4c30d33ad 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -1564,6 +1564,28 @@ static void test_msgs_oaicompat_json_conversion() {
assert_equals<std::string>(res[0].role, "assistant");
assert_equals(true, res[0].content.empty());
assert_equals(true, res[0].tool_calls.empty());
+ assert_equals(false, res[0].content_present);
+ assert_equals(
+ std::string("[{\"role\":\"assistant\"}]"),
+ common_chat_msgs_to_json_oaicompat(res).dump());
+
+ auto null_content =
+ common_chat_msgs_parse_oaicompat(json::parse("[{\"role\":\"assistant\",\"content\":null,\"tool_calls\":[]}]"));
+ assert_equals<size_t>(1, null_content.size());
+ assert_equals(true, null_content[0].content_present);
+ assert_equals(true, null_content[0].content_is_null);
+ assert_equals(
+ std::string("[{\"role\":\"assistant\",\"content\":null}]"),
+ common_chat_msgs_to_json_oaicompat(null_content).dump());
+
+ auto empty_content =
+ common_chat_msgs_parse_oaicompat(json::parse("[{\"role\":\"assistant\",\"content\":\"\",\"tool_calls\":[]}]"));
+ assert_equals<size_t>(1, empty_content.size());
+ assert_equals(true, empty_content[0].content_present);
+ assert_equals(false, empty_content[0].content_is_null);
+ assert_equals(
+ std::string("[{\"role\":\"assistant\",\"content\":\"\"}]"),
+ common_chat_msgs_to_json_oaicompat(empty_content).dump());
try {
common_chat_msgs_parse_oaicompat(json::parse("[{\"role\": \"assistant\"}]"));

39
third_party/llama.cpp/patches/0065-Fix-role-only-OpenAI-chat-test-expectation.patch vendored Normal file
View file

@ -0,0 +1,39 @@
From de3a8bb4306fbb386f178287636dbfaf75f4605e Mon Sep 17 00:00:00 2001
From: Mesh-LLM CI <ci@mesh-llm.local>
Date: Tue, 4 Aug 2026 12:13:45 +1000
Subject: [PATCH] Fix role-only OpenAI chat test expectation
---
tests/test-chat.cpp | 9 ++++++++--
1 file changed, 1 insertion(+), 9 deletions(-)
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index 462020186..c9edf8953 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -1642,7 +1642,7 @@ static void test_msgs_oaicompat_json_conversion() {
- auto res = common_chat_msgs_parse_oaicompat(json::parse("[{\"role\": \"assistant\", \"tool_calls\": []}]"));
+ auto res = common_chat_msgs_parse_oaicompat(json::parse("[{\"role\": \"assistant\"}]"));
assert_equals<size_t>(1, res.size());
assert_equals<std::string>(res[0].role, "assistant");
assert_equals(true, res[0].content.empty());
assert_equals(true, res[0].tool_calls.empty());
assert_equals(false, res[0].content_present);
assert_equals(
@@ -1667,14 +1667,6 @@ static void test_msgs_oaicompat_json_conversion() {
std::string("[{\"role\":\"assistant\",\"content\":\"\"}]"),
common_chat_msgs_to_json_oaicompat(empty_content).dump());
- try {
- common_chat_msgs_parse_oaicompat(json::parse("[{\"role\": \"assistant\"}]"));
- throw std::runtime_error("Expected exception");
- } catch (const std::exception & e) {
- if (std::string(e.what()).find("'content'") == std::string::npos) {
- throw std::runtime_error("Expected exception about missing 'content'");
- }
- }
}
static void test_msg_token_delimiters_split() {
--
2.54.0 (Apple Git-157)