fix(sagent): keep the archive watermark across runs and rearm subscribers on token issue

Modified lib/sagent/src/services/service.rs +20 −2 View lib/sagent/src/services/service.rs

@@ -1,5 +1,5 @@

1use std::collections::HashMap;

2use std::sync::{Arc, OnceLock};

1use std::collections::{HashMap, HashSet};

2use std::sync::{Arc, Mutex, OnceLock};

33

44use async_trait::async_trait;

55use config::AppConfig;

@@ -22,6 +22,7 @@

2222pub struct SagentService {

2323 pub state: AppState,

2424 gateway: Arc<OnceLock<Result<AgentGateway, String>>>,

25 subscribers: Arc<Mutex<HashSet<Uuid>>>,

2526}

2627

2728impl SagentService {

@@ -29,9 +30,26 @@

2930 Self {

3031 state,

3132 gateway: Arc::new(OnceLock::new()),

33 subscribers: Arc::new(Mutex::new(HashSet::new())),

3234 }

3335 }

3436

37 pub(crate) fn claim_subscriber(&self, task_id: Uuid) -> bool {

38 let mut running = self

39 .subscribers

40 .lock()

41 .unwrap_or_else(|err| err.into_inner());

42 running.insert(task_id)

43 }

44

45 pub(crate) fn release_subscriber(&self, task_id: Uuid) {

46 let mut running = self

47 .subscribers

48 .lock()

49 .unwrap_or_else(|err| err.into_inner());

50 running.remove(&task_id);

51 }

52

3553 pub fn config(&self) -> &AppConfig {

3654 &self.state.config

3755 }

Modified lib/sagent/src/services/subscribe.rs +6 −1 View lib/sagent/src/services/subscribe.rs

@@ -54,9 +54,14 @@

5454 }

5555

5656 pub(crate) fn spawn_archive_subscriber(&self, task_id: Uuid) {

57 if !self.claim_subscriber(task_id) {

58 return;

59 }

5760 let service = self.clone();

5861 tokio::spawn(async move {

59 if let Err(error) = service.run_archive_subscriber(task_id).await {

62 let result = service.run_archive_subscriber(task_id).await;

63 service.release_subscriber(task_id);

64 if let Err(error) = result {

6065 tracing::error!(?error, task_id = %task_id, "sagent archive subscriber failed");

6166 }

6267 });

Modified lib/sagent/src/services/tasks.rs +0 −1 View lib/sagent/src/services/tasks.rs

@@ -589,7 +589,6 @@

589589 let mut update: agent::ActiveModel = task.clone().into();

590590 update.session_id = Set(Some(session_id));

591591 update.last_run_id = Set(Some(run_id));

592 update.last_event_seq = Set(0);

593592 update.state = Set(map_run_state(accepted.state));

594593 update.started_at = Set(Some(task.started_at.unwrap_or(now)));

595594 update.finished_at = Set(None);

Modified lib/sagent/src/services/tokens.rs +2 −0 View lib/sagent/src/services/tokens.rs

@@ -34,6 +34,8 @@

3434 .open_session(session_id, caller, &INVOKE_SCOPES)

3535 .await?;

3636

37 self.spawn_archive_subscriber(task.id);

38

3739 Ok(TaskTokenDto {

3840 session_id,

3941 token: handle.token().to_string(),