1use std::collections::{BTreeMap, BTreeSet, VecDeque};
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::{Arc, Mutex};
6use std::task::{Poll, Waker};
7
8use rustev_contract::definition::StepBody;
9use rustev_contract::descriptor::{BackendDescriptor, OutputKind};
10use rustev_contract::generation::{GenerateRequest, ModelBound, RepairRequest};
11use rustev_contract::generation_result::{FieldOrigin, TextLineage, TextOrigin, TextResult};
12use rustev_contract::judgment::Unresolved;
13use rustev_contract::optimization::{
14 OptimizationAdmission, OptimizationMechanism, OptimizationOutcome,
15};
16use rustev_contract::output::TextValue;
17use rustev_contract::plan::PlanStepDetail;
18use rustev_contract::run::{RequestRecord, RequestResult};
19use rustev_contract::{Document, Identified};
20use rustev_core::evaluate::{SemanticRequest, Supplied};
21use rustev_core::generation_eval::{
22 GenerationError, GenerationEvaluation, GenerationRequest, validate_field,
23};
24use rustev_core::seams::{CancelSignal, GenerationFields};
25
26use crate::driver::{DecisionCtx, RequestOutcome, drive, observation, step_deadline, step_ledger};
27use crate::optimization::{CacheMiss, CacheWrite, KeyParts, RequestKey, request_key};
28use crate::{Inner, Ledger, PreparedPlan};
29
30#[derive(Debug, Clone, PartialEq)]
32pub enum GenerationEvent {
33 Field {
34 step: String,
35 item: Vec<String>,
36 field: String,
37 value: TextValue,
38 attempt_id: String,
39 provisional: bool,
40 },
41 Superseded {
42 step: String,
43 item: Vec<String>,
44 field: String,
45 attempt_id: String,
46 origin: FieldOrigin,
47 },
48 Lagged {
49 count: u64,
50 },
51 Complete {
52 text_origin: TextOrigin,
53 },
54}
55
56struct Queue {
57 events: VecDeque<GenerationEvent>,
58 fields: usize,
59 buffer: usize,
60 dropped: u64,
61 closed: bool,
62 waker: Option<Waker>,
63 accepted: Vec<GenerationEvent>,
64}
65
66pub struct GenerationStream(Arc<Mutex<Queue>>);
68impl GenerationStream {
69 pub async fn recv(&mut self) -> Option<GenerationEvent> {
77 std::future::poll_fn(|cx| {
78 let mut q = self.0.lock().unwrap_or_else(|e| e.into_inner());
79 if let Some(event) = q.events.pop_front() {
80 if matches!(event, GenerationEvent::Field { .. }) {
81 q.fields -= 1;
82 if q.dropped > 0 {
83 let count = std::mem::take(&mut q.dropped);
84 let at = q
85 .events
86 .iter()
87 .position(|e| matches!(e, GenerationEvent::Complete { .. }))
88 .unwrap_or(q.events.len());
89 q.events.insert(at, GenerationEvent::Lagged { count });
90 }
91 }
92 Poll::Ready(Some(event))
93 } else if q.closed {
94 Poll::Ready(None)
95 } else {
96 q.waker = Some(cx.waker().clone());
97 Poll::Pending
98 }
99 })
100 .await
101 }
102}
103
104fn stream_capacity(
107 mut bounds: impl Iterator<Item = (u64, usize)>,
108 buffer: usize,
109) -> Result<(usize, usize), GenerationError> {
110 let overflow = || GenerationError::InvalidState("generation stream capacity overflow".into());
111 let n = bounds.try_fold(0usize, |total, (requests, fields)| {
112 let requests = usize::try_from(requests).map_err(|_| overflow())?;
113 let fields = requests.checked_mul(fields).ok_or_else(overflow)?;
114 total.checked_add(fields).ok_or_else(overflow)
115 })?;
116 let capacity = n
117 .checked_mul(2)
118 .and_then(|n| n.checked_add(buffer))
119 .and_then(|n| n.checked_add(1))
120 .ok_or_else(overflow)?;
121 Ok((n, capacity))
122}
123
124pub(crate) struct Streaming(Arc<Mutex<Queue>>);
125impl Streaming {
126 pub(crate) fn new(
127 plan: &PreparedPlan,
128 buffer: usize,
129 ) -> Result<(Self, GenerationStream), GenerationError> {
130 if !(1..=4096).contains(&buffer) {
131 return Err(GenerationError::InvalidState(
132 "stream_buffer must be in 1..=4096".into(),
133 ));
134 }
135 let bounds = plan.compiled.plan.steps.iter().filter_map(|s| {
136 let PlanStepDetail::Generate(g) = &s.detail else {
137 return None;
138 };
139 Some((g.worst_case_requests, g.output.fields.len()))
140 });
141 let (n, capacity) = stream_capacity(bounds, buffer)?;
142 let queue = Arc::new(Mutex::new(Queue {
143 events: VecDeque::with_capacity(capacity),
144 fields: 0,
145 buffer,
146 dropped: 0,
147 closed: false,
148 waker: None,
149 accepted: Vec::with_capacity(n),
150 }));
151 Ok((Self(queue.clone()), GenerationStream(queue)))
152 }
153 pub(crate) fn finish(&self, text: &TextResult) {
154 let mut q = self.0.lock().unwrap_or_else(|e| e.into_inner());
155 for event in std::mem::take(&mut q.accepted) {
156 if let GenerationEvent::Field {
157 step,
158 item,
159 field,
160 value,
161 attempt_id,
162 ..
163 } = event
164 {
165 let final_field = text
166 .steps
167 .iter()
168 .find(|s| s.step == step)
169 .and_then(|s| s.items.iter().find(|i| i.id == item))
170 .and_then(|i| i.generated.get(&field));
171 if let Some(f) = final_field {
178 let same = f.value.as_ref() == Some(&value)
179 && matches!(&f.lineage,
180 Some(TextLineage::Model { attempt, .. }) if attempt == &attempt_id);
181 if !same {
182 q.events.push_back(GenerationEvent::Superseded {
183 step,
184 item,
185 field,
186 attempt_id,
187 origin: f.origin.clone(),
188 });
189 }
190 }
191 }
192 }
193 q.events.push_back(GenerationEvent::Complete {
194 text_origin: text.text_origin,
195 });
196 let waker = q.waker.take();
197 drop(q);
198 if let Some(w) = waker {
199 w.wake();
200 }
201 }
202}
203impl Drop for Streaming {
204 fn drop(&mut self) {
205 let mut q = self.0.lock().unwrap_or_else(|e| e.into_inner());
206 q.closed = true;
207 let waker = q.waker.take();
208 drop(q);
209 if let Some(w) = waker {
210 w.wake();
211 }
212 }
213}
214
215pub(crate) struct RequestStream {
216 document: GenerateRequest,
217 queue: Arc<Mutex<Queue>>,
218 emitted: Mutex<BTreeSet<String>>,
219}
220pub(crate) struct AttemptFields<'a> {
221 request: &'a RequestStream,
222 attempt: &'a str,
223 pub enabled: bool,
224}
225impl RequestStream {
226 pub(crate) fn attempt<'a>(
227 &'a self,
228 attempt: &'a str,
229 descriptor: &BackendDescriptor,
230 ) -> AttemptFields<'a> {
231 AttemptFields {
232 request: self,
233 attempt,
234 enabled: descriptor.operations.iter().any(|o| {
235 o.operation == rustev_contract::definition::Operation::Generate
236 && matches!(
237 o.output,
238 OutputKind::TextRecord {
239 streaming: true,
240 ..
241 }
242 )
243 }),
244 }
245 }
246}
247impl GenerationFields for AttemptFields<'_> {
248 fn field(&self, name: &str, value: TextValue) {
249 let r = self.request;
250 let Some(field) = r.document.output.fields.iter().find(|f| f.name == name) else {
251 return;
252 };
253 if validate_field(field, Some(&value), &r.document.vocabularies).is_err() {
254 return;
255 }
256 if !r
258 .emitted
259 .lock()
260 .unwrap_or_else(|e| e.into_inner())
261 .insert(name.into())
262 {
263 return;
264 }
265 let event = GenerationEvent::Field {
266 step: r.document.step.clone(),
267 item: item_id(&r.document),
268 field: name.into(),
269 value,
270 attempt_id: self.attempt.into(),
271 provisional: true,
272 };
273 let mut q = r.queue.lock().unwrap_or_else(|e| e.into_inner());
274 if q.fields == q.buffer {
275 q.dropped += 1;
276 return;
277 }
278 q.fields += 1;
279 q.accepted.push(event.clone());
280 q.events.push_back(event);
281 let waker = q.waker.take();
282 drop(q);
283 if let Some(w) = waker {
284 w.wake();
285 }
286 }
287}
288fn item_id(doc: &GenerateRequest) -> Vec<String> {
289 let mut id = doc.instance.clone();
290 if matches!(doc.repair, RepairRequest::Round { .. }) {
291 id.pop();
292 }
293 id
294}
295
296pub(crate) struct StageContext<'a> {
297 pub rt: &'a Inner,
298 pub plan: &'a PreparedPlan,
299 pub deadline: u64,
300 pub submitted: u64,
301 pub cancel: &'a CancelSignal,
302 pub ledger: &'a Ledger,
303 pub decision_id: &'a str,
304 pub principal: &'a [u8],
305 pub evaluation_time: rustev_contract::Timestamp,
306 pub streaming: Option<&'a Streaming>,
307 pub decision: Option<&'a crate::optimization::decision::DecisionRun>,
311 pub rerun: bool,
313}
314pub(crate) struct PendingWrite {
315 key: RequestKey,
316 write: CacheWrite,
317}
318pub(crate) fn write_cache(cx: &StageContext<'_>, writes: Vec<PendingWrite>) {
319 if let Some(memory) = cx
320 .plan
321 .optimization
322 .as_ref()
323 .and_then(|p| p.memory_cache.as_ref())
324 {
325 for mut w in writes {
326 w.write.now = cx.rt.clock.now();
327 let _ = cx.rt.optimizer.insert(memory, w.key, w.write);
329 }
330 }
331}
332pub(crate) struct ItemClock {
333 start: Mutex<Option<u64>>,
334 duration: u64,
335}
336impl ItemClock {
337 pub(crate) fn start(&self, now: u64) {
338 self.start
339 .lock()
340 .unwrap_or_else(|e| e.into_inner())
341 .get_or_insert(now);
342 }
343 pub(crate) fn deadline(&self, limit: u64) -> u64 {
344 self.start
345 .lock()
346 .unwrap_or_else(|e| e.into_inner())
347 .map_or(limit, |t| limit.min(t.saturating_add(self.duration)))
348 }
349}
350pub(crate) fn pending_write(out: &RequestOutcome, now: u64) -> Option<PendingWrite> {
351 match (&out.supplied, out.record.attempts.last(), &out.cache) {
352 (Some(Supplied::Output(raw)), Some(a), Some((key, generation)))
353 if a.target == 0 && matches!(a.end, rustev_contract::run::AttemptEnd::Output) =>
354 {
355 Some(PendingWrite {
356 key: key.clone(),
357 write: CacheWrite {
358 output: raw.clone(),
359 source_attempt_id: a.attempt_id.clone(),
360 target: 0,
361 captured_generation: *generation,
362 now,
363 },
364 })
365 }
366 _ => None,
367 }
368}
369struct StepBudget {
370 deadline: u64,
371 item_ms: u64,
372 concurrency: usize,
373 ledger: Ledger,
374}
375struct Finished {
376 out: RequestOutcome,
377 source: Option<String>,
378 write: Option<PendingWrite>,
379}
380type Running<'a> = Pin<Box<dyn Future<Output = Result<Finished, GenerationError>> + Send + 'a>>;
381
382pub(crate) async fn run(
384 cx: &StageContext<'_>,
385 stage: &Mutex<GenerationEvaluation<'_>>,
386 mut capture: Option<&mut crate::capture::Capturing<'_>>,
387) -> Result<(Vec<RequestRecord>, Vec<PendingWrite>, bool), GenerationError> {
388 let start = cx.rt.clock.now();
389 let mut budgets = BTreeMap::new();
390 for s in &cx.plan.compiled.plan.definition.steps {
391 if let StepBody::Generate(g) = &s.body {
392 let ModelBound::Literal(item_ms) = g.limits.per_item_deadline_ms else {
393 return Err(GenerationError::InvalidState(
394 "unresolved generation deadline".into(),
395 ));
396 };
397 budgets.insert(
398 s.id.clone(),
399 StepBudget {
400 deadline: step_deadline(
401 cx.deadline,
402 start,
403 cx.plan.compiled.plan.definition.limits.deadline_ms,
404 g.limits.deadline_share,
405 ),
406 item_ms,
407 concurrency: g.limits.concurrency as usize,
408 ledger: step_ledger(&g.limits.max_cost_units),
409 },
410 );
411 }
412 }
413 let mut active: Vec<(String, Vec<String>, Running<'_>)> = vec![];
414 let mut slots: BTreeMap<(String, Vec<String>), Arc<ItemClock>> = BTreeMap::new();
415 let mut records = BTreeMap::new();
416 let mut seen = BTreeMap::new();
417 let mut writes = vec![];
418 loop {
419 if cx.cancel.is_raised() && active.is_empty() {
420 return Ok((records.into_values().collect(), vec![], true));
421 }
422 let pending = stage.lock().unwrap_or_else(|e| e.into_inner()).pending();
423 for req in &pending {
424 let next = seen.len();
425 seen.entry((req.document.step.clone(), req.document.instance.clone()))
426 .or_insert(next);
427 }
428 for req in pending {
430 if cx.cancel.is_raised() || active.len() >= cx.rt.config.max_parallel_requests {
431 break;
432 }
433 let step = req.owner.clone();
434 let item = item_id(&req.document);
435 if active.iter().any(|(s, i, _)| s == &step && i == &item) {
436 continue;
437 }
438 let Some(budget) = budgets.get(&step) else {
439 return Err(GenerationError::InvalidState(
440 "unbound generation step".into(),
441 ));
442 };
443 let occupied = slots.keys().filter(|(s, _)| s == &step).count();
444 if !slots.contains_key(&(step.clone(), item.clone())) && occupied >= budget.concurrency
445 {
446 continue;
447 }
448 let item_clock = slots
449 .entry((step.clone(), item.clone()))
450 .or_insert_with(|| {
451 Arc::new(ItemClock {
452 start: Mutex::new(None),
453 duration: budget.item_ms,
454 })
455 })
456 .clone();
457 active.push((
458 step,
459 item,
460 Box::pin(request(cx, stage, budget, item_clock, req)),
461 ));
462 }
463 if active.is_empty() {
464 break;
465 }
466 let (index, result) = std::future::poll_fn(|task| {
467 for (i, (_, _, f)) in active.iter_mut().enumerate() {
468 if let Poll::Ready(r) = f.as_mut().poll(task) {
469 return Poll::Ready((i, r));
470 }
471 }
472 Poll::Pending
473 })
474 .await;
475 let (step, item, future) = active.remove(index);
476 drop(future);
477 let Finished { out, source, write } = result?;
478 if let (Some(d), Some(Supplied::Output(raw))) = (cx.decision, &out.supplied) {
479 d.supplied(&out.step, &out.instance, raw);
480 }
481 if let Some(supplied) = out.supplied {
482 let attempt = out
483 .record
484 .attempts
485 .last()
486 .map(|a| a.attempt_id.clone())
487 .or_else(|| {
488 out.record
489 .optimization
490 .as_ref()
491 .and_then(|o| o.source_attempt_id.clone())
492 })
493 .unwrap_or_else(|| {
494 format!(
495 "{}/{}/refused",
496 crate::driver::escape(cx.decision_id),
497 crate::driver::escape(&step)
498 )
499 });
500 let mut g = stage.lock().unwrap_or_else(|e| e.into_inner());
501 if let Some(c) = capture.as_deref_mut() {
502 c.record_generation(&g, cx.plan, &out.record, &supplied);
503 }
504 g.supply_attempt(
505 &out.step,
506 &out.instance,
507 supplied,
508 &attempt,
509 (out.record.attempts.len() as u32).max(1),
510 )?;
511 if let Some(source) = source {
512 g.note_cached_attempt(&attempt, &source);
513 }
514 if !g
515 .pending()
516 .iter()
517 .any(|r| r.owner == step && item_id(&r.document) == item)
518 {
519 slots.remove(&(step, item));
520 }
521 }
522 if let Some(index) = seen.get(&(out.record.step.clone(), out.record.instance.clone())) {
523 records.insert(*index, out.record);
524 }
525 if let Some(write) = write {
526 writes.push(write);
527 }
528 }
529 Ok((
530 records.into_values().collect(),
531 writes,
532 cx.cancel.is_raised(),
533 ))
534}
535
536fn original_decision(attempt_id: &str) -> String {
539 attempt_id
540 .split('/')
541 .next()
542 .map(crate::driver::unescape)
543 .unwrap_or_default()
544}
545
546async fn request(
547 cx: &StageContext<'_>,
548 stage: &Mutex<GenerationEvaluation<'_>>,
549 budget: &StepBudget,
550 item_clock: Arc<ItemClock>,
551 req: GenerationRequest,
552) -> Result<Finished, GenerationError> {
553 let deadline = budget.deadline;
554 let is_screen = req.screen_projection.is_some();
555 let projection = match req.screen_projection {
556 Some(bytes) => bytes,
557 None => req
558 .document
559 .canonical()
560 .map_err(|e| GenerationError::InvalidState(e.to_string()))?,
561 };
562 let semantic = SemanticRequest {
563 step: req.document.step.clone(),
564 instance: req.document.instance.clone(),
565 backend_id: req.backend_id,
566 projection,
567 };
568 let stream = cx.streaming.filter(|_| !is_screen).map(|s| RequestStream {
569 document: req.document,
570 queue: s.0.clone(),
571 emitted: Mutex::new(BTreeSet::new()),
572 });
573 let drive_cx = DecisionCtx {
574 rt: cx.rt,
575 observing: None,
576 plan: cx.plan,
577 deadline,
578 submitted: cx.submitted,
579 cancel: cx.cancel,
580 ledger: cx.ledger,
581 decision_id: cx.decision_id,
582 principal: cx.principal,
583 evaluation_time: cx.evaluation_time,
584 ev: None,
585 generation: Some(stage),
586 step_ledger: Some(&budget.ledger),
587 stream: stream.as_ref(),
588 item_clock: Some(&item_clock),
589 deferred_cache: None,
592 persistent: None,
593 decision: cx.decision,
594 live: None,
595 rerun: cx.rerun,
596 };
597 if let Some(out) = cx
601 .decision
602 .and_then(|d| crate::driver::from_decision(&drive_cx, d, &semantic))
603 {
604 item_clock.start(cx.rt.clock.now());
605 let source = out
606 .record
607 .optimization
608 .as_ref()
609 .and_then(|o| o.source_attempt_id.as_deref())
610 .map(original_decision);
611 return Ok(Finished {
612 out,
613 source,
614 write: None,
615 });
616 }
617 let mut cache = None;
618 let mut diagnostics = vec![];
619 if !is_screen
620 && !cx.cancel.is_raised()
621 && cx.rt.clock.now() < item_clock.deadline(deadline)
622 && let Some(policy) = &cx.plan.optimization
623 && policy.memory_cache.is_some()
624 && let Some(target) = cx
625 .plan
626 .steps
627 .get(&semantic.step)
628 .and_then(|s| s.targets.first())
629 && let Ok(descriptor) = target.registered.backend.descriptor().id()
630 {
631 let profile = crate::optimization::profile_of(&cx.plan.compiled.plan);
632 let key = request_key(KeyParts {
633 plan_id: &cx.plan.compiled.id,
634 profile: &profile,
635 policy,
636 principal: cx.principal,
637 backend_artifact: &target.artifact,
638 backend_descriptor: &descriptor,
639 step: &semantic.step,
640 instance: &semantic.instance,
641 projection: &semantic.projection,
642 });
643 let generation = cx.rt.optimizer.generation(&key.namespace);
644 let lookup = cx.rt.optimizer.lookup(&key, cx.rt.clock.now());
645 if let Ok(hit) = lookup {
646 item_clock.start(cx.rt.clock.now());
647 let mut o = observation(
648 policy,
649 OptimizationMechanism::MemoryCache,
650 OptimizationOutcome::Hit,
651 hit.generation,
652 vec![],
653 );
654 o.entry_id = Some(hit.entry_id);
655 o.age_ms = hit.age_ms;
656 o.source_attempt_id = Some(hit.source_attempt_id.clone());
657 let source = original_decision(&hit.source_attempt_id);
658 let record = RequestRecord {
659 screen: None,
660 step: semantic.step.clone(),
661 instance: semantic.instance.clone(),
662 result: RequestResult::Output { target: hit.target },
663 attempts: vec![],
664 transitions: vec![],
665 optimization: Some(o),
666 };
667 return Ok(Finished {
668 out: RequestOutcome {
669 step: semantic.step,
670 instance: semantic.instance,
671 supplied: Some(Supplied::Output(hit.output)),
672 record,
673 cache: None,
674 },
675 source: Some(source),
676 write: None,
677 });
678 }
679 if let Err(miss) = lookup {
680 let detail = crate::driver::cache_miss_detail(miss);
681 if miss == CacheMiss::Collision && policy.admission == OptimizationAdmission::Refuse {
682 item_clock.start(cx.rt.clock.now());
683 let reason = Unresolved::BackendUnavailable {
684 detail: detail.into(),
685 };
686 let record = RequestRecord {
687 screen: None,
688 step: semantic.step.clone(),
689 instance: semantic.instance.clone(),
690 result: RequestResult::Failed(reason.clone()),
691 attempts: vec![],
692 transitions: vec![],
693 optimization: Some(observation(
694 policy,
695 OptimizationMechanism::MemoryCache,
696 OptimizationOutcome::Refused,
697 generation,
698 vec![detail.into()],
699 )),
700 };
701 return Ok(Finished {
702 out: RequestOutcome {
703 step: semantic.step,
704 instance: semantic.instance,
705 supplied: Some(Supplied::Failed(reason)),
706 record,
707 cache: None,
708 },
709 source: None,
710 write: None,
711 });
712 }
713 diagnostics.push(detail.into());
714 }
715 cache = Some((key, generation));
716 }
717 let mut out = drive(
718 &drive_cx,
719 semantic,
720 cache.as_ref().map(|(k, _)| k.clone()),
721 cache.as_ref().map_or(0, |(_, g)| *g),
722 diagnostics,
723 )
724 .await;
725 crate::driver::record_cache_cost(&mut out);
730 let write = if is_screen {
731 None
732 } else {
733 pending_write(&out, cx.rt.clock.now())
734 };
735 Ok(Finished {
736 out,
737 source: None,
738 write,
739 })
740}
741
742#[cfg(test)]
743mod tests {
744 use super::*;
745
746 #[test]
747 fn stream_capacity_rejects_overflow_before_allocation() {
748 struct Case {
749 label: &'static str,
750 bounds: Vec<(u64, usize)>,
751 buffer: usize,
752 }
753
754 let max = usize::MAX as u64;
755 for Case {
756 label,
757 bounds,
758 buffer,
759 } in [
760 Case {
761 label: "Per-step multiplication",
762 bounds: vec![(max, 2)],
763 buffer: 1,
764 },
765 Case {
766 label: "Sum across steps",
767 bounds: vec![(max, 1), (1, 1)],
768 buffer: 1,
769 },
770 Case {
771 label: "Control-event multiplication",
772 bounds: vec![(max, 1)],
773 buffer: 1,
774 },
775 Case {
776 label: "Field-buffer addition",
777 bounds: vec![(max / 2, 1)],
778 buffer: 2,
780 },
781 Case {
782 label: "Terminal-event addition",
783 bounds: vec![(max / 2, 1)],
784 buffer: 1,
786 },
787 Case {
788 label: "Conversion on narrower targets",
789 bounds: vec![(u64::MAX, 2)],
790 buffer: 1,
791 },
792 ] {
793 assert!(
794 matches!(
795 stream_capacity(bounds.into_iter(), buffer),
796 Err(GenerationError::InvalidState(detail))
797 if detail == "generation stream capacity overflow"
798 ),
799 "{label} (buffer = {buffer}) must reject capacity overflow"
800 );
801 }
802 }
803
804 #[test]
805 fn stream_capacity_reserves_control_events_separately() {
806 assert_eq!(
807 stream_capacity([(2, 3), (4, 2)].into_iter(), 1).unwrap(),
808 (14, 30)
809 );
810 assert_eq!(
811 stream_capacity(std::iter::empty(), 4096).unwrap(),
812 (0, 4097)
813 );
814 }
815}