1use async_trait::async_trait;
2use llm_tool_runtime::{ToolError, ToolErrorClass, ToolReceipt, ToolReceiptSink};
3use semantic_memory_forge::{
4 ForgeToolBudgetContext, ForgeToolReceiptV2, FORGE_TOOL_RECEIPT_V2_SCHEMA,
5};
6use std::sync::Arc;
7
8use crate::{store::db::ToolReceiptRow, ForgeStore};
9
10#[derive(Debug, Clone)]
11pub struct ForgeToolReceiptSink {
12 store: Arc<ForgeStore>,
13}
14
15impl ForgeToolReceiptSink {
16 pub fn new(store: Arc<ForgeStore>) -> Self {
17 Self { store }
18 }
19}
20
21#[async_trait]
22impl ToolReceiptSink for ForgeToolReceiptSink {
23 async fn persist(&self, receipt: &ToolReceipt) -> Result<(), ToolError> {
24 let row = ToolReceiptRow {
25 receipt_id: receipt.receipt_id.clone(),
26 tool_run_id: receipt.tool_run_id.clone(),
27 tool_name: receipt.tool_name.clone(),
28 tool_version: receipt.tool_version.clone(),
29 backend_kind: format!("{:?}", receipt.backend_kind),
30 input_digest: receipt.input_digest.hex().to_string(),
31 output_digest_or_refs_json: serde_json::to_string(&receipt.output_digest_or_refs)
32 .map_err(|err| {
33 ToolError::new(
34 ToolErrorClass::ReceiptPersistence,
35 format!("failed to serialize tool receipt outputs: {}", err),
36 )
37 })?,
38 policy_hash: receipt.policy_hash.hex().to_string(),
39 approval_state: format!("{:?}", receipt.approval_state),
40 host_identity: receipt.host_identity.clone(),
41 started_at: receipt.started_at.clone(),
42 finished_at: receipt.finished_at.clone(),
43 trace_id: receipt.trace_ctx.trace_id.clone(),
44 trace_ctx_json: serde_json::to_string(&receipt.trace_ctx).map_err(|err| {
45 ToolError::new(
46 ToolErrorClass::ReceiptPersistence,
47 format!("failed to serialize trace context: {}", err),
48 )
49 })?,
50 attempt_id: receipt.attempt_id.as_str().to_string(),
51 trial_id: receipt.trial_id.as_str().to_string(),
52 error_class: receipt
53 .error_class
54 .as_ref()
55 .map(|class| format!("{:?}", class)),
56 retry_owner: format!("{:?}", receipt.retry_owner),
57 replay_link: receipt.replay_link.clone(),
58 provider_call_id: receipt.provider_call_id.clone(),
59 raw_payload_json: serde_json::to_string(&forge_tool_receipt(receipt)).map_err(
60 |err| {
61 ToolError::new(
62 ToolErrorClass::ReceiptPersistence,
63 format!("failed to serialize Forge tool receipt payload: {}", err),
64 )
65 },
66 )?,
67 recorded_at: chrono::Utc::now().to_rfc3339(),
68 };
69 self.store.insert_tool_receipt(&row).map_err(|err| {
70 ToolError::new(
71 ToolErrorClass::ReceiptPersistence,
72 format!("failed to persist Forge tool receipt: {}", err),
73 )
74 })?;
75 Ok(())
76 }
77}
78
79fn forge_tool_receipt(receipt: &ToolReceipt) -> ForgeToolReceiptV2 {
80 ForgeToolReceiptV2 {
81 schema_version: FORGE_TOOL_RECEIPT_V2_SCHEMA.into(),
82 receipt_id: receipt.receipt_id.clone(),
83 tool_run_id: receipt.tool_run_id.clone(),
84 tool_name: receipt.tool_name.clone(),
85 tool_version: receipt.tool_version.clone(),
86 backend_kind: format!("{:?}", receipt.backend_kind),
87 input_digest: receipt.input_digest.clone(),
88 output_digest_or_refs: receipt.output_digest_or_refs.clone(),
89 policy_hash: receipt.policy_hash.clone(),
90 approval_state: format!("{:?}", receipt.approval_state),
91 host_identity: receipt.host_identity.clone(),
92 started_at: receipt.started_at.clone(),
93 finished_at: receipt.finished_at.clone(),
94 trace_ctx: receipt.trace_ctx.clone(),
95 attempt_id: receipt.attempt_id.clone(),
96 trial_id: receipt.trial_id.clone(),
97 planner_stage: format!("{:?}", receipt.planner_stage),
98 deadline: receipt.deadline.clone(),
99 workload_class: receipt.workload_class.clone(),
100 budget_context: receipt
101 .budget_context
102 .as_ref()
103 .map(|budget| ForgeToolBudgetContext {
104 budget_kind: budget.budget_kind.clone(),
105 max_steps: budget.max_steps,
106 time_budget_ms: budget.time_budget_ms,
107 cost_budget_units: budget.cost_budget_units,
108 }),
109 parent_receipt_id: receipt.parent_receipt_id.clone(),
110 family_receipt_id: receipt.family_receipt_id.clone(),
111 replay_parent_receipt_id: receipt.replay_parent_receipt_id.clone(),
112 error_class: receipt
113 .error_class
114 .as_ref()
115 .map(|class| format!("{:?}", class)),
116 retry_owner: format!("{:?}", receipt.retry_owner),
117 replay_link: receipt.replay_link.clone(),
118 provider_call_id: receipt.provider_call_id.clone(),
119 raw_payload: serde_json::to_value(receipt).unwrap_or_else(|_| serde_json::json!({})),
120 }
121}