1use crate::{
2 ActorRef, ArtifactRef, Decision, DecisionStatus, EvalResult, Event, EventId, EventStore,
3 ForkId, ForkSnapshot, Frame, Goal, GoalId, GoalStatus, Graph, GraphDiff, GraphSnapshot,
4 Hypothesis, Lineage, ModelCall, NodeId, Observation, PatchId, PatchStatus, ProjectSnapshot,
5 RunId, StateError, StateOp, StatePatch, Task, TaskId, TaskStatus, ToolCall, diff_graphs,
6 fork_events_at, project_event, replay,
7};
8use serde_json::json;
9use std::sync::Arc;
10use tokio::sync::RwLock;
11
12#[derive(Debug)]
13pub struct YoAgentState<S: EventStore> {
14 store: Arc<S>,
15 graph: Arc<RwLock<Graph>>,
16}
17
18impl<S: EventStore> Clone for YoAgentState<S> {
19 fn clone(&self) -> Self {
20 Self {
21 store: self.store.clone(),
22 graph: self.graph.clone(),
23 }
24 }
25}
26
27impl<S: EventStore> YoAgentState<S> {
28 pub async fn load(store: S) -> Result<Self, StateError> {
29 let events = store.scan().await?;
30 let graph = replay(&events)?;
31 Ok(Self {
32 store: Arc::new(store),
33 graph: Arc::new(RwLock::new(graph)),
34 })
35 }
36
37 pub fn store(&self) -> Arc<S> {
38 self.store.clone()
39 }
40
41 pub async fn record_event(&self, event: Event) -> Result<EventId, StateError> {
42 let ids = self.store.append(vec![event.clone()]).await?;
43 {
44 let mut graph = self.graph.write().await;
45 project_event(&mut graph, &event)?;
46 }
47 Ok(ids[0].clone())
48 }
49
50 pub async fn apply_ops(
51 &self,
52 actor: ActorRef,
53 ops: Vec<StateOp>,
54 ) -> Result<EventId, StateError> {
55 let event = Event::new(actor, crate::STATE_OPS_APPLIED, serde_json::to_value(ops)?);
56 self.record_event(event).await
57 }
58
59 pub async fn propose_patch(&self, patch: StatePatch) -> Result<PatchId, StateError> {
60 let patch_id = patch.id.clone();
61 let actor = patch.created_by.clone();
62 self.record_event(Event::new(
63 actor.clone(),
64 "patch.proposed",
65 serde_json::to_value(&patch)?,
66 ))
67 .await?;
68
69 let patch_node_id = NodeId::new(patch.id.0.clone());
70 let mut ops = vec![StateOp::CreateNode {
71 id: patch_node_id.clone(),
72 kind: "patch".to_string(),
73 props: json!({
74 "title": patch.title,
75 "summary": patch.summary,
76 "status": patch.status,
77 "base_state_version": patch.base_state_version,
78 "preconditions": patch.preconditions,
79 "expected_effects": patch.expected_effects,
80 "base_project_ref": patch.base_project_ref,
81 }),
82 }];
83
84 for evidence in patch.evidence {
85 ops.push(StateOp::CreateRelation {
86 from: patch_node_id.clone(),
87 rel: "addresses".to_string(),
88 to: evidence,
89 props: json!({}),
90 });
91 }
92
93 for artifact in patch.artifacts {
94 ops.push(StateOp::AttachArtifact {
95 id: patch_node_id.clone(),
96 artifact,
97 });
98 }
99
100 ops.extend(patch.ops);
101 self.apply_ops(actor, ops).await?;
102 Ok(patch_id)
103 }
104
105 pub async fn update_patch_status(
106 &self,
107 patch_id: PatchId,
108 status: PatchStatus,
109 reason: Option<String>,
110 ) -> Result<EventId, StateError> {
111 let actor = ActorRef::system("yoagent-state");
112 let event = Event::new(
113 actor.clone(),
114 "patch.status_changed",
115 json!({
116 "patch_id": patch_id,
117 "status": status,
118 "reason": reason,
119 }),
120 );
121 self.record_event(event).await?;
122
123 self.apply_ops(
124 actor,
125 vec![StateOp::UpdateNode {
126 id: NodeId::new(patch_id.0),
127 props: json!({
128 "status": status,
129 "status_reason": reason,
130 }),
131 }],
132 )
133 .await
134 }
135
136 pub async fn attach_artifact(
137 &self,
138 node_id: NodeId,
139 artifact: ArtifactRef,
140 ) -> Result<EventId, StateError> {
141 let actor = ActorRef::system("yoagent-state");
142 self.record_event(Event::new(
143 actor.clone(),
144 "artifact.attached",
145 json!({
146 "node_id": node_id,
147 "artifact": artifact,
148 }),
149 ))
150 .await?;
151
152 self.apply_ops(
153 actor,
154 vec![StateOp::AttachArtifact {
155 id: node_id,
156 artifact,
157 }],
158 )
159 .await
160 }
161
162 pub async fn graph(&self) -> GraphSnapshot {
163 self.graph.read().await.clone()
164 }
165
166 pub async fn get_node(&self, node_id: NodeId) -> Option<crate::Node> {
167 self.graph.read().await.get_node(&node_id).cloned()
168 }
169
170 pub async fn outgoing(&self, node_id: NodeId, rel: Option<&str>) -> Vec<crate::Relation> {
171 self.graph.read().await.outgoing(&node_id, rel)
172 }
173
174 pub async fn incoming(&self, node_id: NodeId, rel: Option<&str>) -> Vec<crate::Relation> {
175 self.graph.read().await.incoming(&node_id, rel)
176 }
177
178 pub async fn related(&self, node_id: NodeId) -> Vec<crate::Relation> {
179 self.graph.read().await.related(&node_id)
180 }
181
182 pub async fn patches_for_failure(&self, failure_id: NodeId) -> Vec<crate::Node> {
183 let graph = self.graph.read().await;
184 graph
185 .incoming(&failure_id, Some("addresses"))
186 .into_iter()
187 .filter_map(|rel| graph.get_node(&rel.from).cloned())
188 .filter(|node| node.kind == "patch")
189 .collect()
190 }
191
192 pub async fn evals_for_patch(&self, patch_id: PatchId) -> Vec<crate::Node> {
193 let patch_node_id = NodeId::new(patch_id.0);
194 let graph = self.graph.read().await;
195 graph
196 .outgoing(&patch_node_id, Some("validated_by"))
197 .into_iter()
198 .filter_map(|rel| graph.get_node(&rel.to).cloned())
199 .collect()
200 }
201
202 pub async fn lineage(&self, node_id: NodeId) -> Lineage {
203 let graph = self.graph.read().await;
204 Lineage::from_graph(&graph, &node_id)
205 }
206
207 pub async fn record_run_started(
208 &self,
209 actor: ActorRef,
210 run_id: RunId,
211 task: impl Into<String>,
212 ) -> Result<EventId, StateError> {
213 let task = task.into();
214 self.record_event(Event::new(
215 actor.clone(),
216 "run.started",
217 json!({ "run_id": run_id, "task": task }),
218 ))
219 .await?;
220 self.apply_ops(
221 actor,
222 vec![StateOp::CreateNode {
223 id: NodeId::new(run_id.0),
224 kind: crate::KIND_RUN.to_string(),
225 props: json!({ "task": task, "status": "started" }),
226 }],
227 )
228 .await
229 }
230
231 pub async fn record_run_finished(
232 &self,
233 actor: ActorRef,
234 run_id: RunId,
235 outcome: impl Into<String>,
236 ) -> Result<EventId, StateError> {
237 self.record_event(Event::new(
238 actor,
239 "run.finished",
240 json!({ "run_id": run_id, "outcome": outcome.into() }),
241 ))
242 .await
243 }
244
245 pub async fn record_goal(&self, goal: Goal) -> Result<EventId, StateError> {
246 let actor = goal.owner.clone();
247 self.record_event(Event::new(
248 actor.clone(),
249 "goal.created",
250 serde_json::to_value(&goal)?,
251 ))
252 .await?;
253 self.apply_ops(
254 actor,
255 vec![StateOp::CreateNode {
256 id: NodeId::new(goal.id.0),
257 kind: crate::KIND_GOAL.to_string(),
258 props: json!({
259 "title": goal.title,
260 "summary": goal.summary,
261 "status": goal.status,
262 "owner": goal.owner,
263 "metadata": goal.metadata,
264 }),
265 }],
266 )
267 .await
268 }
269
270 pub async fn update_goal_status(
271 &self,
272 goal_id: GoalId,
273 status: GoalStatus,
274 reason: Option<String>,
275 ) -> Result<EventId, StateError> {
276 let actor = ActorRef::system("yoagent-state");
277 self.record_event(Event::new(
278 actor.clone(),
279 "goal.status_changed",
280 json!({ "goal_id": goal_id, "status": status, "reason": reason }),
281 ))
282 .await?;
283 self.apply_ops(
284 actor,
285 vec![StateOp::UpdateNode {
286 id: NodeId::new(goal_id.0),
287 props: json!({ "status": status, "status_reason": reason }),
288 }],
289 )
290 .await
291 }
292
293 pub async fn record_task(&self, task: Task) -> Result<EventId, StateError> {
294 let actor = task.created_by.clone();
295 self.record_event(Event::new(
296 actor.clone(),
297 "task.created",
298 serde_json::to_value(&task)?,
299 ))
300 .await?;
301 let task_id = NodeId::new(task.id.0);
302 let mut ops = vec![StateOp::CreateNode {
303 id: task_id.clone(),
304 kind: crate::KIND_TASK.to_string(),
305 props: json!({
306 "title": task.title,
307 "summary": task.summary,
308 "status": task.status,
309 "metadata": task.metadata,
310 }),
311 }];
312 if let Some(goal) = task.goal {
313 ops.push(StateOp::CreateRelation {
314 from: task_id,
315 rel: crate::REL_SERVES.to_string(),
316 to: NodeId::new(goal.0),
317 props: json!({}),
318 });
319 }
320 self.apply_ops(actor, ops).await
321 }
322
323 pub async fn update_task_status(
324 &self,
325 task_id: TaskId,
326 status: TaskStatus,
327 reason: Option<String>,
328 ) -> Result<EventId, StateError> {
329 let actor = ActorRef::system("yoagent-state");
330 self.record_event(Event::new(
331 actor.clone(),
332 "task.status_changed",
333 json!({ "task_id": task_id, "status": status, "reason": reason }),
334 ))
335 .await?;
336 self.apply_ops(
337 actor,
338 vec![StateOp::UpdateNode {
339 id: NodeId::new(task_id.0),
340 props: json!({ "status": status, "status_reason": reason }),
341 }],
342 )
343 .await
344 }
345
346 pub async fn record_observation(
347 &self,
348 actor: ActorRef,
349 observation: Observation,
350 ) -> Result<EventId, StateError> {
351 self.record_event(Event::new(
352 actor.clone(),
353 "observation.created",
354 serde_json::to_value(&observation)?,
355 ))
356 .await?;
357 let observation_id = NodeId::new(observation.id.0);
358 let mut ops = vec![StateOp::CreateNode {
359 id: observation_id.clone(),
360 kind: crate::KIND_OBSERVATION.to_string(),
361 props: json!({
362 "title": observation.title,
363 "summary": observation.summary,
364 "metadata": observation.metadata,
365 }),
366 }];
367 if let Some(run_id) = observation.observed_in {
368 ops.push(StateOp::CreateRelation {
369 from: observation_id,
370 rel: "observed_in".to_string(),
371 to: NodeId::new(run_id.0),
372 props: json!({}),
373 });
374 }
375 self.apply_ops(actor, ops).await
376 }
377
378 pub async fn record_hypothesis(
379 &self,
380 actor: ActorRef,
381 hypothesis: Hypothesis,
382 explains: Option<NodeId>,
383 ) -> Result<EventId, StateError> {
384 self.record_event(Event::new(
385 actor.clone(),
386 "hypothesis.created",
387 serde_json::to_value(&hypothesis)?,
388 ))
389 .await?;
390 let hypothesis_id = NodeId::new(hypothesis.id.0);
391 let mut ops = vec![StateOp::CreateNode {
392 id: hypothesis_id.clone(),
393 kind: crate::KIND_HYPOTHESIS.to_string(),
394 props: json!({
395 "title": hypothesis.title,
396 "summary": hypothesis.summary,
397 "confidence": hypothesis.confidence,
398 "metadata": hypothesis.metadata,
399 }),
400 }];
401 if let Some(target) = explains {
402 ops.push(StateOp::CreateRelation {
403 from: hypothesis_id,
404 rel: crate::REL_EXPLAINS.to_string(),
405 to: target,
406 props: json!({}),
407 });
408 }
409 self.apply_ops(actor, ops).await
410 }
411
412 pub async fn record_failure(
413 &self,
414 actor: ActorRef,
415 failure_id: NodeId,
416 title: impl Into<String>,
417 summary: impl Into<String>,
418 ) -> Result<EventId, StateError> {
419 let title = title.into();
420 let summary = summary.into();
421 self.record_event(Event::new(
422 actor.clone(),
423 "failure.observed",
424 json!({
425 "failure_id": failure_id,
426 "title": title,
427 "summary": summary,
428 }),
429 ))
430 .await?;
431 self.apply_ops(
432 actor,
433 vec![StateOp::CreateNode {
434 id: failure_id,
435 kind: "failure".to_string(),
436 props: json!({ "title": title, "summary": summary }),
437 }],
438 )
439 .await
440 }
441
442 pub async fn record_eval_result(
443 &self,
444 actor: ActorRef,
445 eval_id: NodeId,
446 patch_id: PatchId,
447 command: impl Into<String>,
448 passed: bool,
449 ) -> Result<EventId, StateError> {
450 let patch_node_id = NodeId::new(patch_id.0);
451 self.apply_ops(
452 actor,
453 vec![
454 StateOp::CreateNode {
455 id: eval_id.clone(),
456 kind: "eval".to_string(),
457 props: json!({ "command": command.into(), "passed": passed }),
458 },
459 StateOp::CreateRelation {
460 from: patch_node_id,
461 rel: "validated_by".to_string(),
462 to: eval_id,
463 props: json!({ "passed": passed }),
464 },
465 ],
466 )
467 .await
468 }
469
470 pub async fn record_eval(
471 &self,
472 actor: ActorRef,
473 eval: EvalResult,
474 patch_id: Option<PatchId>,
475 ) -> Result<EventId, StateError> {
476 self.record_event(Event::new(
477 actor.clone(),
478 "eval.finished",
479 serde_json::to_value(&eval)?,
480 ))
481 .await?;
482 let eval_node_id = NodeId::new(eval.id.0);
483 let mut ops = vec![StateOp::CreateNode {
484 id: eval_node_id.clone(),
485 kind: crate::KIND_EVAL.to_string(),
486 props: json!({
487 "command": eval.command,
488 "status": eval.status,
489 "score": eval.score,
490 "metadata": eval.metadata,
491 }),
492 }];
493 if let Some(patch_id) = patch_id {
494 ops.push(StateOp::CreateRelation {
495 from: NodeId::new(patch_id.0),
496 rel: crate::REL_VALIDATED_BY.to_string(),
497 to: eval_node_id,
498 props: json!({}),
499 });
500 }
501 self.apply_ops(actor, ops).await
502 }
503
504 pub async fn record_decision(
505 &self,
506 actor: ActorRef,
507 decision_id: NodeId,
508 patch_id: PatchId,
509 approved: bool,
510 reason: impl Into<String>,
511 ) -> Result<EventId, StateError> {
512 let patch_node_id = NodeId::new(patch_id.0);
513 let rel = if approved {
514 "approved_by"
515 } else {
516 "rejected_by"
517 };
518
519 self.apply_ops(
520 actor,
521 vec![
522 StateOp::CreateNode {
523 id: decision_id.clone(),
524 kind: "decision".to_string(),
525 props: json!({ "approved": approved, "reason": reason.into() }),
526 },
527 StateOp::CreateRelation {
528 from: patch_node_id,
529 rel: rel.to_string(),
530 to: decision_id,
531 props: json!({}),
532 },
533 ],
534 )
535 .await
536 }
537
538 pub async fn record_decision_node(
539 &self,
540 actor: ActorRef,
541 decision: Decision,
542 target: Option<NodeId>,
543 ) -> Result<EventId, StateError> {
544 self.record_event(Event::new(
545 actor.clone(),
546 "decision.created",
547 serde_json::to_value(&decision)?,
548 ))
549 .await?;
550 let decision_node_id = NodeId::new(decision.id.0);
551 let mut ops = vec![StateOp::CreateNode {
552 id: decision_node_id.clone(),
553 kind: crate::KIND_DECISION.to_string(),
554 props: json!({
555 "status": decision.status,
556 "reason": decision.reason,
557 "decided_by": decision.decided_by,
558 "metadata": decision.metadata,
559 }),
560 }];
561 if let Some(target) = target {
562 let rel = match decision.status {
563 DecisionStatus::Approved => crate::REL_APPROVED_BY,
564 DecisionStatus::Rejected => crate::REL_REJECTED_BY,
565 _ => crate::REL_REFERENCES,
566 };
567 ops.push(StateOp::CreateRelation {
568 from: target,
569 rel: rel.to_string(),
570 to: decision_node_id,
571 props: json!({}),
572 });
573 }
574 self.apply_ops(actor, ops).await
575 }
576
577 pub async fn record_project_snapshot(
578 &self,
579 actor: ActorRef,
580 snapshot: ProjectSnapshot,
581 ) -> Result<EventId, StateError> {
582 self.record_event(Event::new(
583 actor.clone(),
584 "project.snapshot_recorded",
585 serde_json::to_value(&snapshot)?,
586 ))
587 .await?;
588 let mut ops = vec![StateOp::CreateNode {
589 id: snapshot.id.clone(),
590 kind: crate::KIND_PROJECT_SNAPSHOT.to_string(),
591 props: json!({
592 "project": snapshot.project,
593 "metadata": snapshot.metadata,
594 }),
595 }];
596 for artifact in snapshot.artifacts {
597 ops.push(StateOp::AttachArtifact {
598 id: snapshot.id.clone(),
599 artifact,
600 });
601 }
602 self.apply_ops(actor, ops).await
603 }
604
605 pub async fn record_model_call(
606 &self,
607 actor: ActorRef,
608 call: ModelCall,
609 ) -> Result<EventId, StateError> {
610 self.record_event(Event::new(
611 actor.clone(),
612 "model.called",
613 serde_json::to_value(&call)?,
614 ))
615 .await?;
616 self.apply_ops(
617 actor,
618 vec![
619 StateOp::CreateNode {
620 id: call.id.clone(),
621 kind: crate::KIND_MODEL_CALL.to_string(),
622 props: json!({
623 "run_id": call.run_id,
624 "model": call.model,
625 "prompt_summary": call.prompt_summary,
626 "output_summary": call.output_summary,
627 "metadata": call.metadata,
628 }),
629 },
630 StateOp::CreateRelation {
631 from: call.id,
632 rel: crate::REL_PRODUCED_BY.to_string(),
633 to: NodeId::new(call.run_id.0),
634 props: json!({}),
635 },
636 ],
637 )
638 .await
639 }
640
641 pub async fn record_tool_call(
642 &self,
643 actor: ActorRef,
644 call: ToolCall,
645 ) -> Result<EventId, StateError> {
646 self.record_event(Event::new(
647 actor.clone(),
648 "tool.called",
649 serde_json::to_value(&call)?,
650 ))
651 .await?;
652 self.apply_ops(
653 actor,
654 vec![
655 StateOp::CreateNode {
656 id: call.id.clone(),
657 kind: crate::KIND_TOOL_CALL.to_string(),
658 props: json!({
659 "run_id": call.run_id,
660 "tool": call.tool,
661 "input_summary": call.input_summary,
662 "output_summary": call.output_summary,
663 "success": call.success,
664 "metadata": call.metadata,
665 }),
666 },
667 StateOp::CreateRelation {
668 from: call.id,
669 rel: crate::REL_PRODUCED_BY.to_string(),
670 to: NodeId::new(call.run_id.0),
671 props: json!({}),
672 },
673 ],
674 )
675 .await
676 }
677
678 pub async fn record_frame(&self, actor: ActorRef, frame: Frame) -> Result<EventId, StateError> {
679 self.record_event(Event::new(
680 actor.clone(),
681 "frame.created",
682 serde_json::to_value(&frame)?,
683 ))
684 .await?;
685 let frame_node_id = NodeId::new(frame.id.0);
686 let mut ops = vec![StateOp::CreateNode {
687 id: frame_node_id.clone(),
688 kind: crate::KIND_FRAME.to_string(),
689 props: json!({
690 "title": frame.title,
691 "summary": frame.summary,
692 "metadata": frame.metadata,
693 }),
694 }];
695 if let Some(parent) = frame.parent {
696 ops.push(StateOp::CreateRelation {
697 from: frame_node_id,
698 rel: crate::REL_CONTAINED_IN_FRAME.to_string(),
699 to: NodeId::new(parent.0),
700 props: json!({}),
701 });
702 }
703 self.apply_ops(actor, ops).await
704 }
705
706 pub async fn link(
707 &self,
708 actor: ActorRef,
709 from: NodeId,
710 rel: impl Into<String>,
711 to: NodeId,
712 ) -> Result<EventId, StateError> {
713 self.apply_ops(
714 actor,
715 vec![StateOp::CreateRelation {
716 from,
717 rel: rel.into(),
718 to,
719 props: json!({}),
720 }],
721 )
722 .await
723 }
724
725 pub async fn fork_at_event(
726 &self,
727 fork_id: ForkId,
728 parent_event: Option<EventId>,
729 ) -> Result<ForkSnapshot, StateError> {
730 let events = self.store.scan().await?;
731 fork_events_at(&events, fork_id, parent_event)
732 }
733
734 pub async fn diff_with(&self, other: &Graph) -> GraphDiff {
735 let current = self.graph().await;
736 diff_graphs(¤t, other)
737 }
738}