@ondeinference / ed
Code Commits Issues Pulls Cloud Agent CI/CD
main
main release/1.3.1
rs 209 lines 6.75 KB
Raw
1 //! `serve` over an in-memory channel: the agent and its client in one process, no stdio.
2
3 use std::sync::Arc;
4
5 use agent_client_protocol::schema::ProtocolVersion;
6 use agent_client_protocol::schema::v1::InitializeRequest;
7 use agent_client_protocol::{Agent, Channel, Client, ConnectionTo};
8 use ed_acp::{AgentInfo, LlmEnv, Profile, PromptCtx, ServeOptions, Toolset};
9
10 struct Embedded;
11
12 impl Profile for Embedded {
13 fn info(&self) -> AgentInfo {
14 AgentInfo {
15 name: "embedded-agent",
16 display_name: "Embedded",
17 version: "0.0.1",
18 env: LlmEnv {
19 prefix: "ED_ACP_IN_PROCESS_TEST",
20 dir_name: "ed-acp-in-process-test",
21 default_model: "test-model",
22 },
23 }
24 }
25
26 fn system_prompt(&self, _ctx: &PromptCtx<'_>) -> String {
27 String::new()
28 }
29
30 fn toolsets(&self) -> Vec<Arc<dyn Toolset>> {
31 Vec::new()
32 }
33 }
34
35 #[tokio::test]
36 async fn initialize_over_an_in_memory_channel() {
37 let (agent_end, client_end) = Channel::duplex();
38 let opts = ServeOptions {
39 yolo: false,
40 surface: "app",
41 };
42 let server = tokio::spawn(ed_acp::serve(Arc::new(Embedded), opts, agent_end));
43
44 let info = Client
45 .builder()
46 .connect_with(client_end, |conn: ConnectionTo<Agent>| async move {
47 conn.send_request(InitializeRequest::new(ProtocolVersion::V1))
48 .block_task()
49 .await
50 })
51 .await
52 .expect("initialize");
53
54 let agent = info.agent_info.expect("agentInfo");
55 assert_eq!(agent.name, "embedded-agent");
56 assert_eq!(agent.version, "0.0.1");
57 let caps = info.agent_capabilities;
58 assert!(caps.load_session);
59
60 // The client hung up when `connect_with` returned, so the server stops on its own.
61 tokio::time::timeout(std::time::Duration::from_secs(5), server)
62 .await
63 .expect("server stops after the client disconnects")
64 .expect("server task")
65 .ok();
66 }
67
68 mod provider_auth {
69 use std::sync::atomic::{AtomicUsize, Ordering};
70 use std::sync::{Arc, Mutex};
71
72 use agent_client_protocol::schema::ProtocolVersion;
73 use agent_client_protocol::schema::v1::{
74 AuthMethod, AuthMethodAgent, AuthenticateRequest, InitializeRequest, LogoutRequest,
75 NewSessionRequest,
76 };
77 use agent_client_protocol::{Agent, Channel, Client, ConnectionTo};
78 use async_trait::async_trait;
79 use ed_acp::{
80 AgentInfo, AuthCtx, AuthProvider, LlmConfig, LlmEnv, Profile, PromptCtx, Provider,
81 ServeOptions, Toolset,
82 };
83
84 /// A provider whose "session" is a string it holds in memory.
85 #[derive(Default)]
86 struct State {
87 token: Mutex<Option<String>>,
88 authenticated: AtomicUsize,
89 }
90
91 struct Fake(Arc<State>);
92
93 #[async_trait]
94 impl AuthProvider for Fake {
95 fn method(&self) -> AuthMethod {
96 AuthMethod::Agent(AuthMethodAgent::new("fake", "Sign in to Fake"))
97 }
98
99 fn llm_config(&self) -> Option<LlmConfig> {
100 self.0.token.lock().unwrap().clone().map(|token| LlmConfig {
101 provider: Provider::OpenAi,
102 // Nothing listens here: listing models fails fast and the default is offered.
103 base_url: "http://127.0.0.1:1".into(),
104 api_key: Some(token),
105 model: "fake-model".into(),
106 })
107 }
108
109 async fn authenticate(&self, ctx: AuthCtx<'_>) -> agent_client_protocol::Result<()> {
110 assert!(!ctx.client_offers_url_elicitation());
111 self.0.authenticated.fetch_add(1, Ordering::SeqCst);
112 *self.0.token.lock().unwrap() = Some("token".into());
113 Ok(())
114 }
115
116 async fn logout(&self) -> anyhow::Result<()> {
117 *self.0.token.lock().unwrap() = None;
118 Ok(())
119 }
120 }
121
122 struct WithAuth(Arc<State>);
123
124 impl Profile for WithAuth {
125 fn info(&self) -> AgentInfo {
126 AgentInfo {
127 name: "provider-agent",
128 display_name: "Provider",
129 version: "0.0.1",
130 env: LlmEnv {
131 prefix: "ED_ACP_PROVIDER_TEST",
132 dir_name: "ed-acp-provider-test",
133 default_model: "fake-model",
134 },
135 }
136 }
137
138 fn system_prompt(&self, _ctx: &PromptCtx<'_>) -> String {
139 String::new()
140 }
141
142 fn toolsets(&self) -> Vec<Arc<dyn Toolset>> {
143 Vec::new()
144 }
145
146 fn auth(&self) -> Option<Arc<dyn AuthProvider>> {
147 Some(Arc::new(Fake(self.0.clone())))
148 }
149 }
150
151 #[tokio::test]
152 async fn the_provider_owns_initialize_sessions_authenticate_and_logout() {
153 let fake = Arc::new(State::default());
154 let (agent_end, client_end) = Channel::duplex();
155 let opts = ServeOptions {
156 yolo: false,
157 surface: "app",
158 };
159 tokio::spawn(ed_acp::serve(
160 Arc::new(WithAuth(fake.clone())),
161 opts,
162 agent_end,
163 ));
164 let cwd = std::env::temp_dir();
165
166 Client
167 .builder()
168 .connect_with(client_end, |conn: ConnectionTo<Agent>| async move {
169 // The client cannot run terminal auth, and the method is offered anyway.
170 let init = conn
171 .send_request(InitializeRequest::new(ProtocolVersion::V1))
172 .block_task()
173 .await?;
174 assert_eq!(init.auth_methods.len(), 1);
175 assert_eq!(&*init.auth_methods[0].id().0, "fake");
176
177 let signed_out = conn
178 .send_request(NewSessionRequest::new(cwd.clone()))
179 .block_task()
180 .await;
181 assert!(signed_out.is_err(), "no session while signed out");
182
183 let wrong = conn
184 .send_request(AuthenticateRequest::new("terminal-setup"))
185 .block_task()
186 .await;
187 assert!(wrong.is_err(), "only the provider's method is accepted");
188
189 conn.send_request(AuthenticateRequest::new("fake"))
190 .block_task()
191 .await?;
192 conn.send_request(NewSessionRequest::new(cwd.clone()))
193 .block_task()
194 .await?;
195
196 conn.send_request(LogoutRequest::new()).block_task().await?;
197 let after = conn
198 .send_request(NewSessionRequest::new(cwd.clone()))
199 .block_task()
200 .await;
201 assert!(after.is_err(), "signed out again");
202 Ok(())
203 })
204 .await
205 .expect("conversation");
206
207 assert_eq!(fake.authenticated.load(Ordering::SeqCst), 1);
208 }
209 }