@ondeinference / ed / commits / 5969eea

ed-agent: streaming, history restore, one-shot generate, SessionStore

Ed::stream gives plain chat a stream of text deltas, and Ed::run forwards deltas to the new EventSink::text_delta (a no-op by default, so existing sinks still compile). Adds restore_history, generate, the system prompt and sampling setters, and the SessionStore trait as the seam for host storage. The FFI exposes the same calls, plus stream_message with cancel_stream and an app-id constructor. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-Authored-By: siGit Code v1.6.3 <noreply@sigit.si>

Seto Elkahfi committed Oct 9, 2026 at 00:12 UTC 5969eea2c98108155e530f2ea93ee303ad0051c8

9 files changed +595 -20

Cargo.lock

@@ -1731,6 +1731,7 @@ name = "ed-agent"

1731 version = "1.0.1"

1732 dependencies = [

1733 "async-trait",

1734 + "futures-core",

1735 "log",

1736 "onde",

1737 "serde",

@@ -1750,6 +1751,7 @@ dependencies = [

1751 "serde_json",

1752 "smbcloud-gresiq-sdk",

1753 "thiserror 2.0.21",

1754 + "tokio",

1755 "uniffi",

1756 ]

1757

crates/ed-agent-ffi/Cargo.toml

@@ -18,6 +18,7 @@ crate-type = ["lib", "cdylib", "staticlib"]

18 ed-agent.workspace = true

19 onde.workspace = true

20 async-trait.workspace = true

21 +tokio.workspace = true

22 serde_json.workspace = true

23 thiserror.workspace = true

24 uniffi.workspace = true

crates/ed-agent-ffi/src/lib.rs

+137 -4

@@ -216,6 +216,20 @@ impl From<ChatMessage> for FfiChatMessage {

216 }

217 }

218

219 +impl From<FfiChatMessage> for ChatMessage {

220 + fn from(value: FfiChatMessage) -> Self {

221 + let role = match value.role {

222 + FfiChatRole::System => ed_agent::ChatRole::System,

223 + FfiChatRole::User => ed_agent::ChatRole::User,

224 + FfiChatRole::Assistant => ed_agent::ChatRole::Assistant,

225 + };

226 + Self {

227 + role,

228 + content: value.content,

229 + }

230 + }

231 +}

232 +

233 #[derive(Debug, Clone, uniffi::Record)]

