1use serde::{Deserialize, Serialize};
2
3use crate::events::{ThreadId, TurnId};
4use crate::policy_mode::{PolicyDecision, PolicyMode};
5use crate::tools::{ToolCall, ToolExecutionContext};
6
7pub use crate::extension::{ContextPlannerId, ContextProviderId, PolicyContributorId};
8
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
10pub enum ContextBlockKind {
11 Instruction,
12 RepositoryFact,
13 Memory,
14 Knowledge,
15 RetrievedDocument,
16 Environment,
17 ToolAvailability,
18 SafetyPolicy,
19 TaskMetadata,
20 PriorSummary,
21 EntrypointHint,
22 RetrievalHint,
23 Other(String),
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27pub struct ContextBlock {
28 pub id: String,
29 pub kind: ContextBlockKind,
30 pub text: String,
31 pub priority: i32,
32 pub token_estimate: Option<u32>,
33 pub metadata: serde_json::Value,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
37pub struct ContextQuery {
38 pub thread_id: ThreadId,
39 pub turn_id: TurnId,
40 pub prompt: String,
41 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub workspace: Option<String>,
43 pub token_budget: Option<u32>,
44}
45
46#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
47pub struct ContextPlan {
48 pub blocks: Vec<ContextBlock>,
49}
50
51#[async_trait::async_trait]
52pub trait ContextProvider: Send + Sync + 'static {
53 fn id(&self) -> ContextProviderId;
54
55 async fn blocks(&self, query: &ContextQuery) -> anyhow::Result<Vec<ContextBlock>>;
56}
57
58#[async_trait::async_trait]
59pub trait ContextPlanner: Send + Sync + 'static {
60 fn id(&self) -> ContextPlannerId;
61
62 async fn plan(
63 &self,
64 query: &ContextQuery,
65 provider_blocks: Vec<ContextBlock>,
66 ) -> anyhow::Result<ContextPlan>;
67}
68
69#[derive(Debug, Clone)]
70pub struct PolicyReview {
71 pub call: ToolCall,
72 pub mode: PolicyMode,
73 pub context: ToolExecutionContext,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
77#[serde(rename_all = "snake_case")]
78pub enum PolicyContribution {
79 Abstain,
80 Allow { reason: Option<String> },
81 RequireApproval { reason: Option<String> },
82 Deny { reason: String },
83}
84
85#[async_trait::async_trait]
86pub trait PolicyContributor: Send + Sync + 'static {
87 fn id(&self) -> PolicyContributorId;
88
89 async fn review_tool(&self, review: PolicyReview) -> anyhow::Result<PolicyContribution>;
90}
91
92pub trait PolicyGate: Send + Sync + 'static {
93 fn decide(
94 &self,
95 call: &ToolCall,
96 mode: PolicyMode,
97 context: &ToolExecutionContext,
98 ) -> PolicyDecision;
99}
100
101pub struct SimpleContextPlanner;
102
103#[async_trait::async_trait]
104impl ContextPlanner for SimpleContextPlanner {
105 fn id(&self) -> ContextPlannerId {
106 "default".to_string()
107 }
108
109 async fn plan(
110 &self,
111 _query: &ContextQuery,
112 mut provider_blocks: Vec<ContextBlock>,
113 ) -> anyhow::Result<ContextPlan> {
114 provider_blocks.sort_by_key(|block| std::cmp::Reverse(block.priority));
115 Ok(ContextPlan {
116 blocks: provider_blocks,
117 })
118 }
119}