Skip to main content

phi_ext/phi/
mod.rs

1//! Author-facing SDK for a Phi PXB extension binary — a Rust port of
2//! `ext/go/phi`.
3//!
4//! Build an [`Extension`], register tools / slash commands / event handlers,
5//! then call [`Extension::run`] to speak PXB on stdin/stdout until shutdown:
6//!
7//! ```no_run
8//! use phi_ext::{phi, pxb};
9//!
10//! fn main() -> Result<(), phi::Error> {
11//!     let mut m = phi::Extension::new("hello", "0.1.0");
12//!     m.register_command("hello", phi::Command::new("Say hi", |_args, ctx| {
13//!         ctx.notify("info", "Hello!");
14//!         Ok(())
15//!     }));
16//!     m.subscribe(pxb::Event::SessionStart, |_ev| {});
17//!     m.run()
18//! }
19//! ```
20//!
21//! The run loop is single-threaded. Command handlers receive a [`Context`]
22//! whose methods (`notify`, `confirm`, `submit`, …) forward to the host over
23//! the same pipe — the borrow checker enforces at compile time what the Go
24//! SDK's mutexes enforce at runtime.
25
26use std::collections::HashMap;
27use std::io;
28
29use crate::pxb;
30
31pub use crate::pxb::Error;
32
33mod schema;
34pub use schema::Schema;
35
36type Rd = io::StdinLock<'static>;
37type Wr = io::StdoutLock<'static>;
38type EventHandlers = HashMap<u16, Box<dyn FnMut(pxb::EventNotify)>>;
39
40/// Host metadata filled by the hello handshake (and refreshed by
41/// `SessionMeta` pushes).
42#[derive(Debug, Clone, Default)]
43pub struct HostInfo {
44    pub cwd: String,
45    pub session_id: String,
46    pub extension_dir: String,
47    pub phi_version: String,
48}
49
50/// An LLM-callable tool. `schema` is a typed JSON Schema for parameters
51/// (same role as Go's `Parameters` / Codex's schemars-generated input schema).
52#[allow(clippy::type_complexity)] // execute signature mirrors the Go SDK contract
53pub struct Tool {
54    pub name: String,
55    pub description: String,
56    pub schema: Schema,
57    /// Host RPC wait for `execute`, in seconds. `0` = host default (30s).
58    pub timeout_sec: u32,
59    pub execute: Box<dyn FnMut(&[u8]) -> Result<ToolResult, String>>,
60}
61
62impl Tool {
63    pub fn new(
64        name: impl Into<String>,
65        description: impl Into<String>,
66        schema: impl Into<Schema>,
67        execute: impl FnMut(&[u8]) -> Result<ToolResult, String> + 'static,
68    ) -> Self {
69        Self {
70            name: name.into(),
71            description: description.into(),
72            schema: schema.into(),
73            timeout_sec: 0,
74            execute: Box::new(execute),
75        }
76    }
77
78    /// Sets how long the host waits for this tool's result (1–3600; host clamps).
79    pub fn timeout_sec(mut self, secs: u32) -> Self {
80        self.timeout_sec = secs;
81        self
82    }
83}
84
85/// Outcome of a [`Tool`] execution.
86#[derive(Debug, Clone, Default)]
87pub struct ToolResult {
88    pub content: String,
89    pub detail: String,
90    pub output: String,
91}
92
93/// A slash command. The handler receives the raw argument string and a
94/// [`Context`] for host interaction (notify / confirm / submit / …).
95#[allow(clippy::type_complexity)] // handler signature mirrors the Go SDK contract
96pub struct Command {
97    pub description: String,
98    pub handler: Box<dyn FnMut(&str, &mut Context<'_>) -> Result<(), String>>,
99}
100
101impl Command {
102    pub fn new(
103        description: impl Into<String>,
104        handler: impl FnMut(&str, &mut Context<'_>) -> Result<(), String> + 'static,
105    ) -> Self {
106        Self {
107            description: description.into(),
108            handler: Box::new(handler),
109        }
110    }
111}
112
113/// Modal yes/no dialog shown by the host.
114#[derive(Debug, Clone, Default)]
115pub struct ConfirmRequest {
116    pub title: String,
117    pub message: String,
118    pub yes: String, // host default: "Yes"
119    pub no: String,  // host default: "No"
120    pub danger: bool,
121}
122
123/// The user's choice for a [`ConfirmRequest`].
124#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
125pub struct ConfirmReply {
126    pub ok: bool,
127}
128
129// ── Intercept event payloads (mirror `ext/types.go`) ────────────────────────
130
131#[derive(Debug, Clone, Default)]
132pub struct ToolCallEvent {
133    pub tool_name: String,
134    pub tool_call_id: String,
135    pub input: Vec<u8>,
136}
137
138#[derive(Debug, Clone, Default)]
139pub struct ToolCallResult {
140    pub block: bool,
141    pub reason: String,
142    /// `Some` rewrites the tool input; `None` keeps it.
143    pub input: Option<Vec<u8>>,
144    pub context: String,
145}
146
147#[derive(Debug, Clone, Default)]
148pub struct ToolResultEvent {
149    pub tool_name: String,
150    pub tool_call_id: String,
151    pub input: Vec<u8>,
152    pub content: String,
153    pub is_error: bool,
154    pub err: String,
155}
156
157#[derive(Debug, Clone, Default)]
158pub struct ToolResultResult {
159    /// `Some` rewrites the tool output; `None` keeps it.
160    pub content: Option<String>,
161    pub context: String,
162    /// Ends the agent loop.
163    pub stop: bool,
164    pub reason: String,
165}
166
167#[derive(Debug, Clone, Default)]
168pub struct BeforeAgentStartEvent {
169    pub prompt: String,
170}
171
172#[derive(Debug, Clone, Default)]
173pub struct BeforeAgentStartResult {
174    /// `Some` replaces the user prompt.
175    pub prompt: Option<String>,
176    pub system_prompt_append: String,
177}
178
179#[derive(Debug, Clone, Default)]
180pub struct SessionBeforeSwitchEvent {
181    pub reason: String,
182    pub target_session_id: String,
183}
184
185#[derive(Debug, Clone, Default)]
186pub struct SessionBeforeSwitchResult {
187    pub cancel: bool,
188    pub reason: String,
189    pub toast: String,
190}
191
192#[derive(Debug, Clone, Default)]
193pub struct UserInputEvent {
194    pub text: String,
195}
196
197#[derive(Debug, Clone, Default)]
198pub struct UserInputResult {
199    /// `true` swallows the prompt (no agent loop).
200    pub handled: bool,
201    /// `Some` replaces the prompt text.
202    pub text: Option<String>,
203    pub reason: String,
204}
205
206#[derive(Debug, Clone, Default)]
207pub struct TurnStoppingEvent {
208    pub turn_index: u32,
209}
210
211#[derive(Debug, Clone, Default)]
212pub struct TurnStoppingResult {
213    /// Forces another agent step.
214    pub continue_: bool,
215    /// Injected as a user message when continuing.
216    pub message: String,
217    pub reason: String,
218}
219
220// ── Extension ───────────────────────────────────────────────────────────────
221
222/// Registered intercept / subscribe handlers. Kept as one struct so `run`
223/// can destructure the [`Extension`] into independent fields (each handler
224/// borrows only what it needs, no interior mutability).
225#[derive(Default)]
226struct Handlers {
227    tool_call: Option<Box<dyn FnMut(ToolCallEvent) -> Option<ToolCallResult>>>,
228    tool_result: Option<Box<dyn FnMut(ToolResultEvent) -> Option<ToolResultResult>>>,
229    before_agent_start:
230        Option<Box<dyn FnMut(BeforeAgentStartEvent) -> Option<BeforeAgentStartResult>>>,
231    session_before_switch:
232        Option<Box<dyn FnMut(SessionBeforeSwitchEvent) -> Option<SessionBeforeSwitchResult>>>,
233    user_input: Option<Box<dyn FnMut(UserInputEvent) -> Option<UserInputResult>>>,
234    turn_stopping: Option<Box<dyn FnMut(TurnStoppingEvent) -> Option<TurnStoppingResult>>>,
235    events: EventHandlers,
236}
237
238/// The author-facing registration surface for a PXB extension binary.
239pub struct Extension {
240    name: String,
241    version: String,
242    tools: Vec<Tool>,
243    commands: Vec<(String, Command)>,
244    events: Vec<pxb::Event>,
245    intercept: Vec<pxb::Event>,
246    handlers: Handlers,
247}
248
249impl Extension {
250    /// Constructs a module; call [`run`](Self::run) once everything is
251    /// registered.
252    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
253        Self {
254            name: name.into(),
255            version: version.into(),
256            tools: Vec::new(),
257            commands: Vec::new(),
258            events: Vec::new(),
259            intercept: Vec::new(),
260            handlers: Handlers::default(),
261        }
262    }
263
264    /// Adds an LLM-callable tool.
265    pub fn register_tool(&mut self, tool: Tool) {
266        if !tool.name.is_empty() {
267            self.tools.push(tool);
268        }
269    }
270
271    /// Adds a slash command (cannot override builtins at host level).
272    pub fn register_command(&mut self, name: impl Into<String>, cmd: Command) {
273        let name = name.into();
274        if !name.is_empty() && !self.commands.iter().any(|(n, _)| *n == name) {
275            self.commands.push((name, cmd));
276        }
277    }
278
279    /// Registers a pre-gate tool_call intercept.
280    pub fn on_tool_call(
281        &mut self,
282        f: impl FnMut(ToolCallEvent) -> Option<ToolCallResult> + 'static,
283    ) {
284        self.handlers.tool_call = Some(Box::new(f));
285        push_unique(&mut self.intercept, pxb::Event::ToolCall);
286    }
287
288    /// Registers a post-tool tool_result intercept.
289    pub fn on_tool_result(
290        &mut self,
291        f: impl FnMut(ToolResultEvent) -> Option<ToolResultResult> + 'static,
292    ) {
293        self.handlers.tool_result = Some(Box::new(f));
294        push_unique(&mut self.intercept, pxb::Event::ToolResult);
295    }
296
297    /// May append system prompt text before the agent loop starts.
298    pub fn on_before_agent_start(
299        &mut self,
300        f: impl FnMut(BeforeAgentStartEvent) -> Option<BeforeAgentStartResult> + 'static,
301    ) {
302        self.handlers.before_agent_start = Some(Box::new(f));
303        push_unique(&mut self.intercept, pxb::Event::BeforeAgentStart);
304    }
305
306    /// May cancel a session switch.
307    pub fn on_session_before_switch(
308        &mut self,
309        f: impl FnMut(SessionBeforeSwitchEvent) -> Option<SessionBeforeSwitchResult> + 'static,
310    ) {
311        self.handlers.session_before_switch = Some(Box::new(f));
312        push_unique(&mut self.intercept, pxb::Event::SessionBeforeSwitch);
313    }
314
315    /// May transform or swallow the user prompt before the agent loop.
316    pub fn on_user_input(
317        &mut self,
318        f: impl FnMut(UserInputEvent) -> Option<UserInputResult> + 'static,
319    ) {
320        self.handlers.user_input = Some(Box::new(f));
321        push_unique(&mut self.intercept, pxb::Event::UserInput);
322    }
323
324    /// May steer another agent step when the model stops with no tools.
325    pub fn on_turn_stopping(
326        &mut self,
327        f: impl FnMut(TurnStoppingEvent) -> Option<TurnStoppingResult> + 'static,
328    ) {
329        self.handlers.turn_stopping = Some(Box::new(f));
330        push_unique(&mut self.intercept, pxb::Event::TurnStopping);
331    }
332
333    /// Adds a fire-and-forget lifecycle listener; the payload is the wire
334    /// `EventNotify`. Unknown events are ignored.
335    pub fn subscribe(&mut self, event: pxb::Event, f: impl FnMut(pxb::EventNotify) + 'static) {
336        let code = event.code();
337        if code == 0 {
338            return;
339        }
340        push_unique(&mut self.events, event);
341        self.handlers.events.insert(code, Box::new(f));
342    }
343
344    /// Speaks PXB on stdin/stdout until the host shuts down.
345    pub fn run(self) -> Result<(), Error> {
346        let stdin = io::stdin();
347        let stdout = io::stdout();
348        let mut rd = stdin.lock();
349        let mut wr = stdout.lock();
350
351        let host = handshake(&mut rd, &mut wr, &self)?;
352        register(&mut wr, &self)?;
353
354        // Destructure so each handler owns only the state it mutates,
355        // instead of aliasing `self` (command handlers take `&mut` state).
356        let Extension {
357            tools,
358            commands,
359            handlers,
360            ..
361        } = self;
362        serve(&mut rd, &mut wr, host, tools, commands, handlers)
363    }
364}
365
366// ── Handshake, registration, and frame dispatch ──────────────────────────
367
368/// Exchanges the HELLO handshake and fills [`HostInfo`] from the host's ack.
369fn handshake(rd: &mut Rd, wr: &mut Wr, ext: &Extension) -> Result<HostInfo, Error> {
370    let mut caps = 0u32;
371    if !ext.commands.is_empty() {
372        caps |= pxb::CAP_COMMANDS;
373    }
374    if !ext.tools.is_empty() {
375        caps |= pxb::CAP_TOOLS;
376    }
377    if !ext.events.is_empty() {
378        caps |= pxb::CAP_EVENTS;
379    }
380    if !ext.intercept.is_empty() {
381        caps |= pxb::CAP_INTERCEPT;
382    }
383
384    let hello = pxb::encode_hello(&pxb::Hello {
385        name: ext.name.clone(),
386        version: ext.version.clone(),
387        caps,
388        protocol: pxb::PROTOCOL_VERSION,
389    });
390    pxb::write_frame(wr, pxb::TYPE_HELLO, 0, 0, &hello)?;
391
392    let f = pxb::read_frame(rd)?;
393    if f.header.typ != pxb::TYPE_HELLO_ACK {
394        return Err(Error::UnexpectedFrame {
395            want: "hello_ack",
396            got: f.header.typ,
397        });
398    }
399    let ack = pxb::decode_hello_ack(&f.body)?;
400    Ok(HostInfo {
401        cwd: ack.cwd,
402        session_id: ack.session_id,
403        extension_dir: ack.extension_dir,
404        phi_version: ack.phi_version,
405    })
406}
407
408/// Announces tools, commands, and subscription interest, then signals READY.
409fn register(wr: &mut Wr, ext: &Extension) -> Result<(), Error> {
410    for tool in &ext.tools {
411        let body = pxb::encode_register_tool(&pxb::RegisterTool {
412            name: tool.name.clone(),
413            description: tool.description.clone(),
414            schema_json: tool.schema.to_json_bytes(),
415            timeout_sec: tool.timeout_sec,
416        });
417        pxb::write_frame(wr, pxb::TYPE_REGISTER_TOOL, 0, 0, &body)?;
418    }
419    for (name, cmd) in &ext.commands {
420        let body = pxb::encode_register_command(&pxb::RegisterCommand {
421            name: name.clone(),
422            description: cmd.description.clone(),
423        });
424        pxb::write_frame(wr, pxb::TYPE_REGISTER_COMMAND, 0, 0, &body)?;
425    }
426    if !ext.events.is_empty() || !ext.intercept.is_empty() {
427        let body = pxb::encode_subscribe(&pxb::Subscribe {
428            events: ext.events.iter().map(|e| e.code()).collect(),
429            intercept: ext.intercept.iter().map(|e| e.code()).collect(),
430        });
431        pxb::write_frame(wr, pxb::TYPE_SUBSCRIBE, 0, 0, &body)?;
432    }
433    pxb::write_frame(wr, pxb::TYPE_READY, 0, 0, &[])
434}
435
436/// Dispatches frames until the host shuts down, handing each frame type to a
437/// focused handler that borrows only the state it mutates.
438fn serve(
439    rd: &mut Rd,
440    wr: &mut Wr,
441    mut host: HostInfo,
442    mut tools: Vec<Tool>,
443    mut commands: Vec<(String, Command)>,
444    mut handlers: Handlers,
445) -> Result<(), Error> {
446    let mut pending_submit: Option<String> = None;
447    let mut next_host_id: u32 = 0;
448
449    loop {
450        let f = pxb::read_frame(rd)?;
451        match pxb::FrameType::from_u16(f.header.typ) {
452            pxb::FrameType::Shutdown => {
453                pxb::write_frame(wr, pxb::TYPE_SHUTDOWN_ACK, 0, 0, &[])?;
454                return Ok(());
455            }
456            pxb::FrameType::CommandInvoked => serve_command(
457                rd,
458                wr,
459                &f,
460                &mut host,
461                &mut commands,
462                &mut handlers.events,
463                &mut pending_submit,
464                &mut next_host_id,
465            )?,
466            pxb::FrameType::ToolInvoke => serve_tool(wr, &f, &mut tools)?,
467            pxb::FrameType::Intercept => serve_intercept(wr, &f, &mut handlers)?,
468            pxb::FrameType::Event => {
469                if let Ok(ev) = pxb::decode_event_notify(&f.body) {
470                    dispatch_event(&mut handlers.events, ev);
471                }
472            }
473            pxb::FrameType::SessionMeta => {
474                if let Ok(meta) = pxb::decode_session_meta(&f.body) {
475                    apply_session_meta(&mut host, meta);
476                }
477            }
478            // Unknown frame types are already consumed by length; ignore.
479            _ => {}
480        }
481    }
482}
483
484/// Invokes a registered slash-command handler and replies with its outcome.
485/// An unknown command fails with "unknown command".
486#[allow(clippy::too_many_arguments)] // the loop lends each state piece separately
487fn serve_command(
488    rd: &mut Rd,
489    wr: &mut Wr,
490    frame: &pxb::Frame,
491    host: &mut HostInfo,
492    commands: &mut [(String, Command)],
493    events: &mut EventHandlers,
494    pending_submit: &mut Option<String>,
495    next_host_id: &mut u32,
496) -> Result<(), Error> {
497    let inv = pxb::decode_command_invoked(&frame.body)?;
498    let mut resp = pxb::CommandResponse {
499        ok: true,
500        ..Default::default()
501    };
502    if let Some((_, cmd)) = commands.iter_mut().find(|(n, _)| *n == inv.name) {
503        let mut ctx = Context {
504            cwd: host.cwd.clone(),
505            session_id: host.session_id.clone(),
506            has_ui: true,
507            rd,
508            wr,
509            host,
510            pending_submit,
511            next_host_id,
512            events,
513        };
514        if let Err(e) = (cmd.handler)(&inv.args, &mut ctx) {
515            resp.ok = false;
516            resp.error = e;
517        }
518    } else {
519        resp.ok = false;
520        resp.error = "unknown command".into();
521    }
522    resp.submit = pending_submit.take().unwrap_or_default();
523    let body = pxb::encode_command_response(&resp);
524    pxb::write_frame(
525        wr,
526        pxb::TYPE_COMMAND_RESPONSE,
527        frame.header.flags,
528        frame.header.id,
529        &body,
530    )?;
531    Ok(())
532}
533
534/// Executes a tool and replies with its result, or an error result when the
535/// tool is unknown or its handler failed.
536fn serve_tool(wr: &mut Wr, frame: &pxb::Frame, tools: &mut [Tool]) -> Result<(), Error> {
537    let inv = pxb::decode_tool_invoke(&frame.body)?;
538    let tr = match tools.iter_mut().find(|t| t.name == inv.name) {
539        Some(tool) => match (tool.execute)(&inv.args) {
540            Ok(res) => pxb::ToolResultMsg {
541                content: res.content,
542                detail: res.detail,
543                output: res.output,
544                ..Default::default()
545            },
546            Err(e) => tool_error(e),
547        },
548        None => tool_error("unknown tool"),
549    };
550    let body = pxb::encode_tool_result(&tr);
551    pxb::write_frame(
552        wr,
553        pxb::TYPE_TOOL_RESULT,
554        frame.header.flags,
555        frame.header.id,
556        &body,
557    )?;
558    Ok(())
559}
560
561/// Replies to one intercept request with the registered handler's result.
562fn serve_intercept(wr: &mut Wr, frame: &pxb::Frame, handlers: &mut Handlers) -> Result<(), Error> {
563    let req = pxb::decode_intercept_req(&frame.body)?;
564    let resp = handle_intercept(req, handlers);
565    let body = pxb::encode_intercept_resp(&resp);
566    pxb::write_frame(
567        wr,
568        pxb::TYPE_INTERCEPT_RESPONSE,
569        frame.header.flags,
570        frame.header.id,
571        &body,
572    )?;
573    Ok(())
574}
575
576/// Dispatches one intercept request to the registered handler. A missing
577/// handler (or a handler returning `None`) yields an empty response — the
578/// host treats that as "no change".
579fn handle_intercept(req: pxb::InterceptReq, handlers: &mut Handlers) -> pxb::InterceptResp {
580    let mut resp = pxb::InterceptResp::default();
581    match pxb::Event::from_code(req.event) {
582        pxb::Event::ToolCall => {
583            let Some(f) = handlers.tool_call.as_mut() else {
584                return resp;
585            };
586            let Some(r) = f(ToolCallEvent {
587                tool_name: req.tool_name,
588                tool_call_id: req.tool_call_id,
589                input: req.input,
590            }) else {
591                return resp;
592            };
593            resp.block = r.block;
594            resp.reason = r.reason;
595            resp.context = r.context;
596            if let Some(v) = r.input {
597                resp.input = v;
598            }
599        }
600        pxb::Event::ToolResult => {
601            let Some(f) = handlers.tool_result.as_mut() else {
602                return resp;
603            };
604            let Some(r) = f(ToolResultEvent {
605                tool_name: req.tool_name,
606                tool_call_id: req.tool_call_id,
607                input: req.input,
608                content: req.content,
609                is_error: req.is_error,
610                err: req.err_text,
611            }) else {
612                return resp;
613            };
614            resp.context = r.context;
615            resp.stop = r.stop;
616            resp.reason = r.reason;
617            if let Some(v) = r.content {
618                resp.content = v;
619            }
620        }
621        pxb::Event::BeforeAgentStart => {
622            let Some(f) = handlers.before_agent_start.as_mut() else {
623                return resp;
624            };
625            let Some(r) = f(BeforeAgentStartEvent { prompt: req.prompt }) else {
626                return resp;
627            };
628            resp.system_prompt_append = r.system_prompt_append;
629            if let Some(v) = r.prompt {
630                resp.prompt = v;
631            }
632        }
633        pxb::Event::SessionBeforeSwitch => {
634            let Some(f) = handlers.session_before_switch.as_mut() else {
635                return resp;
636            };
637            let Some(r) = f(SessionBeforeSwitchEvent {
638                reason: req.reason,
639                target_session_id: req.target_id,
640            }) else {
641                return resp;
642            };
643            resp.cancel = r.cancel;
644            resp.reason = r.reason;
645            resp.toast = r.toast;
646        }
647        pxb::Event::UserInput => {
648            let Some(f) = handlers.user_input.as_mut() else {
649                return resp;
650            };
651            let Some(r) = f(UserInputEvent { text: req.prompt }) else {
652                return resp;
653            };
654            resp.handled = r.handled;
655            resp.reason = r.reason;
656            if let Some(v) = r.text {
657                resp.prompt = v;
658            }
659        }
660        pxb::Event::TurnStopping => {
661            let Some(f) = handlers.turn_stopping.as_mut() else {
662                return resp;
663            };
664            let Some(r) = f(TurnStoppingEvent {
665                turn_index: req.turn_index,
666            }) else {
667                return resp;
668            };
669            resp.continue_ = r.continue_;
670            resp.prompt = r.message;
671            resp.reason = r.reason;
672        }
673        _ => {}
674    }
675    resp
676}
677
678/// Applies a session-meta push to host info; empty fields mean "no change".
679fn apply_session_meta(host: &mut HostInfo, meta: pxb::SessionMeta) {
680    if !meta.session_id.is_empty() {
681        host.session_id = meta.session_id;
682    }
683    if !meta.cwd.is_empty() {
684        host.cwd = meta.cwd;
685    }
686}
687
688/// Dispatches an event push to its subscriber, if one is registered.
689fn dispatch_event(handlers: &mut EventHandlers, ev: pxb::EventNotify) {
690    if let Some(handler) = handlers.get_mut(&ev.event) {
691        handler(ev);
692    }
693}
694
695/// An error tool result: the message goes to both `error` and `content` so
696/// the host surfaces it whichever field it renders.
697fn tool_error(message: impl Into<String>) -> pxb::ToolResultMsg {
698    let message = message.into();
699    pxb::ToolResultMsg {
700        is_error: true,
701        error: message.clone(),
702        content: message,
703        ..Default::default()
704    }
705}
706
707/// Interaction surface handed to command handlers. All host traffic goes
708/// over the same PXB pipe the run loop owns — a handler may block the loop
709/// (e.g. [`confirm`](Context::confirm) reads nested frames).
710pub struct Context<'a> {
711    pub cwd: String,
712    pub session_id: String,
713    pub has_ui: bool,
714    rd: &'a mut Rd,
715    wr: &'a mut Wr,
716    host: &'a mut HostInfo,
717    pending_submit: &'a mut Option<String>,
718    next_host_id: &'a mut u32,
719    events: &'a mut EventHandlers,
720}
721
722impl Context<'_> {
723    /// Pushes a toast to the host (`level`: `info` | `warning` | `error`).
724    pub fn notify(&mut self, level: &str, message: &str) {
725        let body = pxb::encode_notify(&pxb::NotifyMsg {
726            level: level.into(),
727            message: message.into(),
728            ..Default::default()
729        });
730        let _ = pxb::write_frame(self.wr, pxb::TYPE_NOTIFY, 0, 0, &body);
731    }
732
733    /// Updates the host footer extension status (empty text clears).
734    pub fn set_status(&mut self, text: &str) {
735        let body = pxb::encode_notify(&pxb::NotifyMsg {
736            status: text.into(),
737            status_set: true,
738            ..Default::default()
739        });
740        let _ = pxb::write_frame(self.wr, pxb::TYPE_NOTIFY, 0, 0, &body);
741    }
742
743    /// Queues a prompt for the host to send after the current slash command
744    /// returns.
745    pub fn submit(&mut self, text: &str) {
746        *self.pending_submit = Some(text.to_string());
747    }
748
749    /// Asks the host to enqueue a user turn (fire-and-forget). Safe to call
750    /// from command handlers on the PXB read loop.
751    pub fn send_user_message(&mut self, text: &str) {
752        if text.is_empty() {
753            return;
754        }
755        let body = pxb::encode_host_request(&pxb::HostRequest {
756            method: "send_user_message".into(),
757            arg: text.into(),
758        });
759        let _ = pxb::write_frame(self.wr, pxb::TYPE_HOST_REQUEST, 0, 0, &body);
760    }
761
762    /// Shows a yes/no dialog on the host and waits for the answer.
763    pub fn confirm(&mut self, title: &str, message: &str) -> ConfirmReply {
764        self.confirm_opts(ConfirmRequest {
765            title: title.into(),
766            message: message.into(),
767            ..Default::default()
768        })
769    }
770
771    /// [`confirm`](Self::confirm) with labels / danger styling.
772    pub fn confirm_opts(&mut self, req: ConfirmRequest) -> ConfirmReply {
773        let Some(id) = self.send_host_request("confirm", &confirm_request_json(&req)) else {
774            return ConfirmReply::default();
775        };
776        // Nested read: keep servicing SessionMeta pushes and subscribed
777        // events while waiting for the HostResult that matches our id.
778        loop {
779            let Ok(f) = pxb::read_frame(self.rd) else {
780                return ConfirmReply::default();
781            };
782            if let Some(reply) = self.nested_reply(f, id) {
783                return reply;
784            }
785        }
786    }
787
788    /// Writes a host request carrying a fresh id; returns that id. `None`
789    /// means the write failed — treat the host as gone.
790    fn send_host_request(&mut self, method: &str, arg: &str) -> Option<u32> {
791        *self.next_host_id = self.next_host_id.wrapping_add(1);
792        let id = *self.next_host_id;
793        let body = pxb::encode_host_request(&pxb::HostRequest {
794            method: method.into(),
795            arg: arg.into(),
796        });
797        if pxb::write_frame(self.wr, pxb::TYPE_HOST_REQUEST, pxb::FLAG_HAS_ID, id, &body).is_err() {
798            return None;
799        }
800        Some(id)
801    }
802
803    /// Handles one frame read while waiting on a host result. Returns the
804    /// [`ConfirmReply`] that ends the wait (a matching HostResult, Shutdown,
805    /// or a broken pipe); `None` means the frame was consumed internally.
806    fn nested_reply(&mut self, f: pxb::Frame, want_id: u32) -> Option<ConfirmReply> {
807        match pxb::FrameType::from_u16(f.header.typ) {
808            pxb::FrameType::HostResult => {
809                if f.header.flags & pxb::FLAG_HAS_ID == 0 || f.header.id != want_id {
810                    return None;
811                }
812                let Ok(res) = pxb::decode_host_result(&f.body) else {
813                    return Some(ConfirmReply::default());
814                };
815                Some(ConfirmReply { ok: res.ok })
816            }
817            pxb::FrameType::SessionMeta => {
818                if let Ok(meta) = pxb::decode_session_meta(&f.body) {
819                    apply_session_meta(self.host, meta);
820                }
821                None
822            }
823            pxb::FrameType::Event => {
824                if let Ok(ev) = pxb::decode_event_notify(&f.body) {
825                    dispatch_event(self.events, ev);
826                }
827                None
828            }
829            pxb::FrameType::Shutdown => {
830                let _ = pxb::write_frame(self.wr, pxb::TYPE_SHUTDOWN_ACK, 0, 0, &[]);
831                Some(ConfirmReply::default())
832            }
833            _ => None,
834        }
835    }
836}
837
838/// Encodes a [`ConfirmRequest`] as the JSON the host parses. The host
839/// unmarshals into Go's `ext.ConfirmRequest` (fields `Title`/`Message`/
840/// `Yes`/`No`/`Danger`), so key names and presence must match exactly —
841/// hence hand-rolled rather than a serde dependency.
842fn confirm_request_json(req: &ConfirmRequest) -> String {
843    let mut s = String::with_capacity(
844        64 + req.title.len() + req.message.len() + req.yes.len() + req.no.len(),
845    );
846    s.push_str(r#"{"Title":"#);
847    push_json_string(&mut s, &req.title);
848    s.push_str(r#","Message":"#);
849    push_json_string(&mut s, &req.message);
850    s.push_str(r#","Yes":"#);
851    push_json_string(&mut s, &req.yes);
852    s.push_str(r#","No":"#);
853    push_json_string(&mut s, &req.no);
854    s.push_str(r#","Danger":"#);
855    s.push_str(if req.danger { "true" } else { "false" });
856    s.push('}');
857    s
858}
859
860/// Appends `s` as a JSON string literal (control chars escaped; `<`, `>`,
861/// `&` are left as-is, which Go escapes but any JSON parser accepts).
862fn push_json_string(out: &mut String, s: &str) {
863    out.push('"');
864    for c in s.chars() {
865        match c {
866            '"' => out.push_str("\\\""),
867            '\\' => out.push_str("\\\\"),
868            '\n' => out.push_str("\\n"),
869            '\r' => out.push_str("\\r"),
870            '\t' => out.push_str("\\t"),
871            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
872            c => out.push(c),
873        }
874    }
875    out.push('"');
876}
877
878fn push_unique(xs: &mut Vec<pxb::Event>, v: pxb::Event) {
879    if !xs.contains(&v) {
880        xs.push(v);
881    }
882}
883
884#[cfg(test)]
885mod tests {
886    use super::*;
887
888    #[test]
889    fn confirm_json_matches_go_field_names() {
890        let req = ConfirmRequest {
891            title: "Delete?".into(),
892            message: "Remove /tmp/x".into(),
893            yes: "Delete".into(),
894            no: "Cancel".into(),
895            danger: true,
896        };
897        assert_eq!(
898            confirm_request_json(&req),
899            r#"{"Title":"Delete?","Message":"Remove /tmp/x","Yes":"Delete","No":"Cancel","Danger":true}"#
900        );
901    }
902
903    #[test]
904    fn confirm_json_escapes_quotes_and_controls() {
905        let req = ConfirmRequest {
906            title: "say \"hi\"\n".into(),
907            ..Default::default()
908        };
909        assert_eq!(
910            confirm_request_json(&req),
911            r#"{"Title":"say \"hi\"\n","Message":"","Yes":"","No":"","Danger":false}"#
912        );
913    }
914
915    #[test]
916    fn push_unique_keeps_first() {
917        let mut xs = Vec::new();
918        push_unique(&mut xs, pxb::Event::ToolCall);
919        push_unique(&mut xs, pxb::Event::ToolCall);
920        push_unique(&mut xs, pxb::Event::AgentEnd);
921        assert_eq!(xs, vec![pxb::Event::ToolCall, pxb::Event::AgentEnd]);
922    }
923}