1use bevy_ecs::prelude::{Commands, Component, Entity, Query, Res, ResMut, Resource, With, Without};
17use leviath_providers::InferenceRequest;
18use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
19
20use crate::persistence::RunMetadata;
21use crate::pipeline::{InferenceStage, Providers};
22use crate::title_bridge::{TitleJob, TitleOutcome, run_title_job};
23
24#[derive(Resource, Clone)]
27pub struct TitleSettings(pub leviath_core::config::TitleConfig);
28
29#[derive(Component, Debug, Clone, Copy)]
32pub struct PendingTitle;
33
34#[derive(Component, Debug, Clone, Copy)]
37pub struct AwaitingTitle;
38
39#[derive(Resource)]
41pub struct TitleResults(pub UnboundedReceiver<TitleOutcome>);
42
43#[derive(Resource)]
45pub struct TitleSink(pub UnboundedSender<TitleOutcome>);
46
47const TITLE_MAX_TOKENS: usize = 64;
49const TITLE_TASK_BUDGET: usize = 2_000;
52const TITLE_MAX_LEN: usize = 80;
54
55const TITLE_SYSTEM_PROMPT: &str = "Reply with only a short title for the given task, \
56 at most 8 words. No quotes, no trailing punctuation, no explanation.";
57
58fn resolve_title_model(
66 settings: &leviath_core::config::TitleConfig,
67 run_model_label: Option<&str>,
68) -> Option<(String, String)> {
69 let run = run_model_label.and_then(|label| label.split_once('/'));
70 let provider = settings
71 .provider
72 .clone()
73 .or_else(|| run.map(|(p, _)| p.to_string()))?;
74 let model = settings.model.clone().or_else(|| match run {
75 Some((run_provider, run_model)) if run_provider == provider => Some(run_model.to_string()),
76 _ => None,
77 })?;
78 Some((provider, model))
79}
80
81fn title_request(task: &str, model: &str) -> InferenceRequest {
83 InferenceRequest {
84 system: vec![leviath_providers::SystemBlock {
93 text: TITLE_SYSTEM_PROMPT.to_string(),
94 cache_hint: leviath_core::CacheHint::Never,
95 }],
96 messages: vec![leviath_providers::Message {
97 role: "user".to_string(),
98 content: leviath_core::truncate_at_boundary(task, TITLE_TASK_BUDGET)
99 .to_string()
100 .into(),
101 cache_breakpoint: false,
102 }],
103 model: model.to_string(),
104 max_tokens: TITLE_MAX_TOKENS,
105 temperature: 0.2,
106 tools: Vec::new(),
107 extra: serde_json::Value::Null,
108 request_timeout_secs: None,
109 }
110}
111
112fn sanitize_title(raw: &str) -> String {
116 let stripped = strip_reasoning(raw);
126 let lines: Vec<&str> = stripped
127 .lines()
128 .map(str::trim)
129 .filter(|l| !l.is_empty())
130 .collect();
131 let unquote = |l: &str| l.trim_matches(['"', '\'', '`']).trim().to_string();
132 let chosen = lines
133 .iter()
134 .map(|l| unquote(l))
135 .find(|l| l.chars().count() <= TITLE_MAX_LEN)
136 .unwrap_or_else(|| lines.first().map(|l| unquote(l)).unwrap_or_default());
137 leviath_core::truncate_at_boundary(&chosen, TITLE_MAX_LEN)
138 .trim_end()
139 .to_string()
140}
141
142fn strip_reasoning(raw: &str) -> String {
146 let mut out = String::with_capacity(raw.len());
147 let mut rest = raw;
148 while let Some((before, after)) = rest.split_once("<think>") {
151 out.push_str(before);
152 match after.split_once("</think>") {
153 Some((_, tail)) => rest = tail,
154 None => return out,
155 }
156 }
157 out.push_str(rest);
158 out
159}
160
161type TitleQuery = (Entity, &'static RunMetadata);
166
167pub fn dispatch_title(
173 agents: Query<TitleQuery, (With<PendingTitle>, Without<AwaitingTitle>)>,
174 settings: Option<Res<TitleSettings>>,
175 stage: Res<InferenceStage>,
176 providers: Res<Providers>,
177 sink: Res<TitleSink>,
178 mut commands: Commands,
179) {
180 crate::tick_scope::clear();
181 for (entity, meta) in agents.iter() {
182 crate::tick_scope::enter(entity);
183 let resolved = settings
184 .as_ref()
185 .filter(|s| s.0.enabled)
186 .and_then(|s| resolve_title_model(&s.0, meta.model.as_deref()));
187 let Some((provider_name, model)) = resolved else {
188 tracing::debug!(run_id = %meta.run_id, "no usable title provider/model; skipping");
189 commands.entity(entity).remove::<PendingTitle>();
190 continue;
191 };
192 let Some(provider) = providers.0.get(&provider_name) else {
193 tracing::debug!(
194 run_id = %meta.run_id,
195 provider = %provider_name,
196 "title provider not registered; skipping"
197 );
198 commands.entity(entity).remove::<PendingTitle>();
199 continue;
200 };
201 let Some(permit) = stage.pools.try_acquire(&model) else {
202 continue; };
204
205 stage.runtime.spawn(run_title_job(
206 TitleJob {
207 entity,
208 provider,
209 request: title_request(&meta.task, &model),
210 permit,
211 },
212 std::time::Duration::from_secs(leviath_providers::DEFAULT_INFERENCE_TIMEOUT_SECS),
213 sink.0.clone(),
214 stage.wake.clone(),
215 ));
216 commands
217 .entity(entity)
218 .remove::<PendingTitle>()
219 .insert(AwaitingTitle);
220 }
221}
222
223pub fn collect_title(
227 mut results: ResMut<TitleResults>,
228 mut agents: Query<&mut RunMetadata, With<AwaitingTitle>>,
229 mut commands: Commands,
230) {
231 crate::tick_scope::clear();
232 while let Ok(outcome) = results.0.try_recv() {
233 let Ok(mut meta) = agents.get_mut(outcome.entity) else {
234 continue; };
236 crate::tick_scope::enter(outcome.entity);
237 if let Ok(raw) = outcome.result {
238 let title = sanitize_title(&raw);
239 if !title.is_empty() {
240 meta.title = Some(title);
241 }
242 }
243 commands.entity(outcome.entity).remove::<AwaitingTitle>();
244 }
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250 use bevy_ecs::schedule::Schedule;
251 use bevy_ecs::world::World;
252 use leviath_providers::{Provider, ProviderError};
253 use std::sync::Arc;
254 use tokio::runtime::Handle;
255 use tokio::sync::Notify;
256 use tokio::sync::mpsc;
257
258 struct Scripted(Result<&'static str, &'static str>);
260
261 #[async_trait::async_trait]
262 impl Provider for Scripted {
263 async fn infer(
264 &self,
265 _r: &InferenceRequest,
266 ) -> leviath_providers::Result<leviath_providers::InferenceResponse> {
267 match self.0 {
268 Ok(reply) => Ok(leviath_providers::InferenceResponse {
269 content: reply.to_string(),
270 tool_calls: vec![],
271 tokens_used: leviath_providers::TokenUsage {
272 prompt_tokens: 1,
273 completion_tokens: 1,
274 total_tokens: 2,
275 cached_tokens: 0,
276 cache_write_tokens: 0,
277 },
278 finish_reason: leviath_providers::FinishReason::Complete,
279 }),
280 Err(msg) => Err(ProviderError::Other(msg.to_string())),
281 }
282 }
283 async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
284 1
285 }
286 fn max_context_tokens(&self, _m: &str) -> usize {
287 100_000
288 }
289 fn name(&self) -> &str {
290 "mock"
291 }
292 fn capabilities(&self, _m: &str) -> leviath_providers::ModelCapabilities {
293 leviath_providers::ModelCapabilities::default()
294 }
295 }
296
297 fn metadata(model: Option<&str>) -> RunMetadata {
298 RunMetadata {
299 run_id: "run-t".to_string(),
300 agent_name: "titled".to_string(),
301 agent_path: "/a".to_string(),
302 task: "summarize the release notes".to_string(),
303 model: model.map(str::to_string),
304 workdir: "/w".to_string(),
305 num_stages: 1,
306 started_at: 0,
307 parent_run_id: None,
308 metadata: Default::default(),
309 callback_url: None,
310 callback_secret: None,
311 title: None,
312 unattended: false,
313 read_paths: None,
314 output_request: None,
315 }
316 }
317
318 fn build_world(
321 reply: Result<&'static str, &'static str>,
322 pools: crate::inference_pool::InferencePools,
323 ) -> (World, mpsc::UnboundedReceiver<TitleOutcome>) {
324 let mut registry = crate::ProviderRegistry::new();
325 registry.register("mock".to_string(), Arc::new(Scripted(reply)));
326 let (title_tx, title_rx) = mpsc::unbounded_channel();
327 let (inf_tx, _inf_rx) = mpsc::unbounded_channel();
328 let (ttx, _trx) = mpsc::unbounded_channel();
329 let (ctx, _crx) = mpsc::unbounded_channel();
330 let (cstx, _csrx) = mpsc::unbounded_channel();
331 let mut world = World::new();
332 world.insert_resource(Providers(registry));
333 world.insert_resource(InferenceStage {
334 pools: Arc::new(pools),
335 outcomes: inf_tx,
336 transition_outcomes: ttx,
337 compaction_outcomes: ctx,
338 content_summary_outcomes: cstx,
339 wake: Arc::new(Notify::new()),
340 runtime: Handle::current(),
341 exact_token_counting: false,
342 });
343 world.insert_resource(TitleSink(title_tx));
344 (world, title_rx)
345 }
346
347 fn run_dispatch(world: &mut World) {
348 let mut schedule = Schedule::default();
349 schedule.add_systems(dispatch_title);
350 schedule.run(world);
351 }
352
353 fn run_collect(world: &mut World) {
354 let mut schedule = Schedule::default();
355 schedule.add_systems(collect_title);
356 schedule.run(world);
357 }
358
359 fn default_pools() -> crate::inference_pool::InferencePools {
360 crate::inference_pool::InferencePools::new(crate::inference_pool::InferencePoolConfig::new())
361 }
362
363 #[tokio::test]
366 async fn title_deadline_frees_the_slot_when_the_provider_hangs() {
367 struct Hang;
368 #[async_trait::async_trait]
369 impl Provider for Hang {
370 async fn infer(
371 &self,
372 _r: &InferenceRequest,
373 ) -> leviath_providers::Result<leviath_providers::InferenceResponse> {
374 std::future::pending().await
375 }
376 async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
377 1
378 }
379 fn max_context_tokens(&self, _m: &str) -> usize {
380 100_000
381 }
382 fn name(&self) -> &str {
383 "hang"
384 }
385 fn capabilities(&self, _m: &str) -> leviath_providers::ModelCapabilities {
386 leviath_providers::ModelCapabilities::default()
387 }
388 }
389
390 assert_eq!(Hang.count_tokens("t", "m").await, 1);
392 assert_eq!(Hang.max_context_tokens("m"), 100_000);
393 assert_eq!(Hang.name(), "hang");
394 let _ = Hang.capabilities("m");
395 let pools = crate::inference_pool::InferencePools::new(
396 crate::inference_pool::InferencePoolConfig::new(),
397 );
398 let (tx, mut rx) = mpsc::unbounded_channel();
399 run_title_job(
400 TitleJob {
401 entity: bevy_ecs::entity::Entity::PLACEHOLDER,
402 provider: Arc::new(Hang),
403 request: title_request("task", "m"),
404 permit: pools.try_acquire("m").expect("free"),
405 },
406 std::time::Duration::from_millis(5),
407 tx,
408 Arc::new(Notify::new()),
409 )
410 .await;
411 let outcome = rx.recv().await.expect("an outcome is always reported");
412 let err = outcome.result.expect_err("the deadline must surface");
413 assert!(err.to_string().contains("deadline"), "{err}");
414 }
415
416 #[tokio::test]
417 async fn dispatch_and_collect_set_the_title() {
418 let (mut world, title_rx) = build_world(Ok("\"Release notes digest\"\n"), default_pools());
419 world.insert_resource(TitleSettings(config(None, None)));
420 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
421
422 run_dispatch(&mut world);
423 assert!(world.get::<PendingTitle>(e).is_none());
424 assert!(world.get::<AwaitingTitle>(e).is_some());
425
426 let mut title_rx = title_rx;
429 let outcome = title_rx.recv().await.expect("job reported");
430 assert_eq!(outcome.entity, e);
431 let (tx, rx) = mpsc::unbounded_channel();
432 tx.send(outcome).unwrap();
433 world.insert_resource(TitleResults(rx));
434
435 run_collect(&mut world);
436 assert_eq!(
437 world.get::<RunMetadata>(e).unwrap().title.as_deref(),
438 Some("Release notes digest")
439 );
440 assert!(world.get::<AwaitingTitle>(e).is_none());
441 }
442
443 #[tokio::test]
444 async fn provider_error_leaves_the_title_unset() {
445 let (mut world, mut title_rx) = build_world(Err("boom"), default_pools());
446 world.insert_resource(TitleSettings(config(None, None)));
447 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
448
449 run_dispatch(&mut world);
450 let outcome = title_rx.recv().await.expect("job reported");
451 assert!(outcome.result.is_err());
452 let (tx, rx) = mpsc::unbounded_channel();
453 tx.send(outcome).unwrap();
454 world.insert_resource(TitleResults(rx));
455
456 run_collect(&mut world);
457 assert_eq!(world.get::<RunMetadata>(e).unwrap().title, None);
458 assert!(world.get::<AwaitingTitle>(e).is_none());
459 }
460
461 #[tokio::test]
462 async fn whitespace_reply_leaves_the_title_unset() {
463 let (mut world, mut title_rx) = build_world(Ok(" \n \n"), default_pools());
464 world.insert_resource(TitleSettings(config(None, None)));
465 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
466
467 run_dispatch(&mut world);
468 let outcome = title_rx.recv().await.expect("job reported");
469 let (tx, rx) = mpsc::unbounded_channel();
470 tx.send(outcome).unwrap();
471 world.insert_resource(TitleResults(rx));
472
473 run_collect(&mut world);
474 assert_eq!(world.get::<RunMetadata>(e).unwrap().title, None);
475 }
476
477 #[tokio::test]
478 async fn collect_skips_a_despawned_agent() {
479 let (mut world, _title_rx) = build_world(Ok("t"), default_pools());
480 let (tx, rx) = mpsc::unbounded_channel();
481 let ghost = world.spawn(metadata(Some("mock/m"))).id();
483 world.despawn(ghost);
484 tx.send(TitleOutcome {
485 entity: ghost,
486 result: Ok("t".to_string()),
487 })
488 .unwrap();
489 world.insert_resource(TitleResults(rx));
490 run_collect(&mut world); }
492
493 #[tokio::test]
494 async fn dispatch_without_settings_drops_the_marker() {
495 let (mut world, _title_rx) = build_world(Ok("t"), default_pools());
496 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
497 run_dispatch(&mut world);
498 assert!(world.get::<PendingTitle>(e).is_none());
499 assert!(world.get::<AwaitingTitle>(e).is_none());
500 }
501
502 #[tokio::test]
503 async fn dispatch_with_disabled_settings_drops_the_marker() {
504 let (mut world, _title_rx) = build_world(Ok("t"), default_pools());
505 world.insert_resource(TitleSettings(leviath_core::config::TitleConfig {
506 enabled: false,
507 provider: None,
508 model: None,
509 }));
510 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
511 run_dispatch(&mut world);
512 assert!(world.get::<PendingTitle>(e).is_none());
513 assert!(world.get::<AwaitingTitle>(e).is_none());
514 }
515
516 #[tokio::test]
517 async fn dispatch_with_unregistered_provider_drops_the_marker() {
518 let (mut world, _title_rx) = build_world(Ok("t"), default_pools());
519 world.insert_resource(TitleSettings(config(Some("nowhere"), Some("m"))));
520 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
521 run_dispatch(&mut world);
522 assert!(world.get::<PendingTitle>(e).is_none());
523 assert!(world.get::<AwaitingTitle>(e).is_none());
524 }
525
526 #[tokio::test]
527 async fn dispatch_retries_while_the_pool_is_full() {
528 let mut cfg = crate::inference_pool::InferencePoolConfig::new();
529 cfg.set_limit("m", 1);
530 let pools = crate::inference_pool::InferencePools::new(cfg);
531 let held = pools.try_acquire("m").unwrap();
532 let (mut world, _title_rx) = build_world(Ok("t"), pools);
533 world.insert_resource(TitleSettings(config(None, None)));
534 let e = world.spawn((metadata(Some("mock/m")), PendingTitle)).id();
535
536 run_dispatch(&mut world);
537 assert!(world.get::<PendingTitle>(e).is_some());
539 assert!(world.get::<AwaitingTitle>(e).is_none());
540 drop(held);
541 }
542
543 fn config(provider: Option<&str>, model: Option<&str>) -> leviath_core::config::TitleConfig {
544 leviath_core::config::TitleConfig {
545 enabled: true,
546 provider: provider.map(str::to_string),
547 model: model.map(str::to_string),
548 }
549 }
550
551 #[test]
552 fn resolve_prefers_the_configured_pair() {
553 assert_eq!(
554 resolve_title_model(
555 &config(Some("openai"), Some("gpt-5-mini")),
556 Some("anthropic/m")
557 ),
558 Some(("openai".to_string(), "gpt-5-mini".to_string()))
559 );
560 }
561
562 #[test]
563 fn resolve_falls_back_to_the_runs_provider_and_model() {
564 assert_eq!(
565 resolve_title_model(&config(None, None), Some("anthropic/claude-x")),
566 Some(("anthropic".to_string(), "claude-x".to_string()))
567 );
568 }
569
570 #[test]
571 fn resolve_borrows_the_runs_model_only_for_the_same_provider() {
572 assert_eq!(
573 resolve_title_model(&config(Some("anthropic"), None), Some("anthropic/claude-x")),
574 Some(("anthropic".to_string(), "claude-x".to_string()))
575 );
576 assert_eq!(
578 resolve_title_model(&config(Some("openai"), None), Some("anthropic/claude-x")),
579 None
580 );
581 }
582
583 #[test]
584 fn resolve_gives_up_without_any_provider_or_model() {
585 assert_eq!(resolve_title_model(&config(None, None), None), None);
586 assert_eq!(resolve_title_model(&config(None, Some("m")), None), None);
587 assert_eq!(
589 resolve_title_model(&config(None, None), Some("bare-label")),
590 None
591 );
592 }
593
594 #[test]
595 fn title_request_truncates_the_task_and_carries_the_model() {
596 let long_task = "x".repeat(5_000);
597 let req = title_request(&long_task, "gpt-5-mini");
598 assert_eq!(req.model, "gpt-5-mini");
599 assert_eq!(req.max_tokens, TITLE_MAX_TOKENS);
600 assert_eq!(req.messages.len(), 1);
604 let expected: leviath_providers::MessageContent =
605 leviath_core::truncate_at_boundary(&long_task, TITLE_TASK_BUDGET)
606 .to_string()
607 .into();
608 assert_eq!(req.messages[0].content, expected);
609 }
610
611 #[tokio::test]
612 async fn scripted_provider_metadata_is_exercised() {
613 let p = Scripted(Ok("t"));
615 assert_eq!(p.name(), "mock");
616 assert_eq!(p.count_tokens("t", "m").await, 1);
617 assert_eq!(p.max_context_tokens("m"), 100_000);
618 let default = leviath_providers::ModelCapabilities::default();
619 assert_eq!(
620 p.capabilities("m").max_output_tokens,
621 default.max_output_tokens
622 );
623 }
624
625 #[test]
626 fn sanitize_takes_the_first_line_unquoted_and_capped() {
627 assert_eq!(
628 sanitize_title("\"Fix the login bug\"\nextra"),
629 "Fix the login bug"
630 );
631 assert_eq!(
632 sanitize_title("\n\n 'Tidy: workspace' \n"),
633 "Tidy: workspace"
634 );
635 assert_eq!(sanitize_title(" \n\t\n"), "");
636 let long = "word ".repeat(40);
637 assert!(sanitize_title(&long).len() <= TITLE_MAX_LEN);
638 }
639
640 #[test]
650 fn the_title_request_carries_its_instruction_as_a_system_block() {
651 let request = title_request("tidy the kitchen", "claude-sonnet-4-6");
652
653 assert_eq!(request.system.len(), 1);
654 assert!(request.system[0].text.contains("short title"));
655 assert!(
656 request.messages.iter().all(|m| m.role != "system"),
657 "no provider is obliged to accept a system role in messages"
658 );
659 assert_eq!(request.messages.len(), 1);
660 assert_eq!(request.messages[0].role, "user");
661 }
662
663 #[test]
667 fn sanitize_skips_reasoning_and_takes_the_line_that_is_a_title() {
668 let reply = "We need to generate a short title for the task. The task: \
669 \"Research leviath.dev and report what the docs cover\" - so \
670 something short and descriptive.\n\nLeviath Docs Coverage";
671 assert_eq!(sanitize_title(reply), "Leviath Docs Coverage");
672
673 let tagged = "<think>The user wants a title. Keep it under eight words \
675 and avoid punctuation at the end.</think>\nRetry Backoff \
676 Refactor";
677 assert_eq!(sanitize_title(tagged), "Retry Backoff Refactor");
678
679 assert_eq!(sanitize_title("<think>thinking with no end"), "");
682 }
683
684 #[test]
687 fn sanitize_leaves_an_ordinary_reply_alone() {
688 assert_eq!(
689 sanitize_title("Tidy the kitchen sink"),
690 "Tidy the kitchen sink"
691 );
692 let long = "x".repeat(TITLE_MAX_LEN * 2);
693 assert_eq!(sanitize_title(&long).chars().count(), TITLE_MAX_LEN);
694 }
695}