Skip to main content

sim_lib_bridge/
ask.rs

1use sim_codec::{Input, Output, decode_with_codec, encode_with_codec};
2use sim_codec_bridge::{
3    BridgeBook, BridgeCallArgument, BridgeCallPayload, BridgeHeader, BridgePacket, BridgePart,
4    BridgeProvenance, CallArgumentMedia, content_id_string, stamp_packet_cid,
5};
6use sim_kernel::{
7    Cx, Datum, EncodeOptions, EvalFabric, Expr, ReadPolicy, Result, Symbol, encode::EncodePosition,
8};
9use sim_lib_agent_runner_core::{
10    InjectionFence, ModelResponse, OutputContract, terminal_model_content,
11};
12use sim_shape::{check_value_report, shape_value};
13use sim_value::{access::field, build::entry};
14
15use crate::model::output_contract_for_packet;
16use crate::parent::parent_token;
17use crate::repair::{AskFailure, RepairPolicy};
18use crate::rx::{effective_caps, rx_check, shape_from_contract_expr};
19use crate::tx::{eval_request_for_checked_packet, prepare_packet};
20
21/// Default codec for packing ASK call arguments and answers.
22pub fn ask_default_codec() -> Symbol {
23    Symbol::qualified("codec", "json")
24}
25
26/// Builds an ASK request packet with model parameters omitted.
27pub fn ask_packet(
28    cx: &mut Cx,
29    name: &str,
30    params: Vec<(String, Expr)>,
31    return_shape: Expr,
32    to: &str,
33) -> Result<BridgePacket> {
34    ask_packet_with_model_params(cx, name, params, Vec::new(), return_shape, to)
35}
36
37/// Builds an ASK request packet.
38///
39/// Argument values are encoded through the default ASK codec at data position,
40/// wrapped in deterministic injection fences, and only then stored in the
41/// packet.
42pub fn ask_packet_with_model_params(
43    cx: &mut Cx,
44    name: &str,
45    params: Vec<(String, Expr)>,
46    model_params: Vec<(String, Expr)>,
47    return_shape: Expr,
48    to: &str,
49) -> Result<BridgePacket> {
50    let codec = ask_default_codec();
51    let mut call = BridgeCallPayload::new(symbol_from_name(name));
52    for (name, value) in params {
53        call = call.with_arg(pack_argument(cx, &name, &codec, &value)?);
54    }
55    for (name, value) in model_params {
56        call = call.with_model_param(symbol_from_name(&name), value);
57    }
58    Ok(BridgePacket {
59        header: BridgeHeader {
60            cid: None,
61            move_kind: Symbol::new("request"),
62            from: "sim".to_owned(),
63            to: vec![to.to_owned()],
64            role: Symbol::new("implementer"),
65            parents: Vec::new(),
66            task: Symbol::new("C1"),
67            output: Symbol::new("O1"),
68            ceiling: ask_capability_ceiling(),
69            context: Vec::new(),
70            provenance: BridgeProvenance::default(),
71        },
72        body: vec![
73            BridgePart {
74                id: Symbol::new("C1"),
75                kind: Symbol::qualified("bridge", "Call"),
76                payload: call.to_expr(),
77            },
78            BridgePart {
79                id: Symbol::new("O1"),
80                kind: Symbol::qualified("bridge", "Return"),
81                payload: Expr::Map(vec![
82                    entry("codec", Expr::Symbol(codec)),
83                    entry("shape", return_shape),
84                ]),
85            },
86        ],
87        warrant: None,
88    })
89}
90
91fn ask_capability_ceiling() -> Vec<Symbol> {
92    [
93        Symbol::qualified("ai", "run"),
94        Symbol::qualified("capability", "ai-runner"),
95        Symbol::qualified("capability", "ai-runner-local"),
96        Symbol::qualified("capability", "ai-runner-network"),
97        Symbol::qualified("capability", "ai-runner-secret"),
98        Symbol::qualified("capability", "exec"),
99        Symbol::qualified("capability", "host.process"),
100    ]
101    .into()
102}
103
104/// Runs an ASK packet with the default bounded repair policy.
105pub fn run_ask(cx: &mut Cx, target: &dyn EvalFabric, packet: BridgePacket) -> Result<BridgePacket> {
106    run_ask_with_policy(cx, target, packet, RepairPolicy::default())
107}
108
109/// Runs an ASK packet with an explicit bounded repair policy.
110pub fn run_ask_with_policy(
111    cx: &mut Cx,
112    target: &dyn EvalFabric,
113    mut packet: BridgePacket,
114    policy: RepairPolicy,
115) -> Result<BridgePacket> {
116    let book = BridgeBook::standard();
117    let max_retries = policy.retries();
118    for attempt in 0..=max_retries {
119        let checked = prepare_packet(cx, &book, &packet)?;
120        let request = eval_request_for_checked_packet(cx, &book, &checked)?;
121        let caps = effective_caps(cx, &checked)?;
122        let reply = cx.with_capabilities(caps, |cx| target.realize(cx, request))?;
123        let response = ModelResponse::try_from(reply.value.object().as_expr(cx)?)?;
124        match answer_packet(cx, &book, &checked, &response)? {
125            Ok(answer) => return Ok(answer),
126            Err(failure) if attempt < max_retries => {
127                packet = repair_packet_for_failure(cx, &checked, &failure, attempt + 1)?;
128            }
129            Err(failure) => {
130                return Err(sim_kernel::Error::Eval(format!(
131                    "bridge ask failed after {} attempt(s): {}",
132                    attempt + 1,
133                    failure.message()
134                )));
135            }
136        }
137    }
138    unreachable!("bounded ASK loop always returns inside the retry range")
139}
140
141pub(crate) fn pack_argument(
142    cx: &mut Cx,
143    name: &str,
144    codec: &Symbol,
145    value: &Expr,
146) -> Result<BridgeCallArgument> {
147    let output = encode_with_codec(
148        cx,
149        codec,
150        value,
151        EncodeOptions {
152            position: EncodePosition::Data,
153            ..EncodeOptions::default()
154        },
155    )?;
156    let (media, datum, body) = match output {
157        Output::Text(text) => (CallArgumentMedia::Text, Datum::String(text.clone()), text),
158        Output::Bytes(bytes) => (
159            CallArgumentMedia::Bytes,
160            Datum::Bytes(bytes.clone()),
161            hex_text(&bytes),
162        ),
163    };
164    let content_id = datum.content_id()?;
165    let fence = InjectionFence::for_content(&content_id);
166    Ok(BridgeCallArgument::new(
167        symbol_from_name(name),
168        codec.clone(),
169        media,
170        content_id_string(&content_id),
171        fence.wrap(name, &body),
172    ))
173}
174
175fn answer_packet(
176    cx: &mut Cx,
177    book: &BridgeBook,
178    parent: &BridgePacket,
179    response: &ModelResponse,
180) -> Result<std::result::Result<BridgePacket, AskFailure>> {
181    let contract = output_contract_for_packet(parent)?;
182    let answer = match decode_terminal_answer(cx, response, &contract)? {
183        Ok(answer) => answer,
184        Err(failure) => return Ok(Err(failure)),
185    };
186    if let Err(failure) = validate_answer(cx, &contract, &answer)? {
187        return Ok(Err(failure));
188    }
189    let packet = stamp_packet_cid(&BridgePacket {
190        header: BridgeHeader {
191            cid: None,
192            move_kind: Symbol::new("reply"),
193            from: parent
194                .header
195                .to
196                .first()
197                .cloned()
198                .unwrap_or_else(|| "model".to_owned()),
199            to: vec![parent.header.from.clone()],
200            role: Symbol::new("implementer"),
201            parents: parent_token(parent).into_iter().collect(),
202            task: Symbol::new("A1"),
203            output: Symbol::new("A1"),
204            ceiling: Vec::new(),
205            context: Vec::new(),
206            provenance: BridgeProvenance::default(),
207        },
208        body: vec![BridgePart {
209            id: Symbol::new("A1"),
210            kind: Symbol::qualified("bridge", "Return"),
211            payload: answer,
212        }],
213        warrant: None,
214    })?;
215    let report = rx_check(cx, book, &packet, Some(parent))?;
216    if !report.accepted() {
217        return Err(sim_kernel::Error::Eval(format!(
218            "bridge ask reply failed rx check: {:?}",
219            report.obligations
220        )));
221    }
222    Ok(Ok(packet))
223}
224
225fn decode_terminal_answer(
226    cx: &mut Cx,
227    response: &ModelResponse,
228    contract: &OutputContract,
229) -> Result<std::result::Result<Expr, AskFailure>> {
230    let input = match terminal_model_content(response) {
231        Ok(Expr::String(text)) => Input::Text(text.clone()),
232        Ok(Expr::Bytes(bytes)) => Input::Bytes(bytes.clone()),
233        Ok(Expr::Map(_)) => match field(terminal_model_content(response)?, "text") {
234            Some(Expr::String(text)) => Input::Text(text.clone()),
235            _ => {
236                return Ok(Err(AskFailure::Decode {
237                    codec: contract.codec.clone(),
238                    message: "terminal content map must carry text".to_owned(),
239                }));
240            }
241        },
242        Ok(other) => {
243            return Ok(Err(AskFailure::Decode {
244                codec: contract.codec.clone(),
245                message: format!("terminal content must be text or bytes, found {other:?}"),
246            }));
247        }
248        Err(err) => {
249            return Ok(Err(AskFailure::Decode {
250                codec: contract.codec.clone(),
251                message: err.to_string(),
252            }));
253        }
254    };
255    if let Input::Text(text) = &input
256        && let Some(failure) = grammar_check_failure(cx, contract, text)?
257    {
258        return Ok(Err(failure));
259    }
260    match decode_with_codec(cx, &contract.codec, input, ReadPolicy::default()) {
261        Ok(answer) => Ok(Ok(answer)),
262        Err(err) => Ok(Err(AskFailure::Decode {
263            codec: contract.codec.clone(),
264            message: err.to_string(),
265        })),
266    }
267}
268
269fn grammar_check_failure(
270    cx: &mut Cx,
271    contract: &OutputContract,
272    text: &str,
273) -> Result<Option<AskFailure>> {
274    if contract.grammar.is_none()
275        && contract.grammar_dialect.is_none()
276        && contract.grammar_graph.is_none()
277    {
278        return Ok(None);
279    }
280    let Some(shape) = shape_from_contract_expr(&contract.shape_expr) else {
281        return Ok(Some(AskFailure::Shape {
282            expected: format!("{:?}", contract.shape_expr),
283            diagnostics: vec!["unsupported return Shape expression".to_owned()],
284        }));
285    };
286    let decoded = match decode_with_codec(
287        cx,
288        &contract.codec,
289        Input::Text(text.to_owned()),
290        ReadPolicy::default(),
291    ) {
292        Ok(decoded) => decoded,
293        Err(err) => {
294            return Ok(Some(AskFailure::Decode {
295                codec: contract.codec.clone(),
296                message: err.to_string(),
297            }));
298        }
299    };
300    let matched = shape.check_expr(cx, &decoded)?;
301    if matched.accepted {
302        Ok(None)
303    } else {
304        Ok(Some(AskFailure::Shape {
305            expected: format!("{:?}", contract.shape_expr),
306            diagnostics: matched
307                .diagnostics
308                .iter()
309                .map(|diagnostic| diagnostic.message.clone())
310                .collect(),
311        }))
312    }
313}
314
315fn validate_answer(
316    cx: &mut Cx,
317    contract: &OutputContract,
318    answer: &Expr,
319) -> Result<std::result::Result<(), AskFailure>> {
320    let Some(shape) = shape_from_contract_expr(&contract.shape_expr) else {
321        return Ok(Err(AskFailure::Shape {
322            expected: format!("{:?}", contract.shape_expr),
323            diagnostics: vec!["unsupported return Shape expression".to_owned()],
324        }));
325    };
326    let shape_ref = shape_value(Symbol::qualified("bridge", "AskReturn"), shape);
327    let value = cx.factory().expr(answer.clone())?;
328    let matched = check_value_report(cx, &shape_ref, value)?;
329    if matched.accepted {
330        Ok(Ok(()))
331    } else {
332        Ok(Err(AskFailure::Shape {
333            expected: format!("{:?}", contract.shape_expr),
334            diagnostics: matched
335                .diagnostics
336                .iter()
337                .map(|diagnostic| diagnostic.message.clone())
338                .collect(),
339        }))
340    }
341}
342
343fn repair_packet_for_failure(
344    cx: &mut Cx,
345    packet: &BridgePacket,
346    failure: &AskFailure,
347    attempt: u8,
348) -> Result<BridgePacket> {
349    let mut repaired = packet.canonicalized();
350    for part in &mut repaired.body {
351        if part.kind != Symbol::qualified("bridge", "Call") {
352            continue;
353        }
354        let payload = BridgeCallPayload::from_expr(&part.payload)?.with_arg(pack_argument(
355            cx,
356            &format!("repair-{attempt}"),
357            &ask_default_codec(),
358            &failure.to_expr(),
359        )?);
360        part.payload = payload.to_expr();
361        return Ok(repaired);
362    }
363    Err(sim_kernel::Error::Eval(
364        "bridge ask repair requires a Call part".to_owned(),
365    ))
366}
367
368fn symbol_from_name(name: &str) -> Symbol {
369    match name.split_once('/') {
370        Some((namespace, name)) if !namespace.is_empty() && !name.is_empty() => {
371            Symbol::qualified(namespace, name)
372        }
373        _ => Symbol::new(name),
374    }
375}
376
377fn hex_text(bytes: &[u8]) -> String {
378    bytes
379        .iter()
380        .map(|byte| format!("{byte:02x}"))
381        .collect::<String>()
382}