1use std::{fmt, sync::Arc, time::Duration};
2
3use runifold_core::{CancellationToken, Instant, InvocationId, RunContext, RunId};
4use runifold_model::{ArtifactScope, ArtifactStore};
5
6#[derive(Clone)]
8pub struct ToolContext {
9 invocation_id: InvocationId,
10 run_id: RunId,
11 deadline: Option<Instant>,
12 cancellation: CancellationToken,
13 artifact_store: Option<Arc<dyn ArtifactStore>>,
14 artifact_scope: Option<ArtifactScope>,
15}
16
17impl ToolContext {
18 pub(crate) fn for_run(run: &RunContext) -> Self {
19 Self {
20 invocation_id: InvocationId::new(),
21 run_id: run.run_id(),
22 deadline: run.deadline(),
23 cancellation: run.cancellation().child_token(),
24 artifact_store: None,
25 artifact_scope: None,
26 }
27 }
28
29 pub const fn invocation_id(&self) -> InvocationId {
31 self.invocation_id
32 }
33
34 pub const fn run_id(&self) -> RunId {
36 self.run_id
37 }
38
39 pub const fn deadline(&self) -> Option<Instant> {
41 self.deadline
42 }
43
44 pub fn remaining(&self) -> Option<Duration> {
46 self.deadline
47 .map(|deadline| deadline.saturating_duration_since(Instant::now()))
48 }
49
50 pub const fn cancellation(&self) -> &CancellationToken {
52 &self.cancellation
53 }
54
55 pub(crate) fn with_artifact_store(
56 mut self,
57 scope: Option<ArtifactScope>,
58 store: Option<Arc<dyn ArtifactStore>>,
59 ) -> Self {
60 self.artifact_scope = scope;
61 self.artifact_store = store;
62 self
63 }
64
65 pub fn artifact_store(&self) -> Option<&Arc<dyn ArtifactStore>> {
68 self.artifact_store.as_ref()
69 }
70
71 pub const fn artifact_scope(&self) -> Option<&ArtifactScope> {
73 self.artifact_scope.as_ref()
74 }
75}
76
77impl fmt::Debug for ToolContext {
78 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
79 formatter
80 .debug_struct("ToolContext")
81 .field("invocation_id", &self.invocation_id)
82 .field("run_id", &self.run_id)
83 .field("deadline", &self.deadline)
84 .field("artifact_store", &self.artifact_store.is_some())
85 .field("artifact_scope", &self.artifact_scope)
86 .finish_non_exhaustive()
87 }
88}