234 pub struct FfiSamplingConfig {

235 pub temperature: Option<f64>,

@@ -333,6 +347,8 @@ pub trait FfiApprovalHandler: Send + Sync {

347 #[uniffi::export(callback_interface)]

348 pub trait FfiEventListener: Send + Sync {

349 fn status_changed(&self, update: FfiStatusUpdate);

350 + /// A piece of the reply as it is generated, during `run`.

351 + fn text_delta(&self, delta: String);

352 fn tool_requested(&self, call: FfiToolCall);

353 fn approval_requested(&self, request: FfiApprovalRequest);

354 fn tool_started(&self, call: FfiToolCall);

@@ -341,6 +357,12 @@ pub trait FfiEventListener: Send + Sync {

357 fn warning(&self, message: String);

358 }

359

360 +/// Receives a streamed reply from `stream_message`.

361 +#[uniffi::export(callback_interface)]

362 +pub trait FfiStreamListener: Send + Sync {

363 + fn on_delta(&self, delta: String);

364 +}

365 +

366 struct ForeignExecutor(Arc<dyn FfiToolExecutor>);

367

368 #[async_trait]

@@ -378,6 +400,10 @@ impl EventSink for FfiSink {

400 });

401 }

402

403 + fn text_delta(&self, delta: &str) {

404 + self.0.text_delta(delta.to_owned());

405 + }

406 +

407 fn tool_requested(&self, call: &ToolCall) {

408 self.0.tool_requested(call.into());

409 }

@@ -410,20 +436,21 @@ impl EventSink for FfiSink {

436 #[derive(uniffi::Object)]

437 pub struct FfiEdAgent {

438 inner: Ed<FfiSink>,

439 + /// Ends the stream in progress, if any.

440 + stream_cancel: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,

441 }

442

415 -#[uniffi::export(async_runtime = "tokio")]

443 impl FfiEdAgent {

417 - #[uniffi::constructor]

418 - pub fn new(

444 + fn build(

445 executor: Box<dyn FfiToolExecutor>,

446 approvals: Box<dyn FfiApprovalHandler>,

447 events: Box<dyn FfiEventListener>,

448 config: Option<FfiAgentConfig>,

449 + app_id: Option<String>,

450 ) -> Arc<Self> {

451 let config = config.unwrap_or_default();

452 Arc::new(Self {

426 - inner: Ed::with_agent(

453 + inner: Ed::with_agent_and_app_id(

454 FfiSink(Arc::from(events)),

455 Arc::new(ForeignExecutor(Arc::from(executor))),

456 Arc::new(ForeignApprovals(Arc::from(approvals))),

@@ -432,9 +459,36 @@ impl FfiEdAgent {

459 max_tool_output_chars: usize::try_from(config.max_tool_output_chars)

460 .unwrap_or(usize::MAX),

461 },

462 + app_id,

463 ),

464 + stream_cancel: std::sync::Mutex::new(None),

465 })

466 }

467 +}

468 +

469 +#[uniffi::export(async_runtime = "tokio")]

470 +impl FfiEdAgent {

471 + #[uniffi::constructor]

472 + pub fn new(

473 + executor: Box<dyn FfiToolExecutor>,

474 + approvals: Box<dyn FfiApprovalHandler>,

475 + events: Box<dyn FfiEventListener>,

476 + config: Option<FfiAgentConfig>,

477 + ) -> Arc<Self> {

478 + Self::build(executor, approvals, events, config, None)

479 + }

480 +

481 + /// Like `new`, with an Onde app id so the engine reports usage under it.

482 + #[uniffi::constructor]

483 + pub fn new_with_app_id(

484 + executor: Box<dyn FfiToolExecutor>,

485 + approvals: Box<dyn FfiApprovalHandler>,

486 + events: Box<dyn FfiEventListener>,

487 + config: Option<FfiAgentConfig>,

488 + app_id: Option<String>,

489 + ) -> Arc<Self> {

490 + Self::build(executor, approvals, events, config, app_id)

491 + }

492

493 pub async fn register_tool(&self, tool: FfiToolDefinition) -> Result<(), FfiEdError> {

494 self.inner

@@ -556,6 +610,85 @@ impl FfiEdAgent {

610 }

611 }

612

613 + /// Stream the reply to `message` into `listener`, for plain chat without tools.

614 + /// Returns when the reply is complete, on error, or after `cancel_stream`.

615 + /// The turn so far stays in history either way.

616 + pub async fn stream_message(

617 + &self,

618 + message: String,

619 + listener: Box<dyn FfiStreamListener>,

620 + ) -> Result<(), FfiEdError> {

621 + let mut stream = self

622 + .inner

623 + .stream(message)

624 + .await

625 + .map_err(|error| FfiEdError::Failure {

626 + reason: error.to_string(),

627 + })?;

628 + let (cancel, mut cancelled) = tokio::sync::oneshot::channel();

629 + *self.stream_cancel.lock().unwrap_or_else(|e| e.into_inner()) = Some(cancel);

630 + loop {

631 + tokio::select! {

632 + next = stream.next() => match next {

633 + Some(Ok(delta)) => listener.on_delta(delta),

634 + Some(Err(error)) => {

635 + return Err(FfiEdError::Failure { reason: error.to_string() });

636 + }

637 + None => return Ok(()),

638 + },

639 + _ = &mut cancelled => return Ok(()),

640 + }

641 + }

642 + }

643 +

644 + /// Stop the `stream_message` in progress. Does nothing when none is.

645 + pub fn cancel_stream(&self) {

646 + if let Some(cancel) = self

647 + .stream_cancel

648 + .lock()

649 + .unwrap_or_else(|e| e.into_inner())

650 + .take()

651 + {

652 + let _ = cancel.send(());

653 + }

654 + }

655 +

656 + /// One-shot generation over `messages`. Leaves the conversation untouched.

657 + pub async fn generate(

658 + &self,

659 + messages: Vec<FfiChatMessage>,

660 + sampling: Option<FfiSamplingConfig>,

661 + ) -> Result<String, FfiEdError> {

662 + self.inner

663 + .generate(

664 + messages.into_iter().map(Into::into).collect(),

665 + sampling.map(Into::into),

666 + )

667 + .await

668 + .map_err(|error| FfiEdError::Failure {

669 + reason: error.to_string(),

670 + })

671 + }

672 +

673 + /// Replace the conversation with `messages`. Needs a loaded model.

674 + pub async fn restore_history(&self, messages: Vec<FfiChatMessage>) {

675 + self.inner

676 + .restore_history(messages.into_iter().map(Into::into).collect())

677 + .await;

678 + }

679 +

680 + pub async fn set_system_prompt(&self, prompt: String) {

681 + self.inner.set_system_prompt(prompt).await;

682 + }

683 +

684 + pub async fn clear_system_prompt(&self) {

685 + self.inner.clear_system_prompt().await;

686 + }

687 +

688 + pub async fn set_sampling(&self, sampling: FfiSamplingConfig) {

689 + self.inner.set_sampling(sampling.into()).await;

690 + }

691 +

692 pub fn cancel(&self) {

693 self.inner.cancel();

694 }

crates/ed-agent/Cargo.toml

@@ -19,6 +19,7 @@ serde_json.workspace = true

19 log.workspace = true

20 thiserror.workspace = true

21 tokio.workspace = true

22 +futures-core = "0.3"

24 [dev-dependencies]

25 tokio = { version = "1", features = ["macros", "rt", "sync"] }

crates/ed-agent/src/ed.rs

+130 -15

@@ -4,7 +4,7 @@ use std::collections::{HashMap, HashSet};

4 use std::future::Future;

5 use std::sync::atomic::{AtomicBool, Ordering};

6 use std::sync::Arc;

7 -use std::time::Duration;

7 +use std::time::{Duration, Instant};

9 use log::{error, info};

10 use onde::inference::types::{format_duration, ChatMessage};

@@ -14,9 +14,10 @@ use onde::inference::{

14 };

16 use crate::{

17 - tool::truncate_output, AgentConfig, AgentError, AgentReply, AgentToolDefinition,

18 - ApprovalDecision, ApprovalHandler, ApprovalRequest, ChatReply, DenyApprovals, EventSink,

19 - NoopSink, RejectingExecutor, ToolCall, ToolExecutionResult, ToolExecutor, ToolRisk,

17 + stream::chunk_error, tool::truncate_output, AgentConfig, AgentError, AgentReply,

18 + AgentToolDefinition, ApprovalDecision, ApprovalHandler, ApprovalRequest, ChatReply,

19 + DenyApprovals, EventSink, NoopSink, RejectingExecutor, TextStream, ToolCall,

20 + ToolExecutionResult, ToolExecutor, ToolRisk,

21 };

23 /// An embeddable chat agent over a single [`ChatEngine`].

@@ -226,6 +227,60 @@ impl<S: EventSink> Ed<S> {

227 reply

228 }

229

230 + /// Stream the reply to one message, for plain chat without tools.

231 + ///

232 + /// Returns once generation has started; the stream yields text deltas as

233 + /// they are produced. Unlike [`run`](Ed::run) this takes no turn lock and

234 + /// reports nothing to the sink, so don't overlap it with another turn.

235 + /// Drop the stream to cancel; a [`cancel`](Ed::cancel) call is not needed.

236 + pub async fn stream(&self, message: impl Into<String>) -> Result<TextStream, InferenceError> {

237 + Ok(TextStream::new(self.engine.stream_message(message).await?))

238 + }

239 +

240 + /// One-shot generation over explicit messages.

241 + ///

242 + /// Leaves the conversation history untouched; the engine's system prompt

243 + /// is still applied. `sampling` overrides the engine's for this call only.

244 + pub async fn generate(

245 + &self,

246 + messages: Vec<ChatMessage>,

247 + sampling: Option<SamplingConfig>,

248 + ) -> Result<String, InferenceError> {

249 + Ok(self.engine.generate(messages, sampling).await?.text)

250 + }

251 +

252 + /// Set the system prompt used for every later turn.

253 + ///

254 + /// Like the other engine setters below, and `restore_history`, this does

255 + /// nothing until a model is loaded: `onde` keeps these on the loaded model.

256 + pub async fn set_system_prompt(&self, prompt: impl Into<String>) {

257 + self.engine.set_system_prompt(prompt).await;

258 + }

259 +

260 + pub async fn clear_system_prompt(&self) {

261 + self.engine.clear_system_prompt().await;

262 + }

263 +

264 + /// Change the default sampling parameters.

265 + pub async fn set_sampling(&self, sampling: SamplingConfig) {

266 + self.engine.set_sampling(sampling).await;

267 + }

268 +

269 + /// Replace the conversation with `messages`.

270 + ///

271 + /// For switching to a saved session or restoring one at launch. Clears

272 + /// the current history first, and with it every `AllowForSession` grant,

273 + /// exactly as [`clear_history`](Ed::clear_history) does: a restored

274 + /// conversation starts without consent from the previous one.

275 + pub async fn restore_history(&self, messages: Vec<ChatMessage>) {

276 + self.clear_history().await;

277 + let count = messages.len();

278 + for message in messages {

279 + self.engine.push_history(message).await;

280 + }

281 + info!("ed: restored {count} messages");

282 + }

283 +

284 /// Register a host capability that the model may request.

285 pub async fn register_tool(&self, tool: AgentToolDefinition) -> Result<(), AgentError> {

286 tool.validate()?;

@@ -345,20 +400,39 @@ impl<S: EventSink> Ed<S> {

400 // `ChatReply` to the sink, and a host listening for both events

401 // would render this turn's answer twice. An agent turn reports

402 // `agent_replied` and nothing else.

348 - let result =

349 - self.engine

350 - .send_message(message)

351 - .await

352 - .map_err(|error| AgentError::Inference {

353 - reason: error.to_string(),

354 - })?;

355 - if result.text.trim().is_empty() {

403 + let started = Instant::now();

404 + let mut rx = self.engine.stream_message(message).await.map_err(|error| {

405 + AgentError::Inference {

406 + reason: error.to_string(),

407 + }

408 + })?;

409 + let mut text = String::new();

410 + while let Some(chunk) = rx.recv().await {

411 + if self.cancelled.load(Ordering::Acquire) {

412 + // Dropping the receiver stops generation.

413 + return Err(AgentError::Cancelled);

414 + }

415 + if chunk.done {

416 + if let Some(error) = chunk_error(&chunk) {

417 + return Err(AgentError::Inference {

418 + reason: error.to_string(),

419 + });

420 + }

421 + break;

422 + }

423 + if !chunk.delta.is_empty() {

424 + self.sink.text_delta(&chunk.delta);

425 + text.push_str(&chunk.delta);

426 + }

427 + }

428 + if text.trim().is_empty() {

429 return Err(AgentError::EmptyReply);

430 }

431 + let elapsed = started.elapsed();

432 let agent_reply = AgentReply {

359 - text: result.text,

360 - duration_seconds: result.duration_secs,

361 - duration: result.duration_display,

433 + text,

434 + duration_seconds: elapsed.as_secs_f64(),

435 + duration: format_duration(elapsed),

436 tool_rounds: 0,

437 };

438 self.sink.agent_replied(&agent_reply);

@@ -482,6 +556,10 @@ impl<S: EventSink> Ed<S> {

556 AgentError::EmptyReply

557 });

558 }

559 + // Tool rounds are not streamed: `onde` only parses tool calls out of a

560 + // complete response, so the answer arrives whole. Report it as one

561 + // delta so a host rendering deltas sees every turn the same way.

562 + self.sink.text_delta(&result.text);

563 let reply = AgentReply {

564 text: result.text,

565 duration_seconds,

@@ -833,4 +911,41 @@ mod tests {

911 assert!(result.content.contains("JSON object"));

912 assert_eq!(executor.0.load(Ordering::Relaxed), 0);

913 }

914 +

915 + #[tokio::test]

916 + async fn restoring_history_revokes_session_grants() {

917 + let ed = Ed::with_agent(

918 + NoopSink,

919 + Arc::new(CountingExecutor(AtomicUsize::new(0))),

920 + Arc::new(FixedApproval(ApprovalDecision::AllowForSession)),

921 + AgentConfig::default(),

922 + );

923 + ed.register_tool(AgentToolDefinition::mutating(

924 + "save_note",

925 + "Save a note",

926 + r#"{"type":"object"}"#,

927 + ))

928 + .await

929 + .unwrap();

930 + let tools = ed.tools.read().await.clone();

931 + ed.execute_tool_call(&call("save_note"), &tools).await;

932 + assert!(ed.session_grants.lock().await.contains("save_note"));

933 +

934 + ed.restore_history(Vec::new()).await;

935 +

936 + assert!(ed.session_grants.lock().await.is_empty());

937 + }

938 +

939 + #[tokio::test]

940 + async fn generation_and_streaming_need_a_loaded_model() {

941 + let ed = Ed::new();

942 + assert!(matches!(

943 + ed.generate(vec![ChatMessage::user("hi")], None).await,

944 + Err(InferenceError::NoModelLoaded)

945 + ));

946 + assert!(matches!(

947 + ed.stream("hi").await,

948 + Err(InferenceError::NoModelLoaded)

949 + ));

950 + }

951 }

crates/ed-agent/src/lib.rs

+6 -1

@@ -15,7 +15,8 @@

15 //! a frontend wants one shape covering both outcomes. That's [`ChatReply`].

16 //!

17 //! Ed deliberately does *not* own sessions, persistence, model catalogues, or

18 -//! downloads. Those differ per application and belong to the host.

18 +//! downloads. Those differ per application and belong to the host. For

19 +//! persistence it defines the [`SessionStore`] seam the host implements.

20 //!

21 //! # Platforms

22 //!

@@ -49,12 +50,16 @@

51 mod ed;

52 mod reply;

53 +mod session;

54 mod sink;

55 +mod stream;

56 mod tool;

58 pub use ed::Ed;

59 pub use reply::ChatReply;

60 +pub use session::{SessionError, SessionMeta, SessionStore};

61 pub use sink::{EventSink, NoopSink, StatusSink};

62 +pub use stream::TextStream;

63 pub use tool::{

64 AgentConfig, AgentError, AgentReply, AgentToolDefinition, ApprovalDecision, ApprovalHandler,

65 ApprovalRequest, DenyApprovals, RejectingExecutor, ToolCall, ToolExecutionResult, ToolExecutor,

crates/ed-agent/src/session.rs new

+178

@@ -0,0 +1,178 @@

1 +//! The seam between Ed and a host's chat storage.

2 +//!

3 +//! Ed owns no storage. A host that keeps several conversations implements

4 +//! [`SessionStore`] over whatever it already has (files, SQLite, a keychain)

5 +//! and hands the messages to [`Ed::restore_history`](crate::Ed::restore_history)

6 +//! when the user switches conversation.

7 +

8 +use async_trait::async_trait;

9 +use onde::inference::types::ChatMessage;

10 +use serde::{Deserialize, Serialize};

11 +

12 +/// What a session list needs to render a row, without loading the messages.

13 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]

14 +pub struct SessionMeta {

15 + pub id: String,

16 + pub title: String,

17 + /// Unix seconds.

18 + pub created_at: i64,

19 + /// Unix seconds.

20 + pub updated_at: i64,

21 + pub message_count: usize,

22 +}

23 +

24 +/// A storage failure. The store decides the wording; Ed only carries it.

25 +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]

26 +#[error("{0}")]

27 +pub struct SessionError(pub String);

28 +

29 +/// Persistent conversations, keyed by an id the host chooses.

30 +///

31 +/// Contract:

32 +/// - `load` of an unknown id is `Ok(None)`, not an error.

33 +/// - `save` creates the session if it doesn't exist and replaces its

34 +/// messages if it does.

35 +/// - `delete` of an unknown id succeeds.

36 +/// - `list` is ordered most recently updated first.

37 +#[async_trait]

38 +pub trait SessionStore: Send + Sync + 'static {

39 + async fn list(&self) -> Result<Vec<SessionMeta>, SessionError>;

40 + async fn load(&self, id: &str) -> Result<Option<Vec<ChatMessage>>, SessionError>;

41 + async fn save(&self, id: &str, messages: &[ChatMessage]) -> Result<(), SessionError>;

42 + async fn delete(&self, id: &str) -> Result<(), SessionError>;

43 +}

44 +

45 +#[cfg(test)]

46 +mod tests {

47 + use std::collections::HashMap;

48 + use std::sync::Mutex;

49 +

50 + use super::*;

51 +

52 + struct Stored {

53 + messages: Vec<ChatMessage>,

54 + created: i64,

55 + updated: i64,

56 + }

57 +

58 + /// The reference implementation of the contract above.

59 + #[derive(Default)]

60 + struct MemoryStore {

61 + sessions: Mutex<HashMap<String, Stored>>,

62 + clock: Mutex<i64>,

63 + }

64 +

65 + impl MemoryStore {

66 + fn tick(&self) -> i64 {

67 + let mut clock = self.clock.lock().unwrap();

68 + *clock += 1;

69 + *clock

70 + }

71 + }

72 +

73 + #[async_trait]

74 + impl SessionStore for MemoryStore {

75 + async fn list(&self) -> Result<Vec<SessionMeta>, SessionError> {

76 + let mut metas: Vec<_> = self

77 + .sessions

78 + .lock()

79 + .unwrap()

80 + .iter()

81 + .map(|(id, stored)| SessionMeta {

82 + id: id.clone(),

83 + title: stored

84 + .messages

85 + .first()

86 + .map(|m| m.content.clone())

87 + .unwrap_or_default(),

88 + created_at: stored.created,

89 + updated_at: stored.updated,

90 + message_count: stored.messages.len(),

91 + })

92 + .collect();

93 + metas.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));

94 + Ok(metas)

95 + }

96 +

97 + async fn load(&self, id: &str) -> Result<Option<Vec<ChatMessage>>, SessionError> {

98 + Ok(self

99 + .sessions

100 + .lock()

101 + .unwrap()

102 + .get(id)

103 + .map(|s| s.messages.clone()))

104 + }

105 +

106 + async fn save(&self, id: &str, messages: &[ChatMessage]) -> Result<(), SessionError> {

107 + let now = self.tick();

108 + let mut sessions = self.sessions.lock().unwrap();

109 + let created = sessions.get(id).map_or(now, |s| s.created);

110 + sessions.insert(

111 + id.to_owned(),

112 + Stored {

113 + messages: messages.to_vec(),

114 + created,

115 + updated: now,

116 + },

117 + );

118 + Ok(())

119 + }

120 +

121 + async fn delete(&self, id: &str) -> Result<(), SessionError> {

122 + self.sessions.lock().unwrap().remove(id);

123 + Ok(())

124 + }

125 + }

126 +

127 + #[tokio::test]

128 + async fn loading_an_unknown_session_is_none() {

129 + let store = MemoryStore::default();

130 + assert!(store.load("nope").await.unwrap().is_none());

131 + }

132 +

133 + #[tokio::test]

134 + async fn save_creates_then_replaces() {

135 + let store = MemoryStore::default();

136 + store.save("a", &[ChatMessage::user("hi")]).await.unwrap();

137 + store

138 + .save(

139 + "a",

140 + &[ChatMessage::user("hi"), ChatMessage::assistant("yo")],

141 + )

142 + .await

143 + .unwrap();

144 +

145 + let loaded = store.load("a").await.unwrap().unwrap();

146 + assert_eq!(loaded.len(), 2);

147 + let listed = store.list().await.unwrap();

148 + assert_eq!(listed.len(), 1);

149 + assert_eq!(listed[0].message_count, 2);

150 + assert!(listed[0].updated_at > listed[0].created_at);

151 + }

152 +

153 + #[tokio::test]

154 + async fn list_is_most_recently_updated_first() {

155 + let store = MemoryStore::default();

156 + store.save("old", &[ChatMessage::user("1")]).await.unwrap();

157 + store.save("new", &[ChatMessage::user("2")]).await.unwrap();

158 + store.save("old", &[ChatMessage::user("3")]).await.unwrap();

159 +

160 + let ids: Vec<_> = store

161 + .list()

162 + .await

163 + .unwrap()

164 + .into_iter()

165 + .map(|m| m.id)

166 + .collect();

167 + assert_eq!(ids, ["old", "new"]);

168 + }

169 +

170 + #[tokio::test]

171 + async fn delete_is_idempotent() {

172 + let store = MemoryStore::default();

173 + store.save("a", &[]).await.unwrap();

174 + store.delete("a").await.unwrap();

175 + store.delete("a").await.unwrap();

176 + assert!(store.list().await.unwrap().is_empty());

177 + }

178 +}

crates/ed-agent/src/sink.rs

+23

@@ -38,6 +38,14 @@ pub trait EventSink: Send + Sync + 'static {

38 /// pick the result up independently.

39 fn replied(&self, _reply: &ChatReply) {}

41 + /// A piece of the assistant's reply, as it is generated.

42 + ///

43 + /// Called zero or more times during [`Ed::run`](crate::Ed::run) before

44 + /// the final reply is reported. Concatenating the deltas gives the reply

45 + /// text. The default does nothing, so sinks written before streaming

46 + /// existed keep compiling and keep seeing only the finished reply.

47 + fn text_delta(&self, _delta: &str) {}

48 +

49 fn tool_requested(&self, _call: &ToolCall) {}

50 fn approval_requested(&self, _request: &ApprovalRequest) {}

51 fn tool_started(&self, _call: &ToolCall) {}

@@ -57,3 +65,18 @@ pub use EventSink as StatusSink;

65 pub struct NoopSink;

67 impl EventSink for NoopSink {}

68 +

69 +#[cfg(test)]

70 +mod tests {

71 + use super::*;

72 +

73 + struct Bare;

74 + impl EventSink for Bare {}

75 +

76 + #[test]

77 + fn text_delta_defaults_to_a_no_op() {

78 + // A sink that predates streaming must still compile and not panic.

79 + Bare.text_delta("hello");

80 + NoopSink.text_delta("hello");

81 + }

82 +}

crates/ed-agent/src/stream.rs new

+117

@@ -0,0 +1,117 @@

1 +//! Token streaming for plain chat.

2 +

3 +use std::pin::Pin;

4 +use std::task::{Context, Poll};

5 +

6 +use futures_core::Stream;

7 +use onde::inference::types::StreamChunk;

8 +use onde::inference::InferenceError;

9 +use tokio::sync::mpsc::Receiver;

10 +

11 +/// The reply to one message, as a stream of text deltas.

12 +///

13 +/// Yields `Ok(delta)` for each piece and ends after the last one. A failure

14 +/// mid-stream is a final `Err`. Dropping the stream stops generation: `onde`

15 +/// notices the closed channel and keeps only what was delivered.

16 +///

17 +/// The finished turn is added to the engine's history when the stream ends,

18 +/// as with [`Ed::send`](crate::Ed::send).

19 +pub struct TextStream {

20 + rx: Receiver<StreamChunk>,

21 + finished: bool,

22 +}

23 +

24 +impl TextStream {

25 + pub(crate) fn new(rx: Receiver<StreamChunk>) -> Self {

26 + Self {

27 + rx,

28 + finished: false,

29 + }

30 + }

31 +}

32 +

33 +impl TextStream {

34 + /// The next delta, `None` when the reply is complete. For callers without a stream

35 + /// combinator library.

36 + pub async fn next(&mut self) -> Option<Result<String, InferenceError>> {

37 + std::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await

38 + }

39 +}

40 +

41 +/// `onde` reports a failed stream as a final chunk whose `finish_reason`

42 +/// starts with `error`.

43 +pub(crate) fn chunk_error(chunk: &StreamChunk) -> Option<InferenceError> {

44 + chunk

45 + .finish_reason

46 + .as_deref()

47 + .filter(|reason| reason.starts_with("error"))

48 + .map(|reason| InferenceError::Inference {

49 + reason: reason.trim_start_matches("error:").trim().to_owned(),

50 + })

51 +}

52 +

53 +impl Stream for TextStream {

54 + type Item = Result<String, InferenceError>;

55 +

56 + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {

57 + loop {

58 + if self.finished {

59 + return Poll::Ready(None);

60 + }

61 + match self.rx.poll_recv(cx) {

62 + Poll::Pending => return Poll::Pending,

63 + Poll::Ready(None) => {

64 + self.finished = true;

65 + return Poll::Ready(None);

66 + }

67 + Poll::Ready(Some(chunk)) => {

68 + if chunk.done {

69 + self.finished = true;

70 + return Poll::Ready(chunk_error(&chunk).map(Err));

71 + }

72 + if chunk.delta.is_empty() {

73 + continue;

74 + }

75 + return Poll::Ready(Some(Ok(chunk.delta)));

76 + }

77 + }

78 + }

79 + }

80 +}

81 +

82 +#[cfg(test)]

83 +mod tests {

84 + use super::*;

85 +

86 + fn chunk(delta: &str, done: bool, reason: Option<&str>) -> StreamChunk {

87 + StreamChunk {

88 + delta: delta.to_owned(),

89 + done,

90 + finish_reason: reason.map(str::to_owned),

91 + }

92 + }

93 +

94 + #[tokio::test]

95 + async fn yields_deltas_then_ends() {

96 + let (tx, rx) = tokio::sync::mpsc::channel(8);

97 + let mut stream = TextStream::new(rx);

98 + tx.send(chunk("he", false, None)).await.unwrap();

99 + tx.send(chunk("llo", false, None)).await.unwrap();

100 + tx.send(chunk("", true, Some("stop"))).await.unwrap();

101 +

102 + assert_eq!(stream.next().await.unwrap().unwrap(), "he");

103 + assert_eq!(stream.next().await.unwrap().unwrap(), "llo");

104 + assert!(stream.next().await.is_none());

105 + }

106 +

107 + #[tokio::test]

108 + async fn a_failed_stream_ends_with_an_error() {

109 + let (tx, rx) = tokio::sync::mpsc::channel(8);

110 + let mut stream = TextStream::new(rx);

111 + tx.send(chunk("", true, Some("error: boom"))).await.unwrap();

112 +

113 + let err = stream.next().await.unwrap().unwrap_err();

114 + assert!(err.to_string().contains("boom"));

115 + assert!(stream.next().await.is_none());

116 + }

117 +}