agents_kit/prompt.rs
1//! Prompt inputs, exact terminal responses, and transient output observation.
2//!
3//! [`crate::prompt::PromptRequest`] applies per-prompt overrides to inherited workspace and
4//! Agent settings. [`crate::prompt::PromptStream`] observes the same accepted owner as unary
5//! prompting: deltas are transient, and only [`crate::prompt::PromptEvent::Completed`] conveys
6//! its authoritative response. Dropping observation does not abort accepted work.
7
8use std::fmt;
9use std::path::{Path, PathBuf};
10use std::pin::Pin;
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13
14use futures_util::{Stream, StreamExt};
15use tokio::sync::{mpsc, oneshot};
16
17use crate::{
18 KitError, KnowledgePolicy, ModelSelector, PersonalitySelector, ReasoningEffort, ToolSelection,
19 WorkspaceContext,
20};
21
22/// Strict self-contained JSON Schema requested for structured prompt output.
23///
24/// ```
25/// use agents_kit::prompt::StructuredOutputRequest;
26/// let output = StructuredOutputRequest::new(serde_json::json!({"type": "integer"}))?;
27/// assert_eq!(output.schema()["type"], "integer");
28/// # Ok::<(), agents_kit::KitError>(())
29/// ```
30#[derive(Clone, Debug, Eq, PartialEq)]
31pub struct StructuredOutputRequest {
32 schema: serde_json::Value,
33}
34
35#[derive(Clone, Copy, Debug)]
36struct StructuredOutputLocalRetriever;
37
38impl fmt::Display for StructuredOutputLocalRetriever {
39 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
40 formatter.write_str("structured output schema references must remain local")
41 }
42}
43
44impl std::error::Error for StructuredOutputLocalRetriever {}
45
46impl jsonschema::Retrieve for StructuredOutputLocalRetriever {
47 fn retrieve(
48 &self,
49 _uri: &jsonschema::Uri<String>,
50 ) -> Result<serde_json::Value, Box<dyn std::error::Error + Send + Sync>> {
51 Err(Box::new(*self))
52 }
53}
54
55impl StructuredOutputRequest {
56 /// Validates and compiles a self-contained Draft 2020-12 schema.
57 ///
58 /// Instance data that happens to contain a `$ref` key is not interpreted
59 /// as a schema reference, while actual external schema references remain
60 /// unavailable.
61 ///
62 /// ```
63 /// use agents_kit::prompt::StructuredOutputRequest;
64 /// use serde_json::json;
65 ///
66 /// let request = StructuredOutputRequest::new(json!({
67 /// "$defs": {"value": {"type": "string"}},
68 /// "allOf": [{"$ref": "#/$defs/value"}],
69 /// "const": {"$ref": "literal instance data"}
70 /// }))?;
71 /// assert!(request.schema().is_object());
72 ///
73 /// assert!(StructuredOutputRequest::new(json!({
74 /// "$ref": "https://example.invalid/schema.json"
75 /// })).is_err());
76 /// # Ok::<(), agents_kit::KitError>(())
77 /// ```
78 ///
79 /// # Errors
80 /// Returns InvalidInput unless the schema is a nonempty object that validates
81 /// against Draft 2020-12 and compiles using only local schema references.
82 pub fn new(schema: serde_json::Value) -> Result<Self, KitError> {
83 if schema.as_object().is_none_or(|object| object.is_empty()) {
84 return Err(KitError::invalid_input(
85 "structured output schema must be a nonempty JSON object",
86 ));
87 }
88 jsonschema::draft202012::meta::validate(&schema).map_err(|_| {
89 KitError::invalid_input("structured output schema is not valid Draft 2020-12")
90 })?;
91 jsonschema::draft202012::options()
92 .with_retriever(StructuredOutputLocalRetriever)
93 .build(&schema)
94 .map_err(|_| {
95 KitError::invalid_input(
96 "structured output schema must compile as a self-contained document",
97 )
98 })?;
99 Ok(Self { schema })
100 }
101
102 /// Returns the validated schema document.
103 #[must_use]
104 ///
105 /// ```
106 /// use agents_kit::prompt::StructuredOutputRequest;
107 /// let schema = serde_json::json!({"type":"boolean"});
108 /// assert_eq!(StructuredOutputRequest::new(schema.clone())?.schema(), &schema);
109 /// # Ok::<(), agents_kit::KitError>(())
110 /// ```
111 pub fn schema(&self) -> &serde_json::Value {
112 &self.schema
113 }
114}
115
116#[derive(Clone, Debug)]
117/// Supplies one foreground prompt and its optional portable overrides.
118///
119/// ```
120/// use agents_kit::{
121/// model::ReasoningEffort,
122/// prompt::PromptRequest,
123/// };
124///
125/// let request = PromptRequest::new("Summarize the current task")?
126/// .with_reasoning_effort(ReasoningEffort::medium());
127/// assert_eq!(request.text(), "Summarize the current task");
128/// # Ok::<(), agents_kit::KitError>(())
129/// ```
130pub struct PromptRequest {
131 text: String,
132 cwd: Option<PathBuf>,
133 model: Option<ModelSelector>,
134 reasoning_effort: Option<ReasoningEffort>,
135 personality: Option<PersonalitySelector>,
136 instructions: Option<String>,
137 timeout: Option<Duration>,
138 knowledge_policy: Option<KnowledgePolicy>,
139 structured_output: Option<StructuredOutputRequest>,
140 tool_selection: Option<ToolSelection>,
141}
142
143impl PromptRequest {
144 /// Validates and constructs a request with no overrides.
145 ///
146 /// Prompt text is retained exactly, including whitespace.
147 ///
148 /// # Errors
149 /// An empty prompt returns InvalidInput.
150 ///
151 /// ```
152 /// use agents_kit::prompt::PromptRequest;
153 /// let request = PromptRequest::new("Hello")?;
154 /// assert_eq!(request.text(), "Hello");
155 /// assert!(PromptRequest::new("").is_err());
156 /// # Ok::<(), agents_kit::KitError>(())
157 /// ```
158 pub fn new(text: impl Into<String>) -> Result<Self, KitError> {
159 let text = text.into();
160 if text.is_empty() {
161 return Err(KitError::invalid_input("prompt text must not be empty"));
162 }
163 Ok(Self {
164 text,
165 cwd: None,
166 model: None,
167 reasoning_effort: None,
168 personality: None,
169 instructions: None,
170 timeout: None,
171 knowledge_policy: None,
172 structured_output: None,
173 tool_selection: None,
174 })
175 }
176
177 /// Sets a validated path relative to the active workspace context.
178 ///
179 /// This path is relative to the active workspace, not the caller process.
180 ///
181 /// # Errors
182 /// Absolute paths, parent traversal, and invalid portable components return InvalidInput.
183 ///
184 /// ```
185 /// use agents_kit::prompt::PromptRequest;
186 /// let request = PromptRequest::new("Hello")?;
187 /// let request = request.with_cwd("src".into())?;
188 /// assert_eq!(request.cwd(), Some(std::path::Path::new("src")));
189 /// # Ok::<(), agents_kit::KitError>(())
190 /// ```
191 pub fn with_cwd(mut self, cwd: PathBuf) -> Result<Self, KitError> {
192 self.cwd = Some(WorkspaceContext::normalized_relative(&cwd)?);
193 Ok(self)
194 }
195
196 /// Sets the model selector for this prompt.
197 ///
198 /// Availability is validated by the selected backend when submitting the prompt.
199 ///
200 /// ```
201 /// use agents_kit::prompt::PromptRequest;
202 /// let request = PromptRequest::new("Hello")?;
203 /// let request = request.with_model(agents_kit::model::ModelSelector::new("model-id")?);
204 /// assert_eq!(request.model().unwrap().as_str(), "model-id");
205 /// # Ok::<(), agents_kit::KitError>(())
206 /// ```
207 pub fn with_model(mut self, model: ModelSelector) -> Self {
208 self.model = Some(model);
209 self
210 }
211
212 /// Sets the reasoning effort for this prompt.
213 ///
214 /// The value must be supported by the selected backend and model; setting it
215 /// does not validate availability. See the [backend restrictions](crate::model#supplied-selector-restrictions).
216 ///
217 /// ```
218 /// use agents_kit::prompt::PromptRequest;
219 /// let request = PromptRequest::new("Hello")?;
220 /// let request = request.with_reasoning_effort(agents_kit::model::ReasoningEffort::low());
221 /// assert_eq!(request.reasoning_effort(), Some(&agents_kit::model::ReasoningEffort::low()));
222 /// # Ok::<(), agents_kit::KitError>(())
223 /// ```
224 pub fn with_reasoning_effort(mut self, reasoning_effort: ReasoningEffort) -> Self {
225 self.reasoning_effort = Some(reasoning_effort);
226 self
227 }
228
229 /// Sets the personality selector for this prompt.
230 #[must_use]
231 ///
232 /// This retains the caller value without validating backend support. The supplied
233 /// Codex, OpenCode and Custom implementations currently reject explicit personality
234 /// selection; caller backends may advertise and support it. See [`crate::model`].
235 ///
236 /// ```
237 /// use agents_kit::prompt::PromptRequest;
238 /// let request = PromptRequest::new("Hello")?;
239 /// let request = request.with_personality(agents_kit::model::PersonalitySelector::new("concise")?);
240 /// assert_eq!(request.personality().unwrap().as_str(), "concise");
241 /// # Ok::<(), agents_kit::KitError>(())
242 /// ```
243 pub fn with_personality(mut self, personality: PersonalitySelector) -> Self {
244 self.personality = Some(personality);
245 self
246 }
247
248 /// Sets nonempty prompt-level instructions.
249 ///
250 /// Text is preserved without trimming.
251 ///
252 /// # Errors
253 /// Empty instructions return InvalidInput.
254 ///
255 /// ```
256 /// use agents_kit::prompt::PromptRequest;
257 /// let request = PromptRequest::new("Hello")?;
258 /// let request = request.with_instructions("Be concise.")?;
259 /// assert_eq!(request.instructions(), Some("Be concise."));
260 /// # Ok::<(), agents_kit::KitError>(())
261 /// ```
262 pub fn with_instructions(mut self, instructions: impl Into<String>) -> Result<Self, KitError> {
263 let instructions = instructions.into();
264 if instructions.is_empty() {
265 return Err(KitError::invalid_input(
266 "prompt instructions must not be empty",
267 ));
268 }
269 self.instructions = Some(instructions);
270 Ok(self)
271 }
272
273 /// Sets the effective managed-knowledge policy for this prompt.
274 ///
275 /// This configures private retrieval; it does not install a model-visible tool.
276 ///
277 /// # Errors
278 /// An Only policy with no scopes or duplicate scopes returns InvalidInput.
279 ///
280 /// ```
281 /// use agents_kit::prompt::PromptRequest;
282 /// let request = PromptRequest::new("Hello")?;
283 /// let request = request.with_knowledge_policy(agents_kit::knowledge::KnowledgePolicy::Disabled)?;
284 /// assert_eq!(request.knowledge_policy(), Some(&agents_kit::knowledge::KnowledgePolicy::Disabled));
285 /// # Ok::<(), agents_kit::KitError>(())
286 /// ```
287 pub fn with_knowledge_policy(mut self, policy: KnowledgePolicy) -> Result<Self, KitError> {
288 policy.validate()?;
289 self.knowledge_policy = Some(policy);
290 Ok(self)
291 }
292
293 /// Sets a nonzero foreground prompt timeout that starts after admission.
294 ///
295 /// The timeout bounds execution after admission; it does not renew lifecycle deadlines.
296 ///
297 /// # Errors
298 /// A zero timeout returns InvalidInput.
299 ///
300 /// ```
301 /// use agents_kit::prompt::PromptRequest;
302 /// let request = PromptRequest::new("Hello")?;
303 /// let request = request.with_timeout(std::time::Duration::from_secs(20))?;
304 /// assert_eq!(request.timeout(), Some(std::time::Duration::from_secs(20)));
305 /// # Ok::<(), agents_kit::KitError>(())
306 /// ```
307 pub fn with_timeout(mut self, timeout: Duration) -> Result<Self, KitError> {
308 if timeout.is_zero() {
309 return Err(KitError::invalid_input(
310 "prompt timeout must be greater than zero",
311 ));
312 }
313 self.timeout = Some(timeout);
314 Ok(self)
315 }
316
317 /// Requests strict output matching the supplied schema.
318 #[must_use]
319 ///
320 /// Selected backends may reject unsupported structured-output execution.
321 ///
322 /// ```
323 /// use agents_kit::prompt::PromptRequest;
324 /// let request = PromptRequest::new("Hello")?;
325 /// let schema = agents_kit::prompt::StructuredOutputRequest::new(serde_json::json!({"type":"string"}))?;
326 /// let request = request.with_structured_output(schema);
327 /// assert!(request.structured_output().is_some());
328 /// # Ok::<(), agents_kit::KitError>(())
329 /// ```
330 pub fn with_structured_output(mut self, structured_output: StructuredOutputRequest) -> Self {
331 self.structured_output = Some(structured_output);
332 self
333 }
334
335 /// Sets the portable tool-selection override.
336 ///
337 /// A restricted catalog is distinct from requiring a call.
338 ///
339 /// # Errors
340 /// Empty named selections, duplicate names, or invalid portable names return InvalidInput.
341 ///
342 /// ```
343 /// use agents_kit::prompt::PromptRequest;
344 /// let request = PromptRequest::new("Hello")?;
345 /// let request = request.with_tool_selection(agents_kit::tool::ToolSelection::None)?;
346 /// assert_eq!(request.tool_selection(), Some(&agents_kit::tool::ToolSelection::None));
347 /// # Ok::<(), agents_kit::KitError>(())
348 /// ```
349 pub fn with_tool_selection(mut self, selection: ToolSelection) -> Result<Self, KitError> {
350 selection.validate()?;
351 self.tool_selection = Some(selection);
352 Ok(self)
353 }
354
355 /// Returns the required prompt text.
356 ///
357 /// ```
358 /// use agents_kit::prompt::PromptRequest;
359 /// let request = PromptRequest::new("Hello")?;
360 /// assert_eq!(request.text(), "Hello");
361 /// # Ok::<(), agents_kit::KitError>(())
362 /// ```
363 pub fn text(&self) -> &str {
364 &self.text
365 }
366
367 /// Returns the optional workspace-relative cwd override.
368 ///
369 /// ```
370 /// use agents_kit::prompt::PromptRequest;
371 /// let request = PromptRequest::new("Hello")?;
372 /// assert_eq!(request.cwd(), None);
373 /// # Ok::<(), agents_kit::KitError>(())
374 /// ```
375 pub fn cwd(&self) -> Option<&Path> {
376 self.cwd.as_deref()
377 }
378
379 /// Returns the optional model override.
380 ///
381 /// ```
382 /// use agents_kit::prompt::PromptRequest;
383 /// let request = PromptRequest::new("Hello")?;
384 /// assert_eq!(request.model(), None);
385 /// # Ok::<(), agents_kit::KitError>(())
386 /// ```
387 pub fn model(&self) -> Option<&ModelSelector> {
388 self.model.as_ref()
389 }
390
391 /// Returns the optional reasoning-effort override.
392 ///
393 /// ```
394 /// use agents_kit::prompt::PromptRequest;
395 /// let request = PromptRequest::new("Hello")?;
396 /// assert_eq!(request.reasoning_effort(), None);
397 /// # Ok::<(), agents_kit::KitError>(())
398 /// ```
399 pub fn reasoning_effort(&self) -> Option<&ReasoningEffort> {
400 self.reasoning_effort.as_ref()
401 }
402
403 /// Returns the optional personality override.
404 ///
405 /// ```
406 /// use agents_kit::prompt::PromptRequest;
407 /// let request = PromptRequest::new("Hello")?;
408 /// assert_eq!(request.personality(), None);
409 /// # Ok::<(), agents_kit::KitError>(())
410 /// ```
411 pub fn personality(&self) -> Option<&PersonalitySelector> {
412 self.personality.as_ref()
413 }
414
415 /// Returns the optional prompt-level instructions.
416 ///
417 /// ```
418 /// use agents_kit::prompt::PromptRequest;
419 /// let request = PromptRequest::new("Hello")?;
420 /// assert_eq!(request.instructions(), None);
421 /// # Ok::<(), agents_kit::KitError>(())
422 /// ```
423 pub fn instructions(&self) -> Option<&str> {
424 self.instructions.as_deref()
425 }
426
427 /// Returns the optional managed-knowledge policy override.
428 ///
429 /// ```
430 /// use agents_kit::prompt::PromptRequest;
431 /// let request = PromptRequest::new("Hello")?;
432 /// assert_eq!(request.knowledge_policy(), None);
433 /// # Ok::<(), agents_kit::KitError>(())
434 /// ```
435 pub fn knowledge_policy(&self) -> Option<&KnowledgePolicy> {
436 self.knowledge_policy.as_ref()
437 }
438
439 /// Returns the optional foreground prompt timeout.
440 ///
441 /// ```
442 /// use agents_kit::prompt::PromptRequest;
443 /// let request = PromptRequest::new("Hello")?;
444 /// assert_eq!(request.timeout(), None);
445 /// # Ok::<(), agents_kit::KitError>(())
446 /// ```
447 pub fn timeout(&self) -> Option<Duration> {
448 self.timeout
449 }
450
451 /// Returns the optional strict structured-output request.
452 ///
453 /// ```
454 /// use agents_kit::prompt::PromptRequest;
455 /// let request = PromptRequest::new("Hello")?;
456 /// assert_eq!(request.structured_output(), None);
457 /// # Ok::<(), agents_kit::KitError>(())
458 /// ```
459 pub fn structured_output(&self) -> Option<&StructuredOutputRequest> {
460 self.structured_output.as_ref()
461 }
462
463 /// Returns the optional portable tool-selection override.
464 ///
465 /// ```
466 /// use agents_kit::prompt::PromptRequest;
467 /// let request = PromptRequest::new("Hello")?;
468 /// assert_eq!(request.tool_selection(), None);
469 /// # Ok::<(), agents_kit::KitError>(())
470 /// ```
471 pub fn tool_selection(&self) -> Option<&ToolSelection> {
472 self.tool_selection.as_ref()
473 }
474}
475
476#[derive(Clone, Debug, Eq, PartialEq)]
477/// Contains the exact terminal output of one foreground prompt.
478///
479/// ```
480/// use agents_kit::prompt::{PromptResponse, UsageInfo};
481/// let response = PromptResponse::new("Done".into(), Some(UsageInfo::new(10, 2)));
482/// assert_eq!(response.output(), "Done");
483/// assert_eq!(response.usage().unwrap().input_tokens(), 10);
484/// ```
485pub struct PromptResponse {
486 output: String,
487 structured_output: Option<serde_json::Value>,
488 usage: Option<UsageInfo>,
489}
490
491impl PromptResponse {
492 /// Constructs a terminal response from output and whole-call usage.
493 ///
494 /// ```
495 /// use agents_kit::prompt::PromptResponse;
496 /// let response = PromptResponse::new("Done".into(), None);
497 /// assert_eq!(response.output(), "Done");
498 /// ```
499 pub fn new(output: String, usage: Option<UsageInfo>) -> Self {
500 Self {
501 output,
502 structured_output: None,
503 usage,
504 }
505 }
506
507 /// Constructs a terminal response from exact text, structured value, and usage.
508 #[must_use]
509 ///
510 /// ```
511 /// use agents_kit::prompt::PromptResponse;
512 /// let response = PromptResponse::from_parts("7".into(), Some(serde_json::json!(7)), None);
513 /// assert_eq!(response.structured_output(), Some(&serde_json::json!(7)));
514 /// ```
515 pub fn from_parts(
516 output: String,
517 structured_output: Option<serde_json::Value>,
518 usage: Option<UsageInfo>,
519 ) -> Self {
520 Self {
521 output,
522 structured_output,
523 usage,
524 }
525 }
526
527 /// Returns the exact terminal output.
528 ///
529 /// ```
530 /// use agents_kit::prompt::PromptResponse;
531 /// let response = PromptResponse::new(" Done\n".into(), None);
532 /// assert_eq!(response.output(), " Done\n");
533 /// ```
534 pub fn output(&self) -> &str {
535 &self.output
536 }
537
538 /// Returns whole-call usage when the backend can report it honestly.
539 ///
540 /// ```
541 /// use agents_kit::prompt::PromptResponse;
542 /// let response = PromptResponse::new("Done".into(), None);
543 /// assert_eq!(response.usage(), None);
544 /// ```
545 pub fn usage(&self) -> Option<&UsageInfo> {
546 self.usage.as_ref()
547 }
548
549 /// Returns the separately parsed strict structured value, when requested.
550 ///
551 /// ```
552 /// use agents_kit::prompt::PromptResponse;
553 /// let response = PromptResponse::new("7".into(), None);
554 /// assert_eq!(response.structured_output(), None);
555 /// ```
556 pub fn structured_output(&self) -> Option<&serde_json::Value> {
557 self.structured_output.as_ref()
558 }
559
560 /// Consumes the response into exact text, structured value, and usage.
561 #[must_use]
562 ///
563 /// ```
564 /// use agents_kit::prompt::PromptResponse;
565 /// let response = PromptResponse::from_parts("7".into(), Some(serde_json::json!(7)), None);
566 /// let (text, value, usage) = response.into_parts();
567 /// assert_eq!(text, "7");
568 /// assert_eq!(value, Some(serde_json::json!(7)));
569 /// assert_eq!(usage, None);
570 /// ```
571 pub fn into_parts(self) -> (String, Option<serde_json::Value>, Option<UsageInfo>) {
572 (self.output, self.structured_output, self.usage)
573 }
574}
575
576/// Transient text observation or the authoritative terminal prompt response.
577///
578/// Intermediate assistant text can differ from the completed response, for
579/// example when a prompt contains tool calls before its final answer. Only
580/// `Completed` confirms terminal durability; deltas are not durable history.
581#[derive(Clone, Debug, Eq, PartialEq)]
582pub enum PromptEvent {
583 /// Real source text in its observed order, excluding reasoning and tools.
584 OutputDelta {
585 /// Exact assistant text fragment.
586 text: String,
587 },
588 /// Exact terminal output, structured value and whole-prompt usage.
589 Completed(PromptResponse),
590}
591
592/// Non-Clone observation of one accepted prompt's transient output.
593///
594/// Dropping this value releases observation only. An implementation's accepted
595/// prompt remains owned by its conversation and is cancelled through
596/// [`ConversationBackend::abort`](crate::ConversationBackend::abort).
597/// Completion or error fuses the stream; premature EOF becomes one protocol
598/// error. A caller source owns bounded production and accepted-work cleanup.
599///
600/// ```
601/// use agents_kit::prompt::{PromptEvent, PromptResponse, PromptStream};
602/// let runtime = tokio::runtime::Builder::new_current_thread().build()?;
603/// runtime.block_on(async {
604/// let mut output = PromptStream::from_stream(futures_util::stream::iter([
605/// Ok(PromptEvent::OutputDelta { text: "ready".into() }),
606/// Ok(PromptEvent::Completed(PromptResponse::new("done".into(), None))),
607/// ]));
608/// assert!(matches!(output.recv().await.transpose()?, Some(PromptEvent::OutputDelta { .. })));
609/// let Some(PromptEvent::Completed(response)) = output.recv().await.transpose()? else { panic!("terminal response"); };
610/// assert_eq!(response.output(), "done");
611/// assert!(output.recv().await.is_none());
612/// Ok::<(), agents_kit::KitError>(())
613/// })?;
614/// # Ok::<(), Box<dyn std::error::Error>>(())
615/// ```
616pub struct PromptStream {
617 source: Option<PromptSource>,
618}
619
620enum PromptSource {
621 Caller(Pin<Box<dyn Stream<Item = Result<PromptEvent, KitError>> + Send + 'static>>),
622 Buffered {
623 admission: PromptAdmission,
624 fragments: mpsc::Receiver<String>,
625 terminal: oneshot::Receiver<Result<PromptResponse, KitError>>,
626 },
627}
628
629enum PromptAdmission {
630 Pending(oneshot::Receiver<()>),
631 Accepted,
632 Rejected,
633}
634
635pub(crate) struct PromptOutput {
636 state: Mutex<PromptOutputState>,
637}
638
639enum PromptOutputState {
640 Open {
641 admission: Option<oneshot::Sender<()>>,
642 fragments: mpsc::Sender<String>,
643 terminal: oneshot::Sender<Result<PromptResponse, KitError>>,
644 },
645 Closed,
646}
647
648impl PromptOutput {
649 const FRAGMENTS: usize = 64;
650 const FRAGMENT_BYTES: usize = 16 * 1024;
651
652 pub(crate) fn channel() -> (Arc<Self>, PromptStream) {
653 let (output, stream) = Self::gated_channel();
654 output.admit();
655 (output, stream)
656 }
657
658 pub(crate) fn gated_channel() -> (Arc<Self>, PromptStream) {
659 let (admission, admitted) = oneshot::channel();
660 let (fragments, receiver) = mpsc::channel(Self::FRAGMENTS);
661 let (terminal, completion) = oneshot::channel();
662 (
663 Arc::new(Self {
664 state: Mutex::new(PromptOutputState::Open {
665 admission: Some(admission),
666 fragments,
667 terminal,
668 }),
669 }),
670 PromptStream {
671 source: Some(PromptSource::Buffered {
672 admission: PromptAdmission::Pending(admitted),
673 fragments: receiver,
674 terminal: completion,
675 }),
676 },
677 )
678 }
679
680 pub(crate) fn admit(&self) {
681 if let PromptOutputState::Open { admission, .. } = &mut *self
682 .state
683 .lock()
684 .unwrap_or_else(|poisoned| poisoned.into_inner())
685 && let Some(admission) = admission.take()
686 {
687 let _ = admission.send(());
688 }
689 }
690
691 pub(crate) fn fragment(&self, mut text: &str) {
692 let mut state = self
693 .state
694 .lock()
695 .unwrap_or_else(|poisoned| poisoned.into_inner());
696 let PromptOutputState::Open { fragments, .. } = &*state else {
697 return;
698 };
699 loop {
700 let mut end = text.len().min(Self::FRAGMENT_BYTES);
701 while !text.is_char_boundary(end) {
702 end -= 1;
703 }
704 if let Err(error) = fragments.try_send(text[..end].to_owned()) {
705 let previous = std::mem::replace(&mut *state, PromptOutputState::Closed);
706 if let PromptOutputState::Open { terminal, .. } = previous {
707 let failure = match error {
708 mpsc::error::TrySendError::Full(_) => {
709 KitError::integration("prompt output observation overflow")
710 .with_detail(crate::error::ErrorDetail::ResponseLimit {
711 limit: Self::FRAGMENTS,
712 })
713 .with_context(
714 crate::error::Operation::Prompt,
715 crate::error::ErrorPhase::Publication,
716 )
717 }
718 mpsc::error::TrySendError::Closed(_) => {
719 KitError::cancelled("prompt output observer was dropped")
720 }
721 };
722 let _ = terminal.send(Err(failure));
723 }
724 return;
725 }
726 text = &text[end..];
727 if text.is_empty() {
728 return;
729 }
730 }
731 }
732
733 pub(crate) fn finish(&self, result: Result<PromptResponse, KitError>) {
734 let previous = std::mem::replace(
735 &mut *self
736 .state
737 .lock()
738 .unwrap_or_else(|poisoned| poisoned.into_inner()),
739 PromptOutputState::Closed,
740 );
741 if let PromptOutputState::Open { terminal, .. } = previous {
742 let _ = terminal.send(result);
743 }
744 }
745}
746
747impl PromptStream {
748 /// Adapts a caller's concrete value stream without normalizing a backend.
749 ///
750 /// The source may be non-Unpin. It must produce ordered real text and one
751 /// authoritative completed response or error, retain accepted work after
752 /// observer drop, and keep pending output bounded.
753 ///
754 /// See [`PromptStream`] for a complete caller-stream example.
755 pub fn from_stream<S>(stream: S) -> Self
756 where
757 S: Stream<Item = Result<PromptEvent, KitError>> + Send + 'static,
758 {
759 Self {
760 source: Some(PromptSource::Caller(Box::pin(stream))),
761 }
762 }
763
764 /// Receives the next observation, or None after its single terminal item.
765 ///
766 /// Cancelling this read does not abort the accepted prompt. A premature
767 /// source EOF returns Protocol once rather than claiming successful output.
768 ///
769 /// # Errors
770 /// A source failure is yielded once with prompt/publication context. EOF
771 /// before completion yields Protocol with IncompleteStream; terminal results
772 /// fuse subsequent reads to None.
773 ///
774 /// ```
775 /// use agents_kit::{error::ErrorKind, prompt::PromptStream};
776 /// # tokio::runtime::Builder::new_current_thread().build()?.block_on(async {
777 /// let mut output = PromptStream::from_stream(futures_util::stream::empty());
778 /// let error = output.recv().await.expect("one terminal error").unwrap_err();
779 /// assert_eq!(error.kind(), ErrorKind::Protocol);
780 /// assert!(output.recv().await.is_none());
781 /// # });
782 /// # Ok::<(), Box<dyn std::error::Error>>(())
783 /// ```
784 pub async fn recv(&mut self) -> Option<Result<PromptEvent, KitError>> {
785 let item = match self.source.as_mut()? {
786 PromptSource::Caller(source) => source.next().await,
787 PromptSource::Buffered {
788 admission,
789 fragments,
790 terminal,
791 } => {
792 if let PromptAdmission::Pending(receiver) = admission {
793 *admission = if receiver.await.is_ok() {
794 PromptAdmission::Accepted
795 } else {
796 PromptAdmission::Rejected
797 };
798 }
799 if matches!(admission, PromptAdmission::Accepted)
800 && let Some(text) = fragments.recv().await
801 {
802 Some(Ok(PromptEvent::OutputDelta { text }))
803 } else {
804 terminal
805 .await
806 .ok()
807 .map(|result| result.map(PromptEvent::Completed))
808 }
809 }
810 };
811 match item {
812 Some(Ok(PromptEvent::OutputDelta { text })) => {
813 Some(Ok(PromptEvent::OutputDelta { text }))
814 }
815 Some(terminal) => {
816 self.source.take();
817 Some(terminal.map_err(|error| {
818 error
819 .at_operation(crate::error::Operation::Prompt)
820 .in_phase(crate::error::ErrorPhase::Publication)
821 }))
822 }
823 None => {
824 self.source.take();
825 Some(Err(KitError::protocol(
826 "prompt output ended before its terminal response",
827 )
828 .with_context(
829 crate::error::Operation::Prompt,
830 crate::error::ErrorPhase::Publication,
831 )
832 .with_detail(crate::error::ErrorDetail::IncompleteStream)))
833 }
834 }
835 }
836}
837
838#[derive(Clone, Copy, Debug, Eq, PartialEq)]
839/// Reports token usage covering every native model call in one prompt.
840///
841/// ```
842/// use agents_kit::prompt::UsageInfo;
843/// let usage = UsageInfo::new(10, 2);
844/// assert_eq!(usage.input_tokens(), 10);
845/// assert_eq!(usage.output_tokens(), 2);
846/// ```
847pub struct UsageInfo {
848 input_tokens: u64,
849 output_tokens: u64,
850}
851
852impl UsageInfo {
853 /// Constructs whole-call input and output token counts.
854 ///
855 /// ```
856 /// use agents_kit::prompt::UsageInfo;
857 /// let usage = UsageInfo::new(10, 2);
858 /// assert_eq!(usage.output_tokens(), 2);
859 /// ```
860 pub fn new(input_tokens: u64, output_tokens: u64) -> Self {
861 Self {
862 input_tokens,
863 output_tokens,
864 }
865 }
866
867 /// Returns the total input-token count.
868 ///
869 /// ```
870 /// use agents_kit::prompt::UsageInfo;
871 /// assert_eq!(UsageInfo::new(10, 2).input_tokens(), 10);
872 /// ```
873 pub fn input_tokens(&self) -> u64 {
874 self.input_tokens
875 }
876
877 /// Returns the total output-token count.
878 ///
879 /// ```
880 /// use agents_kit::prompt::UsageInfo;
881 /// assert_eq!(UsageInfo::new(10, 2).output_tokens(), 2);
882 /// ```
883 pub fn output_tokens(&self) -> u64 {
884 self.output_tokens
885 }
886}