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