Skip to main content

forge_engine/
tool_receipts.rs

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}