Skip to main content

tea_context/providers/
workspace.rs

1use std::str::FromStr;
2
3use crate::{
4    BudgetBehavior, CacheScope, ContextError, ContextErrorCode, ContextProvider,
5    ContextProviderFuture, ContextProviderId, ContextRequest, PromptAuthority, PromptModule,
6    PromptModuleId, PromptPriority, PromptProvenance, PromptSegment, PromptSegmentId, TrustLevel,
7};
8
9/// Maximum caller-supplied workspace instruction documents.
10pub const MAX_WORKSPACE_INSTRUCTIONS: usize = 128;
11
12/// One caller-loaded workspace instruction document.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct WorkspaceInstruction {
15    id: PromptSegmentId,
16    content: String,
17    locator: String,
18    trust: TrustLevel,
19}
20
21impl WorkspaceInstruction {
22    /// Creates one explicit workspace document snapshot.
23    ///
24    /// # Errors
25    ///
26    /// Returns an error for invalid content or source locator.
27    pub fn new(
28        id: PromptSegmentId,
29        content: impl Into<String>,
30        locator: impl Into<String>,
31        trust: TrustLevel,
32    ) -> Result<Self, ContextError> {
33        let content = content.into();
34        let locator = locator.into();
35        if locator.is_empty() || locator.len() > 2048 || locator.chars().any(char::is_control) {
36            return Err(ContextError::new(
37                ContextErrorCode::InvalidValue,
38                "workspace instruction locator is invalid",
39            ));
40        }
41        // Reuse segment construction as the authoritative content bound.
42        PromptSegment::new(
43            id.clone(),
44            content.clone(),
45            PromptProvenance::new(
46                ContextProviderId::from_str("builtin.workspace_instructions")
47                    .map_err(value_error)?,
48                "workspace_file",
49                Some(locator.clone()),
50            )
51            .map_err(value_error)?,
52            trust,
53            CacheScope::Session,
54            BudgetBehavior::Omit,
55        )
56        .map_err(value_error)?;
57        Ok(Self {
58            id,
59            content,
60            locator,
61            trust,
62        })
63    }
64}
65
66/// Provider over caller-supplied workspace instruction snapshots.
67#[derive(Debug, Clone)]
68pub struct WorkspaceInstructionProvider {
69    id: ContextProviderId,
70    instructions: Vec<WorkspaceInstruction>,
71}
72
73impl WorkspaceInstructionProvider {
74    /// Creates a canonical instruction provider.
75    ///
76    /// # Errors
77    ///
78    /// Returns an error for too many or duplicate instruction identities.
79    pub fn new(mut instructions: Vec<WorkspaceInstruction>) -> Result<Self, ContextError> {
80        if instructions.len() > MAX_WORKSPACE_INSTRUCTIONS {
81            return Err(ContextError::new(
82                ContextErrorCode::BoundsExceeded,
83                "workspace instruction collection is too large",
84            ));
85        }
86        instructions.sort_by(|left, right| left.id.cmp(&right.id));
87        if instructions
88            .windows(2)
89            .any(|items| items[0].id == items[1].id)
90        {
91            return Err(ContextError::new(
92                ContextErrorCode::DuplicateIdentity,
93                "workspace instruction ID is duplicated",
94            ));
95        }
96        Ok(Self {
97            id: ContextProviderId::from_str("builtin.workspace_instructions")
98                .map_err(value_error)?,
99            instructions,
100        })
101    }
102}
103
104impl ContextProvider for WorkspaceInstructionProvider {
105    fn id(&self) -> &ContextProviderId {
106        &self.id
107    }
108
109    fn provide(&self, _request: ContextRequest) -> ContextProviderFuture<'_> {
110        let id = self.id.clone();
111        let instructions = self.instructions.clone();
112        Box::pin(async move {
113            if instructions.is_empty() {
114                return Ok(Vec::new());
115            }
116            let segments = instructions
117                .into_iter()
118                .map(|instruction| {
119                    PromptSegment::new(
120                        instruction.id,
121                        instruction.content,
122                        PromptProvenance::new(
123                            id.clone(),
124                            "workspace_file",
125                            Some(instruction.locator),
126                        )
127                        .map_err(value_error)?,
128                        instruction.trust,
129                        CacheScope::Session,
130                        BudgetBehavior::Omit,
131                    )
132                    .map_err(value_error)
133                })
134                .collect::<Result<Vec<_>, ContextError>>()?;
135            Ok(vec![
136                PromptModule::new(
137                    PromptModuleId::from_str("workspace.instructions").map_err(value_error)?,
138                    PromptAuthority::Workspace,
139                    PromptPriority::new(0),
140                    segments,
141                )
142                .map_err(value_error)?,
143            ])
144        })
145    }
146}
147
148fn value_error(error: impl std::fmt::Display) -> ContextError {
149    ContextError::new(ContextErrorCode::InvalidValue, error.to_string())
150}