harn_vm/llm/
api.rs

1//! LLM API entry points and re-exports. The transport layer, request/
2//! response parsing, auth, context-window discovery, and option/result
3//! types each live in their own submodule under [`self`]; this file only
4//! wires them together and hosts the `vm_call_llm_full*` chat entry
5//! points that provider-specific completion / agent paths dispatch into.
6
7mod auth;
8mod completion;
9mod context_window;
10mod errors;
11mod ollama;
12mod openai_normalize;
13pub(crate) mod options;
14mod partial_tool_args;
15mod response;
16mod result;
17mod schema_stream;
18mod telemetry;
19mod thinking;
20mod transport;
21
22use crate::value::{ErrorCategory, VmError, VmValue};
23
24use super::mock::{
25    fixture_hash, get_replay_mode, load_fixture, mock_llm_response, record_cli_llm_result,
26    save_fixture, LlmReplayMode,
27};
28
29// ─── Public surface (crate-wide) ────────────────────────────────────────
30
31pub(crate) use auth::apply_auth_headers;
32pub(crate) use completion::vm_call_completion_full;
33pub use context_window::fetch_provider_max_context;
34pub(crate) use errors::{
35    classify_llm_error, classify_provider_http_error, retry_after_header, LlmErrorInfo,
36    LlmErrorKind, LlmErrorReason,
37};
38pub(crate) use ollama::apply_ollama_runtime_settings;
39pub(crate) use ollama::ollama_unload_grace_duration_from_env;
40pub use ollama::{
41    normalize_ollama_keep_alive, ollama_readiness, ollama_runtime_settings_from_env,
42    warm_ollama_model, warm_ollama_model_with_settings, OllamaReadinessOptions,
43    OllamaReadinessResult, OllamaRuntimeSettings, OllamaWarmupResult, HARN_OLLAMA_KEEP_ALIVE_ENV,
44    HARN_OLLAMA_NUM_CTX_ENV, OLLAMA_DEFAULT_KEEP_ALIVE, OLLAMA_DEFAULT_NUM_CTX, OLLAMA_HOST_ENV,
45};
46pub(crate) use openai_normalize::normalize_openai_style_messages;
47pub(crate) use options::{
48    push_unique_anthropic_beta_feature, DeltaSender, LlmApiMode, LlmCallOptions, LlmRequestPayload,
49    LlmRouteAlternative, LlmRouteFallback, LlmRoutePolicy, LlmRoutingDecision, OutputFormat,
50    ReasoningEffort, ReminderLifecycleEmission, ThinkingConfig, ToolSearchConfig, ToolSearchMode,
51    ToolSearchVariant,
52};
53pub(crate) use response::parse_llm_response as parse_llm_response_for_provider;
54pub(crate) use response::parse_openai_responses_response;
55pub(crate) use result::{vm_build_llm_result, LlmResult};
56pub(crate) use schema_stream::{
57    aborted_result_value as schema_stream_aborted_result_value, parse_schema_stream_abort,
58    SchemaStreamAbort, StreamSchemaWatch,
59};
60pub(crate) use telemetry::elapsed_ms;
61pub use telemetry::{source as telemetry_source, OllamaPsModel, ProviderTelemetry};
62pub(crate) use thinking::{split_openai_thinking_blocks, ThinkingStreamSplitter};
63pub(crate) use transport::vm_call_llm_api_with_body;
64
65use transport::vm_call_llm_api;
66
67#[derive(Debug, Clone)]
68struct OffthreadLlmError {
69    message: String,
70    category: Option<ErrorCategory>,
71}
72
73impl OffthreadLlmError {
74    fn from_vm_error(err: VmError) -> Self {
75        match err {
76            VmError::CategorizedError { message, category } => Self {
77                message,
78                category: Some(category),
79            },
80            VmError::Thrown(VmValue::String(message)) => {
81                Self::from_display_message(message.to_string())
82            }
83            other => Self::from_display_message(other.to_string()),
84        }
85    }
86
87    fn from_display_message(message: String) -> Self {
88        if let Some((category, stripped)) = parse_displayed_categorized_error(&message) {
89            return Self {
90                message: stripped.to_string(),
91                category: Some(category),
92            };
93        }
94        Self {
95            message,
96            category: None,
97        }
98    }
99
100    fn into_vm_error(self) -> VmError {
101        match self.category {
102            Some(category) => VmError::CategorizedError {
103                message: self.message,
104                category,
105            },
106            None => VmError::Thrown(VmValue::String(arcstr::ArcStr::from(self.message))),
107        }
108    }
109}
110
111fn parse_displayed_categorized_error(message: &str) -> Option<(ErrorCategory, &str)> {
112    let body = message.strip_prefix("Error [")?;
113    let (category, rest) = body.split_once("]: ")?;
114    Some((ErrorCategory::parse(category), rest))
115}
116
117/// Execute an LLM call. Always goes through the streaming path with a
118/// discarding receiver so all callers share one code path for status/error
119/// handling; non-streaming callers just drop the receiver.
120pub(crate) async fn vm_call_llm_full(opts: &LlmCallOptions) -> Result<LlmResult, VmError> {
121    super::cost::check_llm_preflight_budget(opts)?;
122    let (delta_tx, mut delta_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
123    let mut first_token = super::first_token::FirstTokenTimer::for_current_span();
124    let mut deltas_open = true;
125    let mut call = Box::pin(vm_call_llm_full_inner(opts, Some(delta_tx)));
126    let result = loop {
127        tokio::select! {
128            maybe_delta = delta_rx.recv(), if deltas_open => {
129                match maybe_delta {
130                    Some(_) => first_token.observe_delta(),
131                    None => deltas_open = false,
132                }
133            }
134            result = &mut call => break result?,
135        }
136    };
137    while delta_rx.try_recv().is_ok() {
138        first_token.observe_delta();
139    }
140    super::cost::record_llm_usage_for_provider(
141        &result.provider,
142        &result.model,
143        result.input_tokens,
144        result.output_tokens,
145        result.served_fast,
146    )?;
147    Ok(result)
148}
149
150/// Execute an LLM call, streaming text deltas to `delta_tx`.
151pub(crate) async fn vm_call_llm_full_streaming(
152    opts: &LlmCallOptions,
153    delta_tx: DeltaSender,
154) -> Result<LlmResult, VmError> {
155    super::cost::check_llm_preflight_budget(opts)?;
156    let result = vm_call_llm_full_inner(opts, Some(delta_tx)).await?;
157    super::cost::record_llm_usage_for_provider(
158        &result.provider,
159        &result.model,
160        result.input_tokens,
161        result.output_tokens,
162        result.served_fast,
163    )?;
164    Ok(result)
165}
166
167/// Execute provider I/O on Tokio's multithreaded scheduler while keeping
168/// VM-local values and transcript assembly on the caller's LocalSet.
169pub(crate) async fn vm_call_llm_full_streaming_offthread(
170    opts: &LlmCallOptions,
171    delta_tx: DeltaSender,
172) -> Result<LlmResult, VmError> {
173    super::cost::check_llm_preflight_budget(opts)?;
174    let request = LlmRequestPayload::from(opts);
175    let cached = super::trigger_predicate::lookup_cached_result(&request).is_some();
176    let intercepted = crate::llm::providers::MockProvider::should_intercept_request(&request)
177        || crate::llm::fake::FakeLlmProvider::should_intercept(&request.provider);
178    let replay_mode = get_replay_mode();
179    if !cached && !intercepted && replay_mode == LlmReplayMode::Replay {
180        let hash = fixture_hash(&request.model, &request.messages, request.system.as_deref());
181        if load_fixture(&hash).is_none() {
182            return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
183                format!("No fixture found for LLM call (hash: {hash}). Run with --record first."),
184            ))));
185        }
186    }
187    if !cached && !intercepted && replay_mode != LlmReplayMode::Replay {
188        super::ensure_real_llm_allowed(&request.provider)?;
189    }
190    request.emit_reminder_lifecycle();
191    let result = tokio::task::spawn(async move {
192        vm_call_llm_full_inner_offthread(&request, Some(delta_tx)).await
193    })
194    .await
195    .map_err(|join_err| {
196        VmError::Thrown(VmValue::String(arcstr::ArcStr::from(format!(
197            "llm_call background task failed: {join_err}"
198        ))))
199    })?
200    .map_err(OffthreadLlmError::into_vm_error)?;
201    super::cost::record_llm_usage_for_provider(
202        &result.provider,
203        &result.model,
204        result.input_tokens,
205        result.output_tokens,
206        result.served_fast,
207    )?;
208    Ok(result)
209}
210
211async fn vm_call_llm_full_inner(
212    opts: &LlmCallOptions,
213    delta_tx: Option<DeltaSender>,
214) -> Result<LlmResult, VmError> {
215    let request = LlmRequestPayload::from(opts);
216    vm_call_llm_full_inner_request(&request, delta_tx).await
217}
218
219async fn vm_call_llm_full_inner_request(
220    request: &LlmRequestPayload,
221    delta_tx: Option<DeltaSender>,
222) -> Result<LlmResult, VmError> {
223    if let Some(result) = super::trigger_predicate::lookup_cached_result(request) {
224        request.emit_reminder_lifecycle();
225        record_cli_llm_result(request, &result);
226        if let Some(tx) = delta_tx {
227            if !result.text.is_empty() {
228                let _ = tx.send(result.text.clone());
229            }
230        }
231        return Ok(result);
232    }
233
234    if crate::llm::providers::MockProvider::should_intercept_request(request) {
235        request.emit_reminder_lifecycle();
236        let result = mock_llm_response(request)?;
237        super::trigger_predicate::note_result(request, &result);
238        record_cli_llm_result(request, &result);
239        if let Some(tx) = delta_tx {
240            // A mock may script an ordered chunk sequence to emulate a real
241            // token stream; otherwise fall back to a single full-text delta so
242            // streaming callers still see the visible text (the graceful
243            // non-streaming path). `stream_chunks.concat() == result.text`.
244            if let Some(chunks) = super::mock::take_mock_stream_chunks() {
245                for chunk in chunks {
246                    let _ = tx.send(chunk);
247                }
248                return Ok(result);
249            }
250            if !result.text.is_empty() {
251                let _ = tx.send(result.text.clone());
252            }
253            return Ok(result);
254        }
255        return Ok(result);
256    }
257
258    if crate::llm::fake::FakeLlmProvider::should_intercept(&request.provider) {
259        // Bypass fixture/replay so the script-driven fake never collides
260        // with HARN_LLM_REPLAY/RECORD being set from an outer harness.
261        request.emit_reminder_lifecycle();
262        let result = crate::llm::fake::FakeLlmProvider
263            .chat_impl(request, delta_tx)
264            .await?;
265        super::trigger_predicate::note_result(request, &result);
266        record_cli_llm_result(request, &result);
267        return Ok(result);
268    }
269
270    let replay_mode = get_replay_mode();
271    let hash = fixture_hash(&request.model, &request.messages, request.system.as_deref());
272
273    if replay_mode == LlmReplayMode::Replay {
274        if let Some(result) = load_fixture(&hash) {
275            request.emit_reminder_lifecycle();
276            super::trigger_predicate::note_result(request, &result);
277            return Ok(result);
278        }
279        return Err(VmError::Thrown(VmValue::String(arcstr::ArcStr::from(
280            format!("No fixture found for LLM call (hash: {hash}). Run with --record first."),
281        ))));
282    }
283
284    super::ensure_real_llm_allowed(&request.provider)?;
285    request.emit_reminder_lifecycle();
286
287    let result = match vm_call_llm_api(request, delta_tx).await {
288        Ok(result) => result,
289        Err(primary_error) => match try_fallback_provider(request).await {
290            Some(result) => result,
291            None => return Err(primary_error),
292        },
293    };
294
295    if replay_mode == LlmReplayMode::Record {
296        save_fixture(&hash, &result);
297    }
298    super::trigger_predicate::note_result(request, &result);
299    record_cli_llm_result(request, &result);
300
301    Ok(result)
302}
303
304async fn vm_call_llm_full_inner_offthread(
305    request: &LlmRequestPayload,
306    delta_tx: Option<DeltaSender>,
307) -> Result<LlmResult, OffthreadLlmError> {
308    if let Some(result) = super::trigger_predicate::lookup_cached_result(request) {
309        record_cli_llm_result(request, &result);
310        return Ok(result);
311    }
312
313    if crate::llm::providers::MockProvider::should_intercept_request(request) {
314        let result = mock_llm_response(request).map_err(OffthreadLlmError::from_vm_error)?;
315        super::trigger_predicate::note_result(request, &result);
316        record_cli_llm_result(request, &result);
317        return Ok(result);
318    }
319
320    if crate::llm::fake::FakeLlmProvider::should_intercept(&request.provider) {
321        let result = crate::llm::fake::FakeLlmProvider
322            .chat_impl(request, delta_tx)
323            .await
324            .map_err(OffthreadLlmError::from_vm_error)?;
325        super::trigger_predicate::note_result(request, &result);
326        record_cli_llm_result(request, &result);
327        return Ok(result);
328    }
329
330    let replay_mode = get_replay_mode();
331    let hash = fixture_hash(&request.model, &request.messages, request.system.as_deref());
332
333    if replay_mode == LlmReplayMode::Replay {
334        return load_fixture(&hash)
335            .inspect(|result| {
336                super::trigger_predicate::note_result(request, result);
337            })
338            .ok_or_else(|| {
339                OffthreadLlmError::from_display_message(format!(
340                    "No fixture found for LLM call (hash: {hash}). Run with --record first."
341                ))
342            });
343    }
344
345    super::ensure_real_llm_allowed(&request.provider).map_err(OffthreadLlmError::from_vm_error)?;
346
347    let primary_error = match vm_call_llm_api(request, delta_tx).await {
348        Ok(result) => {
349            if replay_mode == LlmReplayMode::Record {
350                save_fixture(&hash, &result);
351            }
352            super::trigger_predicate::note_result(request, &result);
353            record_cli_llm_result(request, &result);
354            return Ok(result);
355        }
356        Err(primary_error) => OffthreadLlmError::from_vm_error(primary_error),
357    };
358
359    let result = match try_fallback_provider(request).await {
360        Some(result) => result,
361        None => return Err(primary_error),
362    };
363
364    if replay_mode == LlmReplayMode::Record {
365        save_fixture(&hash, &result);
366    }
367    super::trigger_predicate::note_result(request, &result);
368    record_cli_llm_result(request, &result);
369
370    Ok(result)
371}
372
373/// Attempt the request on the configured fallback provider.
374async fn try_fallback_provider(request: &LlmRequestPayload) -> Option<LlmResult> {
375    for route in &request.route_fallbacks {
376        if route.provider == request.provider && route.model == request.model {
377            continue;
378        }
379        let Ok(fb_key) = super::helpers::resolve_api_key(&route.provider) else {
380            continue;
381        };
382
383        let mut fb_request = request.clone();
384        fb_request.provider = route.provider.clone();
385        fb_request.model = route.model.clone();
386        fb_request.api_key = fb_key;
387        if super::ensure_real_llm_allowed(&fb_request.provider).is_err() {
388            continue;
389        }
390        if let Ok(result) = vm_call_llm_api(&fb_request, None).await {
391            return Some(result);
392        }
393    }
394
395    let mut fallback_providers = Vec::<String>::new();
396    for provider in &request.fallback_chain {
397        if provider != &request.provider && !fallback_providers.contains(provider) {
398            fallback_providers.push(provider.clone());
399        }
400    }
401    if let Some(pdef) = crate::llm_config::provider_config(&request.provider) {
402        if let Some(fallback_provider) = pdef.fallback {
403            if fallback_provider != request.provider
404                && !fallback_providers.contains(&fallback_provider)
405            {
406                fallback_providers.push(fallback_provider);
407            }
408        }
409    }
410    if fallback_providers.is_empty() {
411        return None;
412    }
413
414    for fallback_provider in fallback_providers {
415        let Ok(fb_key) = super::helpers::resolve_api_key(&fallback_provider) else {
416            continue;
417        };
418
419        let mut fb_request = request.clone();
420        fb_request.provider = fallback_provider;
421        fb_request.api_key = fb_key;
422        if super::ensure_real_llm_allowed(&fb_request.provider).is_err() {
423            continue;
424        }
425        if let Ok(result) = vm_call_llm_api(&fb_request, None).await {
426            return Some(result);
427        }
428    }
429
430    None
431}
432
433#[cfg(test)]
434mod tests {
435    use super::options::base_opts;
436    use super::{
437        vm_call_llm_full, vm_call_llm_full_streaming, vm_call_llm_full_streaming_offthread,
438        LlmRequestPayload, ThinkingConfig,
439    };
440    use crate::llm::env_guard;
441
442    struct ScopedEnvVar {
443        key: &'static str,
444        previous: Option<String>,
445    }
446
447    impl ScopedEnvVar {
448        fn set(key: &'static str, value: &str) -> Self {
449            let previous = std::env::var(key).ok();
450            unsafe {
451                std::env::set_var(key, value);
452            }
453            Self { key, previous }
454        }
455
456        fn remove(key: &'static str) -> Self {
457            let previous = std::env::var(key).ok();
458            unsafe {
459                std::env::remove_var(key);
460            }
461            Self { key, previous }
462        }
463    }
464
465    impl Drop for ScopedEnvVar {
466        fn drop(&mut self) {
467            match &self.previous {
468                Some(value) => unsafe { std::env::set_var(self.key, value) },
469                None => unsafe { std::env::remove_var(self.key) },
470            }
471        }
472    }
473
474    fn allow_stubbed_llm_transport() -> ScopedEnvVar {
475        ScopedEnvVar::remove(crate::llm::LLM_CALLS_DISABLED_ENV)
476    }
477
478    #[test]
479    fn openai_compat_prefill_appends_assistant_and_sets_chat_template_kwargs() {
480        use crate::llm::providers::OpenAiCompatibleProvider;
481
482        let mut opts = base_opts("local");
483        opts.model = "Qwen/Qwen3.5-Coder-32B".to_string();
484        opts.prefill = Some("<done>##DONE##</done>".to_string());
485        let payload = LlmRequestPayload::from(&opts);
486        let body = OpenAiCompatibleProvider::build_request_body(&payload, false);
487
488        let messages = body["messages"].as_array().expect("messages array");
489        let last = messages.last().expect("at least one message");
490        assert_eq!(last["role"].as_str(), Some("assistant"));
491        assert_eq!(last["content"].as_str(), Some("<done>##DONE##</done>"));
492
493        let kw = &body["chat_template_kwargs"];
494        assert_eq!(kw["add_generation_prompt"].as_bool(), Some(false));
495        assert_eq!(kw["continue_final_message"].as_bool(), Some(true));
496    }
497
498    #[test]
499    fn openai_compat_without_prefill_omits_continue_flags() {
500        use crate::llm::providers::OpenAiCompatibleProvider;
501
502        let opts = base_opts("openai");
503        let payload = LlmRequestPayload::from(&opts);
504        let body = OpenAiCompatibleProvider::build_request_body(&payload, false);
505
506        let kw = &body["chat_template_kwargs"];
507        assert!(kw.get("add_generation_prompt").is_none());
508        assert!(kw.get("continue_final_message").is_none());
509    }
510
511    #[test]
512    fn anthropic_prefill_appends_assistant_for_legacy_model() {
513        use crate::llm::providers::AnthropicProvider;
514
515        let mut opts = base_opts("anthropic");
516        opts.model = "claude-sonnet-4-20250514".to_string();
517        opts.prefill = Some("<done>##DONE##</done>".to_string());
518        let payload = LlmRequestPayload::from(&opts);
519        let body = AnthropicProvider::build_request_body(&payload);
520
521        let messages = body["messages"].as_array().expect("messages array");
522        let last = messages.last().expect("at least one message");
523        assert_eq!(last["role"].as_str(), Some("assistant"));
524        assert_eq!(last["content"].as_str(), Some("<done>##DONE##</done>"));
525    }
526
527    #[test]
528    fn anthropic_prefill_skipped_for_deprecated_4_6_model() {
529        use crate::llm::providers::AnthropicProvider;
530
531        let mut opts = base_opts("anthropic");
532        opts.model = "claude-opus-4-6".to_string();
533        opts.prefill = Some("<done>##DONE##</done>".to_string());
534        let payload = LlmRequestPayload::from(&opts);
535        let body = AnthropicProvider::build_request_body(&payload);
536
537        let messages = body["messages"].as_array().expect("messages array");
538        // User message only; prefill dropped silently.
539        assert_eq!(messages.len(), 1);
540        assert_eq!(messages[0]["role"].as_str(), Some("user"));
541    }
542
543    #[test]
544    fn anthropic_prefill_skipped_for_opus_4_7() {
545        use crate::llm::providers::AnthropicProvider;
546
547        let mut opts = base_opts("anthropic");
548        opts.model = "claude-opus-4-7".to_string();
549        opts.prefill = Some("<done>##DONE##</done>".to_string());
550        let payload = LlmRequestPayload::from(&opts);
551        let body = AnthropicProvider::build_request_body(&payload);
552
553        let messages = body["messages"].as_array().expect("messages array");
554        assert_eq!(messages.len(), 1);
555        assert_eq!(messages[0]["role"].as_str(), Some("user"));
556    }
557
558    #[test]
559    fn anthropic_sampling_params_stripped_for_opus_4_7() {
560        use crate::llm::providers::AnthropicProvider;
561
562        let mut opts = base_opts("anthropic");
563        opts.model = "claude-opus-4-7".to_string();
564        // base_opts already supplies temperature/top_p/top_k.
565        let payload = LlmRequestPayload::from(&opts);
566        let body = AnthropicProvider::build_request_body(&payload);
567
568        assert!(
569            body.get("temperature").is_none(),
570            "Opus 4.7 body must omit temperature (returns HTTP 400 otherwise)"
571        );
572        assert!(body.get("top_p").is_none(), "Opus 4.7 body must omit top_p");
573        assert!(body.get("top_k").is_none(), "Opus 4.7 body must omit top_k");
574    }
575
576    #[test]
577    fn anthropic_sampling_params_preserved_for_opus_4_6() {
578        use crate::llm::providers::AnthropicProvider;
579
580        let mut opts = base_opts("anthropic");
581        opts.model = "claude-opus-4-6".to_string();
582        let payload = LlmRequestPayload::from(&opts);
583        let body = AnthropicProvider::build_request_body(&payload);
584
585        assert_eq!(body["temperature"].as_f64(), Some(0.2));
586        assert_eq!(body["top_p"].as_f64(), Some(0.8));
587        assert_eq!(body["top_k"].as_i64(), Some(40));
588    }
589
590    #[test]
591    fn disabled_llm_calls_reject_real_provider_before_transport() {
592        let _guard = env_guard();
593        let _disabled = ScopedEnvVar::set(crate::llm::LLM_CALLS_DISABLED_ENV, "1");
594        let runtime = tokio::runtime::Builder::new_current_thread()
595            .enable_all()
596            .build()
597            .expect("runtime");
598
599        let err = runtime
600            .block_on(vm_call_llm_full(&base_opts("local")))
601            .expect_err("local provider should be blocked before transport");
602        let message = err.to_string();
603        assert!(message.contains("HARN_LLM_CALLS_DISABLED"), "{message}");
604        assert!(message.contains("provider `local`"), "{message}");
605    }
606
607    #[test]
608    fn offthread_error_preserves_schema_stream_abort_category() {
609        let abort = super::SchemaStreamAbort {
610            provider: "openrouter".to_string(),
611            model: "mistralai/devstral-small".to_string(),
612            reason: "expected JSON value, got '`'".to_string(),
613            path: "$".to_string(),
614            chunks_consumed: 1,
615        };
616
617        let err = super::OffthreadLlmError::from_vm_error(abort.into_vm_error()).into_vm_error();
618        let parsed = super::parse_schema_stream_abort(&err)
619            .expect("schema stream abort must survive off-thread conversion");
620
621        assert_eq!(parsed.provider, "openrouter");
622        assert_eq!(parsed.model, "mistralai/devstral-small");
623        assert_eq!(parsed.path, "$");
624        assert_eq!(parsed.chunks_consumed, 1);
625    }
626
627    #[test]
628    fn disabled_llm_calls_still_allow_mock_provider() {
629        let _guard = env_guard();
630        let _disabled = ScopedEnvVar::set(crate::llm::LLM_CALLS_DISABLED_ENV, "1");
631        let runtime = tokio::runtime::Builder::new_current_thread()
632            .enable_all()
633            .build()
634            .expect("runtime");
635
636        let result = runtime
637            .block_on(vm_call_llm_full(&base_opts("mock")))
638            .expect("mock provider remains available");
639        assert_eq!(result.provider, "mock");
640    }
641
642    #[test]
643    fn fake_provider_routes_through_full_pipeline_with_streaming_deltas() {
644        use crate::llm::fake::{
645            install_fake_llm_script, FakeLlmEvent, FakeLlmScript, FakeStopReason,
646        };
647
648        let _guard = env_guard();
649        // Even with HARN_LLM_CALLS_DISABLED, the fake must pass through —
650        // it never hits the network, so it must not be gated by that env.
651        let _disabled = ScopedEnvVar::set(crate::llm::LLM_CALLS_DISABLED_ENV, "1");
652        let runtime = tokio::runtime::Builder::new_current_thread()
653            .enable_all()
654            .build()
655            .expect("runtime");
656
657        let _script = install_fake_llm_script(FakeLlmScript::streaming(vec![
658            FakeLlmEvent::Token("alpha".into()),
659            FakeLlmEvent::Token(" beta".into()),
660            FakeLlmEvent::Done(FakeStopReason::EndTurn),
661        ]));
662
663        runtime.block_on(async {
664            let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
665            let result = vm_call_llm_full_streaming(&base_opts("fake"), tx)
666                .await
667                .expect("fake provider routes through dispatch");
668            assert_eq!(result.provider, "fake");
669            assert_eq!(result.text, "alpha beta");
670            let mut deltas = Vec::new();
671            while let Ok(delta) = rx.try_recv() {
672                deltas.push(delta);
673            }
674            assert_eq!(deltas, vec!["alpha".to_string(), " beta".to_string()]);
675        });
676    }
677
678    #[test]
679    fn anthropic_thinking_rewritten_to_adaptive_for_opus_4_7() {
680        use crate::llm::providers::AnthropicProvider;
681
682        let mut opts = base_opts("anthropic");
683        opts.model = "claude-opus-4-7".to_string();
684        opts.thinking = ThinkingConfig::Enabled {
685            budget_tokens: None,
686        };
687        let payload = LlmRequestPayload::from(&opts);
688        let body = AnthropicProvider::build_request_body(&payload);
689
690        let thinking = &body["thinking"];
691        assert_eq!(thinking["type"].as_str(), Some("adaptive"));
692        assert!(
693            thinking.get("budget_tokens").is_none(),
694            "Opus 4.7 adaptive thinking must not carry budget_tokens"
695        );
696    }
697
698    #[test]
699    fn anthropic_thinking_budget_discarded_for_opus_4_7() {
700        use crate::llm::providers::AnthropicProvider;
701
702        let mut opts = base_opts("anthropic");
703        opts.model = "claude-opus-4-7".to_string();
704        opts.thinking = ThinkingConfig::Enabled {
705            budget_tokens: Some(32000),
706        };
707        let payload = LlmRequestPayload::from(&opts);
708        let body = AnthropicProvider::build_request_body(&payload);
709
710        let thinking = &body["thinking"];
711        assert_eq!(thinking["type"].as_str(), Some("adaptive"));
712        assert!(thinking.get("budget_tokens").is_none());
713    }
714
715    #[test]
716    fn anthropic_thinking_preserves_extended_for_opus_4_6() {
717        use crate::llm::providers::AnthropicProvider;
718
719        let mut opts = base_opts("anthropic");
720        opts.model = "claude-opus-4-6".to_string();
721        opts.thinking = ThinkingConfig::Enabled {
722            budget_tokens: Some(16000),
723        };
724        let payload = LlmRequestPayload::from(&opts);
725        let body = AnthropicProvider::build_request_body(&payload);
726
727        let thinking = &body["thinking"];
728        assert_eq!(thinking["type"].as_str(), Some("enabled"));
729        assert_eq!(thinking["budget_tokens"].as_i64(), Some(16000));
730    }
731
732    #[test]
733    fn anthropic_prefill_preserved_for_or_opus_dotted_older_generations() {
734        use crate::llm::providers::AnthropicProvider;
735
736        // Dotted "claude-opus-4.5" style should NOT hit the 4.6 gate.
737        let mut opts = base_opts("anthropic");
738        opts.model = "anthropic/claude-opus-4.5".to_string();
739        opts.prefill = Some("<done>##DONE##</done>".to_string());
740        let payload = LlmRequestPayload::from(&opts);
741        let body = AnthropicProvider::build_request_body(&payload);
742
743        let messages = body["messages"].as_array().expect("messages array");
744        assert_eq!(messages.len(), 2);
745        assert_eq!(messages.last().unwrap()["role"].as_str(), Some("assistant"));
746    }
747
748    #[test]
749    fn anthropic_prefill_skipped_for_or_opus_4_7_dotted() {
750        use crate::llm::providers::AnthropicProvider;
751
752        let mut opts = base_opts("anthropic");
753        opts.model = "anthropic/claude-opus-4.7".to_string();
754        opts.prefill = Some("<done>##DONE##</done>".to_string());
755        let payload = LlmRequestPayload::from(&opts);
756        let body = AnthropicProvider::build_request_body(&payload);
757
758        let messages = body["messages"].as_array().expect("messages array");
759        assert_eq!(messages.len(), 1);
760        assert_eq!(messages[0]["role"].as_str(), Some("user"));
761    }
762
763    /// Cooperative accept: blocks the stub thread on a real
764    /// `accept()` call until a client connects, then returns the
765    /// stream. Shutdown wakes the thread by self-connecting to the
766    /// listener (see [`LlmStub::drop`]) — when the resulting `accept`
767    /// returns, the shutdown flag is checked and the synthetic stream
768    /// is discarded.
769    ///
770    /// This replaces a previous nonblocking polling loop with a 5ms
771    /// sleep tick. Polling introduced two flake modes under nextest's
772    /// 50× concurrent flake-detection profile: (1) a real client
773    /// connection could land between two polls and reqwest could time
774    /// out on the SYN-ACK before the stub thread woke; (2) under
775    /// heavy CPU contention the 5ms tick could stretch to tens of
776    /// milliseconds, compounding (1). Blocking accept removes the
777    /// scheduling-latency variable entirely.
778    fn accept_with_shutdown(
779        listener: &std::net::TcpListener,
780        label: &str,
781        shutdown: &std::sync::atomic::AtomicBool,
782    ) -> Option<std::net::TcpStream> {
783        let (stream, _peer) = listener
784            .accept()
785            .unwrap_or_else(|e| panic!("{label}: accept failed: {e}"));
786        if shutdown.load(std::sync::atomic::Ordering::Acquire) {
787            drop(stream);
788            return None;
789        }
790        stream
791            .set_read_timeout(Some(std::time::Duration::from_secs(30)))
792            .ok();
793        stream
794            .set_write_timeout(Some(std::time::Duration::from_secs(30)))
795            .ok();
796        Some(stream)
797    }
798
799    /// Wake a stub thread blocked in [`accept_with_shutdown`] by
800    /// opening a one-shot self-connection to its listener. The thread
801    /// then observes the shutdown flag and exits without serving the
802    /// connection. The connect uses a short timeout so a deferred
803    /// shutdown (e.g. drop during panic unwind on a saturated CI
804    /// worker) cannot wedge the test process.
805    fn wake_accept_for_shutdown(addr: std::net::SocketAddr) {
806        let _ = std::net::TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500));
807    }
808
809    /// RAII guard for an in-process LLM stub. Dropping signals the stub
810    /// thread to exit and joins it so no FDs leak past the test, even on
811    /// panic.
812    struct LlmStub {
813        addr: std::net::SocketAddr,
814        shutdown: std::sync::Arc<std::sync::atomic::AtomicBool>,
815        handle: Option<std::thread::JoinHandle<()>>,
816        /// Maximum number of `accept()` calls the stub thread can be
817        /// parked on. Single-shot stubs use 1; `spawn_llm_stub_many`
818        /// uses its connection count. Drop fires that many self-
819        /// connections so every parked accept observes shutdown.
820        pending_accepts: usize,
821    }
822
823    impl LlmStub {
824        fn addr(&self) -> std::net::SocketAddr {
825            self.addr
826        }
827    }
828
829    impl Drop for LlmStub {
830        fn drop(&mut self) {
831            self.shutdown
832                .store(true, std::sync::atomic::Ordering::Release);
833            // Self-connect to unblock any thread parked inside
834            // `accept_with_shutdown`. Multiple stubs in
835            // `spawn_llm_stub_many` may need waking, so issue one
836            // wake per outstanding accept slot.
837            for _ in 0..self.pending_accepts.max(1) {
838                wake_accept_for_shutdown(self.addr);
839            }
840            if let Some(handle) = self.handle.take() {
841                let _ = handle.join();
842            }
843        }
844    }
845
846    /// Bind a localhost listener and run `body` on a background thread once
847    /// a client connects. Wraps the listener in an [`LlmStub`] guard whose
848    /// lifetime bounds the stub thread.
849    fn spawn_llm_stub<F>(label: &'static str, body: F) -> LlmStub
850    where
851        F: FnOnce(&mut std::net::TcpStream) + Send + 'static,
852    {
853        use std::net::TcpListener;
854        let listener = TcpListener::bind("127.0.0.1:0").expect("bind llm stub");
855        let addr = listener.local_addr().expect("stub addr");
856        let shutdown = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
857        let shutdown_thread = shutdown.clone();
858        let handle = std::thread::spawn(move || {
859            let Some(mut stream) = accept_with_shutdown(&listener, label, &shutdown_thread) else {
860                return;
861            };
862            body(&mut stream);
863        });
864        LlmStub {
865            addr,
866            shutdown,
867            handle: Some(handle),
868            pending_accepts: 1,
869        }
870    }
871
872    fn spawn_llm_stub_many<F>(label: &'static str, connections: usize, mut body: F) -> LlmStub
873    where
874        F: FnMut(usize, &mut std::net::TcpStream) + Send + 'static,
875    {
876        use std::net::TcpListener;
877        let listener = TcpListener::bind("127.0.0.1:0").expect("bind llm stub");
878        let addr = listener.local_addr().expect("stub addr");
879        let shutdown = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
880        let shutdown_thread = shutdown.clone();
881        let handle = std::thread::spawn(move || {
882            for attempt in 0..connections {
883                let Some(mut stream) = accept_with_shutdown(&listener, label, &shutdown_thread)
884                else {
885                    return;
886                };
887                body(attempt, &mut stream);
888            }
889        });
890        LlmStub {
891            addr,
892            shutdown,
893            handle: Some(handle),
894            pending_accepts: connections,
895        }
896    }
897
898    fn spawn_ollama_stub() -> LlmStub {
899        spawn_llm_stub("ollama stub", |stream| {
900            use std::io::{Read, Write};
901            let mut buf = vec![0u8; 8192];
902            let n = stream.read(&mut buf).expect("read request");
903            let request = String::from_utf8_lossy(&buf[..n]);
904            assert!(request.starts_with("POST /api/chat HTTP/1.1\r\n"));
905
906            let body = concat!(
907                "{\"message\":{\"role\":\"assistant\",\"content\":\"hello \"},\"done\":false,\"model\":\"stub-model\"}\n",
908                "{\"message\":{\"role\":\"assistant\",\"content\":\"world\"},\"done\":false}\n",
909                "{\"done\":true,\"prompt_eval_count\":3,\"eval_count\":2,\"model\":\"stub-model\"}\n"
910            );
911            let response = format!(
912                "HTTP/1.1 200 OK\r\ncontent-type: application/x-ndjson\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
913                body.len(),
914                body
915            );
916            stream
917                .write_all(response.as_bytes())
918                .expect("write response");
919        })
920    }
921
922    fn spawn_ollama_empty_then_success_stub(
923        request_count: std::sync::Arc<std::sync::atomic::AtomicUsize>,
924    ) -> LlmStub {
925        spawn_llm_stub_many("ollama retry stub", 2, move |attempt, stream| {
926            use std::io::{Read, Write};
927            request_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
928            let mut buf = vec![0u8; 8192];
929            let n = stream.read(&mut buf).expect("read request");
930            let request = String::from_utf8_lossy(&buf[..n]);
931            assert!(request.starts_with("POST /api/chat HTTP/1.1\r\n"));
932
933            let body = if attempt == 0 {
934                "{\"message\":{\"role\":\"assistant\",\"content\":\"\"},\"done\":true,\"prompt_eval_count\":5,\"eval_count\":3}\n"
935            } else {
936                concat!(
937                    "{\"message\":{\"role\":\"assistant\",\"content\":\"retried\"},\"done\":false,\"model\":\"stub-model\"}\n",
938                    "{\"done\":true,\"prompt_eval_count\":5,\"eval_count\":1,\"model\":\"stub-model\"}\n"
939                )
940            };
941            let response = format!(
942                "HTTP/1.1 200 OK\r\ncontent-type: application/x-ndjson\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
943                body.len(),
944                body
945            );
946            stream
947                .write_all(response.as_bytes())
948                .expect("write response");
949        })
950    }
951
952    fn spawn_ollama_stub_with_body_capture(
953        captured: std::sync::Arc<std::sync::Mutex<Option<String>>>,
954    ) -> LlmStub {
955        spawn_llm_stub("ollama stub (capture)", move |stream| {
956            use std::io::{Read, Write};
957            let mut buf = vec![0u8; 16384];
958            let n = stream.read(&mut buf).expect("read request");
959            let request = String::from_utf8_lossy(&buf[..n]).to_string();
960            let body = request
961                .split("\r\n\r\n")
962                .nth(1)
963                .unwrap_or_default()
964                .to_string();
965            *captured.lock().expect("capture body") = Some(body);
966
967            let body = concat!(
968                "{\"message\":{\"role\":\"assistant\",\"content\":\"ok\"},\"done\":false}\n",
969                "{\"done\":true,\"prompt_eval_count\":1,\"eval_count\":1}\n"
970            );
971            let response = format!(
972                "HTTP/1.1 200 OK\r\ncontent-type: application/x-ndjson\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
973                body.len(),
974                body
975            );
976            stream
977                .write_all(response.as_bytes())
978                .expect("write response");
979        })
980    }
981
982    fn spawn_ollama_raw_generate_stub(
983        captured: std::sync::Arc<std::sync::Mutex<Option<String>>>,
984    ) -> LlmStub {
985        spawn_llm_stub("ollama raw stub", move |stream| {
986            use std::io::{Read, Write};
987            let mut buf = vec![0u8; 16384];
988            let n = stream.read(&mut buf).expect("read request");
989            let request = String::from_utf8_lossy(&buf[..n]).to_string();
990            assert!(request.starts_with("POST /api/generate HTTP/1.1\r\n"));
991            let body = request
992                .split("\r\n\r\n")
993                .nth(1)
994                .unwrap_or_default()
995                .to_string();
996            *captured.lock().expect("capture body") = Some(body);
997
998            let body = concat!(
999                "{\"response\":\"<tool_call>\\nedit({ path: \\\"a.rs\\\" })\\n</tool_call>\",\"done\":false,\"model\":\"qwen3.5:stub\"}\n",
1000                "{\"done\":true,\"prompt_eval_count\":7,\"eval_count\":11,\"model\":\"qwen3.5:stub\",\"done_reason\":\"stop\"}\n"
1001            );
1002            let response = format!(
1003                "HTTP/1.1 200 OK\r\ncontent-type: application/x-ndjson\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
1004                body.len(),
1005                body
1006            );
1007            stream
1008                .write_all(response.as_bytes())
1009                .expect("write response");
1010        })
1011    }
1012
1013    fn spawn_anthropic_stub_with_request_capture(
1014        captured: std::sync::Arc<std::sync::Mutex<Option<String>>>,
1015    ) -> LlmStub {
1016        spawn_llm_stub("anthropic stub (capture)", move |stream| {
1017            use std::io::{Read, Write};
1018            let mut buf = vec![0u8; 16384];
1019            let n = stream.read(&mut buf).expect("read request");
1020            let request = String::from_utf8_lossy(&buf[..n]).to_string();
1021            *captured.lock().expect("capture request") = Some(request);
1022
1023            let body = concat!(
1024                r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-opus-4-6","#,
1025                r#""content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","#,
1026                r#""usage":{"input_tokens":1,"output_tokens":1}}"#
1027            );
1028            let response = format!(
1029                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
1030                body.len(),
1031                body
1032            );
1033            stream
1034                .write_all(response.as_bytes())
1035                .expect("write response");
1036        })
1037    }
1038
1039    #[test]
1040    fn anthropic_interleaved_thinking_beta_header_is_sent_for_supported_model() {
1041        let _guard = env_guard();
1042        let _allow_llm_transport = allow_stubbed_llm_transport();
1043        let runtime = tokio::runtime::Builder::new_current_thread()
1044            .enable_all()
1045            .build()
1046            .expect("runtime");
1047
1048        runtime.block_on(async {
1049            let captured = std::sync::Arc::new(std::sync::Mutex::new(None));
1050            let server = spawn_anthropic_stub_with_request_capture(captured.clone());
1051            let mut overlay = crate::llm_config::ProvidersConfig::default();
1052            overlay.providers.insert(
1053                "anthropic".to_string(),
1054                crate::llm_config::ProviderDef {
1055                    base_url: format!("http://{}", server.addr()),
1056                    auth_style: "none".to_string(),
1057                    auth_env: crate::llm_config::AuthEnv::None,
1058                    extra_headers: std::collections::BTreeMap::from([(
1059                        "anthropic-version".to_string(),
1060                        "2023-06-01".to_string(),
1061                    )]),
1062                    chat_endpoint: "/messages".to_string(),
1063                    ..Default::default()
1064                },
1065            );
1066            crate::llm_config::set_user_overrides(Some(overlay));
1067            crate::llm::helpers::reset_provider_key_cache();
1068
1069            let mut opts = base_opts("anthropic");
1070            opts.model = "claude-opus-4-6".to_string();
1071            opts.stream = false;
1072            opts.thinking = ThinkingConfig::Enabled {
1073                budget_tokens: Some(8000),
1074            };
1075            let result = vm_call_llm_full(&opts)
1076                .await
1077                .expect("stubbed Anthropic response");
1078
1079            crate::llm_config::clear_user_overrides();
1080            crate::llm::helpers::reset_provider_key_cache();
1081            drop(server);
1082
1083            assert_eq!(result.text, "ok");
1084            let request = captured
1085                .lock()
1086                .expect("captured request")
1087                .clone()
1088                .expect("request captured")
1089                .to_lowercase();
1090            assert!(
1091                request.contains("anthropic-beta: interleaved-thinking-2025-05-14\r\n"),
1092                "{request}"
1093            );
1094        });
1095    }
1096
1097    #[test]
1098    fn offthread_streaming_completes_inside_localset() {
1099        let _guard = env_guard();
1100        let _allow_llm_transport = allow_stubbed_llm_transport();
1101        let runtime = tokio::runtime::Builder::new_multi_thread()
1102            .enable_all()
1103            .worker_threads(2)
1104            .build()
1105            .expect("runtime");
1106
1107        runtime.block_on(async {
1108            let server = spawn_ollama_stub();
1109            let addr = server.addr();
1110            let prev_ollama_host = std::env::var("OLLAMA_HOST").ok();
1111            unsafe {
1112                std::env::set_var("OLLAMA_HOST", format!("http://{addr}"));
1113            }
1114
1115            let local = tokio::task::LocalSet::new();
1116            let result = local
1117                .run_until(async {
1118                    let opts = base_opts("ollama");
1119                    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1120                    let result = vm_call_llm_full_streaming_offthread(&opts, tx)
1121                        .await
1122                        .expect("llm call should succeed");
1123
1124                    let mut deltas = Vec::new();
1125                    while let Ok(delta) = rx.try_recv() {
1126                        deltas.push(delta);
1127                    }
1128                    (result, deltas)
1129                })
1130                .await;
1131
1132            match prev_ollama_host {
1133                Some(value) => unsafe {
1134                    std::env::set_var("OLLAMA_HOST", value);
1135                },
1136                None => unsafe {
1137                    std::env::remove_var("OLLAMA_HOST");
1138                },
1139            }
1140
1141            drop(server);
1142
1143            let (result, deltas) = result;
1144            assert_eq!(result.text, "hello world");
1145            assert_eq!(result.model, "stub-model");
1146            assert_eq!(result.input_tokens, 3);
1147            assert_eq!(result.output_tokens, 2);
1148            assert_eq!(deltas.join(""), "hello world");
1149        });
1150    }
1151
1152    #[test]
1153    fn ollama_empty_content_done_frame_retries_once() {
1154        let _guard = env_guard();
1155        let _allow_llm_transport = allow_stubbed_llm_transport();
1156        let runtime = tokio::runtime::Builder::new_multi_thread()
1157            .enable_all()
1158            .worker_threads(2)
1159            .build()
1160            .expect("runtime");
1161
1162        runtime.block_on(async {
1163            let request_count = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
1164            let server = spawn_ollama_empty_then_success_stub(request_count.clone());
1165            let addr = server.addr();
1166            let prev_ollama_host = std::env::var("OLLAMA_HOST").ok();
1167            unsafe {
1168                std::env::set_var("OLLAMA_HOST", format!("http://{addr}"));
1169            }
1170
1171            let local = tokio::task::LocalSet::new();
1172            let result = local
1173                .run_until(async {
1174                    let opts = base_opts("ollama");
1175                    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1176                    let result = vm_call_llm_full_streaming_offthread(&opts, tx)
1177                        .await
1178                        .expect("retry should recover from empty done frame");
1179
1180                    let mut deltas = Vec::new();
1181                    while let Ok(delta) = rx.try_recv() {
1182                        deltas.push(delta);
1183                    }
1184                    (result, deltas)
1185                })
1186                .await;
1187
1188            match prev_ollama_host {
1189                Some(value) => unsafe { std::env::set_var("OLLAMA_HOST", value) },
1190                None => unsafe { std::env::remove_var("OLLAMA_HOST") },
1191            }
1192
1193            drop(server);
1194
1195            let (result, deltas) = result;
1196            assert_eq!(request_count.load(std::sync::atomic::Ordering::SeqCst), 2);
1197            assert_eq!(result.text, "retried");
1198            assert_eq!(deltas.join(""), "retried");
1199        });
1200    }
1201
1202    #[test]
1203    fn ollama_chat_applies_env_runtime_overrides() {
1204        let _guard = env_guard();
1205        let _allow_llm_transport = allow_stubbed_llm_transport();
1206        let runtime = tokio::runtime::Builder::new_multi_thread()
1207            .enable_all()
1208            .worker_threads(2)
1209            .build()
1210            .expect("runtime");
1211
1212        runtime.block_on(async {
1213            let captured = std::sync::Arc::new(std::sync::Mutex::new(None));
1214            let server = spawn_ollama_stub_with_body_capture(captured.clone());
1215            let addr = server.addr();
1216            let prev_ollama_host = std::env::var("OLLAMA_HOST").ok();
1217            let prev_num_ctx = std::env::var("HARN_OLLAMA_NUM_CTX").ok();
1218            let prev_keep_alive = std::env::var("HARN_OLLAMA_KEEP_ALIVE").ok();
1219            unsafe {
1220                std::env::set_var("OLLAMA_HOST", format!("http://{addr}"));
1221                std::env::set_var("HARN_OLLAMA_NUM_CTX", "131072");
1222                std::env::set_var("HARN_OLLAMA_KEEP_ALIVE", "forever");
1223            }
1224
1225            let local = tokio::task::LocalSet::new();
1226            let result = local
1227                .run_until(async {
1228                    let opts = base_opts("ollama");
1229                    let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
1230                    vm_call_llm_full_streaming_offthread(&opts, tx)
1231                        .await
1232                        .expect("llm call should succeed")
1233                })
1234                .await;
1235
1236            match prev_ollama_host {
1237                Some(value) => unsafe { std::env::set_var("OLLAMA_HOST", value) },
1238                None => unsafe { std::env::remove_var("OLLAMA_HOST") },
1239            }
1240            match prev_num_ctx {
1241                Some(value) => unsafe { std::env::set_var("HARN_OLLAMA_NUM_CTX", value) },
1242                None => unsafe { std::env::remove_var("HARN_OLLAMA_NUM_CTX") },
1243            }
1244            match prev_keep_alive {
1245                Some(value) => unsafe { std::env::set_var("HARN_OLLAMA_KEEP_ALIVE", value) },
1246                None => unsafe { std::env::remove_var("HARN_OLLAMA_KEEP_ALIVE") },
1247            }
1248
1249            drop(server);
1250            assert_eq!(result.text, "ok");
1251            let body = captured
1252                .lock()
1253                .expect("captured body")
1254                .clone()
1255                .expect("request body");
1256            let json: serde_json::Value = serde_json::from_str(&body).expect("valid request json");
1257            assert_eq!(json["keep_alive"].as_i64(), Some(-1));
1258            assert_eq!(json["options"]["num_ctx"].as_u64(), Some(131072));
1259        });
1260    }
1261
1262    #[test]
1263    fn ollama_qwen_text_tool_route_bypasses_chat_parser_with_raw_generate() {
1264        let _guard = env_guard();
1265        let _allow_llm_transport = allow_stubbed_llm_transport();
1266        let runtime = tokio::runtime::Builder::new_multi_thread()
1267            .enable_all()
1268            .worker_threads(2)
1269            .build()
1270            .expect("runtime");
1271
1272        runtime.block_on(async {
1273            let captured = std::sync::Arc::new(std::sync::Mutex::new(None));
1274            let server = spawn_ollama_raw_generate_stub(captured.clone());
1275            let addr = server.addr();
1276            let prev_ollama_host = std::env::var("OLLAMA_HOST").ok();
1277            unsafe {
1278                std::env::set_var("OLLAMA_HOST", format!("http://{addr}"));
1279            }
1280
1281            let local = tokio::task::LocalSet::new();
1282            let result = local
1283                .run_until(async {
1284                    let mut opts = base_opts("ollama");
1285                    opts.model = "qwen3.5:35b-a3b-coding-nvfp4".to_string();
1286                    opts.native_tools = None;
1287                    opts.output_format = crate::llm::api::OutputFormat::Text;
1288                    opts.response_format = None;
1289                    opts.json_schema = None;
1290                    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1291                    let result = vm_call_llm_full_streaming_offthread(&opts, tx)
1292                        .await
1293                        .expect("raw-generate route should succeed");
1294                    let mut deltas = Vec::new();
1295                    while let Ok(delta) = rx.try_recv() {
1296                        deltas.push(delta);
1297                    }
1298                    (result, deltas)
1299                })
1300                .await;
1301
1302            match prev_ollama_host {
1303                Some(value) => unsafe { std::env::set_var("OLLAMA_HOST", value) },
1304                None => unsafe { std::env::remove_var("OLLAMA_HOST") },
1305            }
1306
1307            drop(server);
1308            let (result, deltas) = result;
1309            assert_eq!(
1310                result.text,
1311                "<tool_call>\nedit({ path: \"a.rs\" })\n</tool_call>"
1312            );
1313            assert_eq!(deltas.join(""), result.text);
1314            assert_eq!(result.model, "qwen3.5:stub");
1315            assert_eq!(result.input_tokens, 7);
1316            assert_eq!(result.output_tokens, 11);
1317            assert_eq!(result.stop_reason.as_deref(), Some("stop"));
1318
1319            let body = captured
1320                .lock()
1321                .expect("captured body")
1322                .clone()
1323                .expect("request body");
1324            let json: serde_json::Value = serde_json::from_str(&body).expect("valid request json");
1325            assert_eq!(json["raw"].as_bool(), Some(true));
1326            assert!(json["prompt"]
1327                .as_str()
1328                .unwrap_or_default()
1329                .contains("<|im_start|>assistant\n"));
1330            assert!(json.get("chat_template_kwargs").is_none());
1331        });
1332    }
1333
1334    #[test]
1335    fn ollama_warmup_applies_shared_runtime_settings() {
1336        let _guard = env_guard();
1337        let _allow_llm_transport = allow_stubbed_llm_transport();
1338        let runtime = tokio::runtime::Builder::new_multi_thread()
1339            .enable_all()
1340            .worker_threads(2)
1341            .build()
1342            .expect("runtime");
1343
1344        runtime.block_on(async {
1345            let captured = std::sync::Arc::new(std::sync::Mutex::new(None));
1346            let server = spawn_ollama_stub_with_body_capture(captured.clone());
1347            let addr = server.addr();
1348            let _num_ctx = ScopedEnvVar::set("HARN_OLLAMA_NUM_CTX", "65536");
1349            let _keep_alive = ScopedEnvVar::set("HARN_OLLAMA_KEEP_ALIVE", "forever");
1350
1351            super::ollama::warm_ollama_model("qwen3.5:35b", Some(&format!("http://{addr}")))
1352                .await
1353                .expect("warmup should succeed");
1354
1355            drop(server);
1356            let body = captured
1357                .lock()
1358                .expect("captured body")
1359                .clone()
1360                .expect("request body");
1361            let json: serde_json::Value = serde_json::from_str(&body).expect("valid request json");
1362            assert_eq!(json["model"].as_str(), Some("qwen3.5:35b"));
1363            assert_eq!(json["keep_alive"].as_i64(), Some(-1));
1364            assert_eq!(json["options"]["num_ctx"].as_u64(), Some(65536));
1365        });
1366    }
1367
1368    /// Bind a stub listener that serves one canned HTTP error response.
1369    /// The returned [`LlmStub`] guard owns the listener and the worker
1370    /// thread, so dropping it (test exit, panic) signals shutdown and
1371    /// joins — a stuck or misrouted client can never wedge the suite.
1372    fn spawn_openai_error_stub(
1373        status_line: &'static str,
1374        extra_headers: &'static str,
1375        body: &'static str,
1376    ) -> LlmStub {
1377        spawn_llm_stub("openai error stub", move |stream| {
1378            use std::io::{Read, Write};
1379            let mut buf = vec![0u8; 16384];
1380            let _ = stream.read(&mut buf);
1381            let response = format!(
1382                "{status_line}\r\ncontent-type: application/json\r\ncontent-length: {}\r\n{extra_headers}connection: close\r\n\r\n{body}",
1383                body.len()
1384            );
1385            let _ = stream.write_all(response.as_bytes());
1386            let _ = stream.flush();
1387        })
1388    }
1389
1390    /// Single-entrypoint helper that serializes env-var mutation and the
1391    /// LLM call behind `env_lock`, so parallel streaming error tests can't
1392    /// clobber each other's `LOCAL_LLM_BASE_URL` and leak an unconnected
1393    /// stub whose `join()` would hang the test binary.
1394    fn run_streaming_error_case(
1395        status_line: &'static str,
1396        extra_headers: &'static str,
1397        body: &'static str,
1398    ) -> String {
1399        let _guard = env_guard();
1400        let _allow_llm_transport = allow_stubbed_llm_transport();
1401        let server = spawn_openai_error_stub(status_line, extra_headers, body);
1402        let addr = server.addr();
1403        let prev = std::env::var("LOCAL_LLM_BASE_URL").ok();
1404        unsafe {
1405            std::env::set_var("LOCAL_LLM_BASE_URL", format!("http://{addr}"));
1406        }
1407        let runtime = tokio::runtime::Builder::new_multi_thread()
1408            .enable_all()
1409            .worker_threads(2)
1410            .build()
1411            .expect("runtime");
1412        let err = runtime.block_on(async {
1413            let local = tokio::task::LocalSet::new();
1414            local
1415                .run_until(async {
1416                    let mut opts = base_opts("local");
1417                    opts.tools = None;
1418                    opts.native_tools = None;
1419                    opts.tool_choice = None;
1420                    opts.output_format = crate::llm::api::OutputFormat::Text;
1421                    opts.response_format = None;
1422                    opts.json_schema = None;
1423                    opts.output_schema = None;
1424                    let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
1425                    let call = tokio::time::timeout(
1426                        // Must stay inside the stub's accept window.
1427                        std::time::Duration::from_secs(30),
1428                        vm_call_llm_full_streaming_offthread(&opts, tx),
1429                    )
1430                    .await;
1431                    match call {
1432                        Ok(Ok(_)) => panic!("expected streaming call to fail"),
1433                        Ok(Err(err)) => err.to_string(),
1434                        Err(elapsed) => panic!("streaming call timed out ({elapsed})"),
1435                    }
1436                })
1437                .await
1438        });
1439        match prev {
1440            Some(v) => unsafe { std::env::set_var("LOCAL_LLM_BASE_URL", v) },
1441            None => unsafe { std::env::remove_var("LOCAL_LLM_BASE_URL") },
1442        }
1443        drop(server);
1444        err
1445    }
1446
1447    #[test]
1448    fn streaming_path_classifies_context_overflow() {
1449        let err = run_streaming_error_case(
1450            "HTTP/1.1 400 Bad Request",
1451            "",
1452            r#"{"error":{"message":"This model's maximum context length is 8192 tokens. However, your prompt is too long."}}"#,
1453        );
1454        assert!(err.contains("[context_overflow]"), "err was: {err}");
1455        assert!(err.contains("local HTTP 400"), "err was: {err}");
1456    }
1457
1458    #[test]
1459    fn streaming_path_classifies_rate_limit_with_retry_after() {
1460        let err = run_streaming_error_case(
1461            "HTTP/1.1 429 Too Many Requests",
1462            "retry-after: 7\r\n",
1463            r#"{"error":{"type":"rate_limit_error","message":"slow down"}}"#,
1464        );
1465        assert!(err.contains("[rate_limited]"), "err was: {err}");
1466        assert!(err.contains("(retry-after: 7)"), "err was: {err}");
1467    }
1468
1469    #[test]
1470    fn streaming_path_classifies_opaque_500_as_http_error() {
1471        let err = run_streaming_error_case(
1472            "HTTP/1.1 500 Internal Server Error",
1473            "",
1474            r#"{"error":"upstream exploded"}"#,
1475        );
1476        assert!(err.contains("[http_error]"), "err was: {err}");
1477        assert!(err.contains("upstream exploded"), "err was: {err}");
1478    }
1479}
harn_vm/llm/api.rs

harn_vm/llm/
api.rs
