1use 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#[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#[allow(clippy::type_complexity)] pub struct Tool {
54 pub name: String,
55 pub description: String,
56 pub schema: Schema,
57 pub timeout_sec: u32,
59 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 pub fn timeout_sec(mut self, secs: u32) -> Self {
83 self.timeout_sec = secs;
84 self
85 }
86
87 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#[derive(Debug, Clone, Default)]
96pub struct ToolResult {
97 pub content: String,
98 pub detail: String,
99 pub output: String,
100}
101
102#[allow(clippy::type_complexity)] pub 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#[derive(Debug, Clone, Default)]
124pub struct ConfirmRequest {
125 pub title: String,
126 pub message: String,
127 pub yes: String, pub no: String, pub danger: bool,
130}
131
132#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
134pub struct ConfirmReply {
135 pub ok: bool,
136}
137
138#[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 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 pub content: Option<String>,
170 pub context: String,
171 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 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 pub handled: bool,
210 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 pub continue_: bool,
224 pub message: String,
226 pub reason: String,
227}
228
229#[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
247pub 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 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 pub fn register_tool(&mut self, tool: Tool) {
275 if !tool.name.is_empty() {
276 self.tools.push(tool);
277 }
278 }
279
280 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 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 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 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 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 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 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 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 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 let Extension {
366 tools,
367 commands,
368 handlers,
369 ..
370 } = self;
371 serve(&mut rd, &mut wr, host, tools, commands, handlers)
372 }
373}
374
375fn 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
417fn 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
446fn 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 _ => {}
491 }
492 }
493}
494
495#[allow(clippy::too_many_arguments)] fn 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
545fn 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
572fn 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
592fn 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
607fn 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
709fn 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
719fn dispatch_event(handlers: &mut EventHandlers, ev: pxb::EventNotify) {
721 if let Some(handler) = handlers.get_mut(&ev.event) {
722 handler(ev);
723 }
724}
725
726fn 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
738pub 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 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 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 pub fn submit(&mut self, text: &str) {
777 *self.pending_submit = Some(text.to_string());
778 }
779
780 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 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 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 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 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 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
869fn 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
891fn 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}