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

harn_vm/llm/
api.rs
