1use std::{path::PathBuf, sync::Arc};
4
5use anyhow::Result;
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9use crate::{ClientCapabilities, ContentBlock, Error, ProtocolVersion, SessionId};
10
11pub trait Agent {
12 fn initialize(
13 &self,
14 arguments: InitializeRequest,
15 ) -> impl Future<Output = Result<InitializeResponse, Error>>;
16
17 fn authenticate(
18 &self,
19 arguments: AuthenticateRequest,
20 ) -> impl Future<Output = Result<(), Error>>;
21
22 fn new_session(
23 &self,
24 arguments: NewSessionRequest,
25 ) -> impl Future<Output = Result<NewSessionResponse, Error>>;
26
27 fn load_session(
28 &self,
29 arguments: LoadSessionRequest,
30 ) -> impl Future<Output = Result<(), Error>>;
31
32 fn prompt(
33 &self,
34 arguments: PromptRequest,
35 ) -> impl Future<Output = Result<PromptResponse, Error>>;
36
37 fn cancel(&self, args: CancelNotification) -> impl Future<Output = Result<(), Error>>;
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
43#[serde(rename_all = "camelCase")]
44pub struct InitializeRequest {
45 pub protocol_version: ProtocolVersion,
47 #[serde(default)]
49 pub client_capabilities: ClientCapabilities,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
53#[serde(rename_all = "camelCase")]
54pub struct InitializeResponse {
55 pub protocol_version: ProtocolVersion,
60 #[serde(default)]
62 pub agent_capabilities: AgentCapabilities,
63 #[serde(default)]
65 pub auth_methods: Vec<AuthMethod>,
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
71#[serde(rename_all = "camelCase")]
72pub struct AuthenticateRequest {
73 pub method_id: AuthMethodId,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq, Hash)]
77#[serde(transparent)]
78pub struct AuthMethodId(pub Arc<str>);
79
80#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
81#[serde(rename_all = "camelCase")]
82pub struct AuthMethod {
83 pub id: AuthMethodId,
84 pub name: String,
85 pub description: Option<String>,
86}
87
88#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
91#[serde(rename_all = "camelCase")]
92pub struct NewSessionRequest {
93 pub mcp_servers: Vec<McpServer>,
94 pub cwd: PathBuf,
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
98#[serde(rename_all = "camelCase")]
99pub struct NewSessionResponse {
100 pub session_id: SessionId,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
106#[serde(rename_all = "camelCase")]
107pub struct LoadSessionRequest {
108 pub mcp_servers: Vec<McpServer>,
109 pub cwd: PathBuf,
110 pub session_id: SessionId,
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
116#[serde(rename_all = "camelCase")]
117pub struct McpServer {
118 pub name: String,
119 pub command: PathBuf,
120 pub args: Vec<String>,
121 pub env: Vec<EnvVariable>,
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
125#[serde(rename_all = "camelCase")]
126pub struct EnvVariable {
127 pub name: String,
128 pub value: String,
129}
130
131#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
134#[serde(rename_all = "camelCase")]
135pub struct PromptRequest {
136 pub session_id: SessionId,
137 pub prompt: Vec<ContentBlock>,
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
141#[serde(rename_all = "camelCase")]
142pub struct PromptResponse {
143 pub stop_reason: StopReason,
144}
145
146#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize, JsonSchema)]
147#[serde(rename_all = "snake_case")]
148pub enum StopReason {
149 EndTurn,
151 MaxTokens,
153 MaxTurnRequests,
156 Refusal,
160 Cancelled,
162}
163
164#[derive(Default, Debug, Clone, Serialize, Deserialize, JsonSchema)]
167#[serde(rename_all = "camelCase")]
168pub struct AgentCapabilities {
169 #[serde(default)]
171 load_session: bool,
172}
173
174#[derive(Debug, Clone, Serialize, Deserialize)]
177pub struct AgentMethodNames {
178 pub initialize: &'static str,
179 pub authenticate: &'static str,
180 pub session_new: &'static str,
181 pub session_load: &'static str,
182 pub session_prompt: &'static str,
183 pub session_cancel: &'static str,
184}
185
186pub const AGENT_METHOD_NAMES: AgentMethodNames = AgentMethodNames {
187 initialize: INITIALIZE_METHOD_NAME,
188 authenticate: AUTHENTICATE_METHOD_NAME,
189 session_new: SESSION_NEW_METHOD_NAME,
190 session_load: SESSION_LOAD_METHOD_NAME,
191 session_prompt: SESSION_PROMPT_METHOD_NAME,
192 session_cancel: SESSION_CANCEL_METHOD_NAME,
193};
194
195pub const INITIALIZE_METHOD_NAME: &str = "initialize";
196pub const AUTHENTICATE_METHOD_NAME: &str = "authenticate";
197pub const SESSION_NEW_METHOD_NAME: &str = "session/new";
198pub const SESSION_LOAD_METHOD_NAME: &str = "session/load";
199pub const SESSION_PROMPT_METHOD_NAME: &str = "session/prompt";
200pub const SESSION_CANCEL_METHOD_NAME: &str = "session/cancel";
201
202#[derive(Debug, Serialize, Deserialize, JsonSchema)]
204#[serde(untagged)]
205pub enum ClientRequest {
206 InitializeRequest(InitializeRequest),
207 AuthenticateRequest(AuthenticateRequest),
208 NewSessionRequest(NewSessionRequest),
209 LoadSessionRequest(LoadSessionRequest),
210 PromptRequest(PromptRequest),
211}
212
213#[derive(Debug, Serialize, Deserialize, JsonSchema)]
215#[serde(untagged)]
216pub enum AgentResponse {
217 InitializeResponse(InitializeResponse),
218 AuthenticateResponse,
219 NewSessionResponse(NewSessionResponse),
220 LoadSessionResponse,
221 PromptResponse(PromptResponse),
222}
223
224#[derive(Debug, Serialize, Deserialize, JsonSchema)]
226#[serde(untagged)]
227pub enum ClientNotification {
228 CancelNotification(CancelNotification),
229}
230
231#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
232#[serde(rename_all = "camelCase")]
233pub struct CancelNotification {
234 pub session_id: SessionId,
235}