Skip to main content

kcode_k1_codex_runtime/
lib.rs

1pub use kcode_k1_codex_conversations::{
2    Config, DynamicTool, Error, ErrorKind, Event, ToolCall, ToolResult,
3};
4use kcode_k1_codex_conversations::{
5    Diagnostics, Pending, State, ToolToken, parse_rpc_id, thread_start_params, turn_start_params,
6};
7use kcode_k1_codex_events::{
8    DynamicToolCall as DecodedToolCall, Inbound, ResponseOutcome, RpcId, ScopedEvent, ScopedKind,
9    ServerError, TurnCompleted, TurnStatus, decode, rejection_payload,
10};
11use kcode_k1_codex_transport::{
12    ClientInfo, Config as WireConfig, Error as WireError, ErrorKind as WireKind, Transport,
13};
14use serde_json::{Value, json};
15use std::{
16    fmt,
17    sync::{Arc, Weak},
18    time::Duration,
19};
20use tokio::sync::{mpsc, oneshot};
21
22const EARLY_LIMIT: usize = 1024;
23type Events = mpsc::UnboundedSender<Event>;
24type Reply = oneshot::Sender<Result<(), Error>>;
25type CloseReply = Option<Reply>;
26type StartReply = oneshot::Sender<Result<u64, Error>>;
27
28enum Command {
29    Start {
30        key: String,
31        input: String,
32        events: Events,
33        reply: StartReply,
34    },
35    Respond {
36        key: String,
37        turn: u64,
38        call: String,
39        result: ToolResult,
40        reply: Reply,
41    },
42    Abandon {
43        key: String,
44        turn: u64,
45    },
46    Close {
47        key: String,
48        reply: Reply,
49    },
50}
51
52struct Client {
53    commands: mpsc::UnboundedSender<Command>,
54    diagnostics: Diagnostics,
55}
56
57#[derive(Clone)]
58pub struct Adapter {
59    client: Arc<Client>,
60}
61
62impl Adapter {
63    pub async fn open(config: Config) -> Result<Self, Error> {
64        open(config).await
65    }
66
67    pub async fn start_turn(
68        &self,
69        conversation_key: impl Into<String>,
70        input: impl Into<String>,
71    ) -> Result<Turn, Error> {
72        let key = conversation_key.into();
73        let (events, receiver) = mpsc::unbounded_channel();
74        let (reply, answer) = oneshot::channel();
75        self.client
76            .commands
77            .send(Command::Start {
78                key: key.clone(),
79                input: input.into(),
80                events,
81                reply,
82            })
83            .map_err(|_| self.unavailable())?;
84        let serial = answer.await.map_err(|_| self.unavailable())??;
85        Ok(Turn {
86            key,
87            serial,
88            client: Arc::downgrade(&self.client),
89            events: receiver,
90            diagnostics: self.client.diagnostics.clone(),
91            terminal: false,
92        })
93    }
94
95    pub async fn close_conversation(&self, key: impl Into<String>) -> Result<(), Error> {
96        let (reply, answer) = oneshot::channel();
97        self.client
98            .commands
99            .send(Command::Close {
100                key: key.into(),
101                reply,
102            })
103            .map_err(|_| self.unavailable())?;
104        answer.await.map_err(|_| self.unavailable())?
105    }
106
107    pub fn diagnostics(&self) -> Vec<u8> {
108        self.client.diagnostics.snapshot()
109    }
110
111    fn unavailable(&self) -> Error {
112        self.client
113            .diagnostics
114            .error(ErrorKind::Unavailable, "Codex app-server is unavailable")
115    }
116}
117
118pub struct Turn {
119    key: String,
120    serial: u64,
121    client: Weak<Client>,
122    events: mpsc::UnboundedReceiver<Event>,
123    diagnostics: Diagnostics,
124    terminal: bool,
125}
126
127impl fmt::Debug for Turn {
128    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
129        formatter
130            .debug_struct("Turn")
131            .field("key", &self.key)
132            .field("serial", &self.serial)
133            .finish_non_exhaustive()
134    }
135}
136
137impl Turn {
138    pub async fn next_event(&mut self) -> Option<Event> {
139        let event = self.events.recv().await;
140        if event
141            .as_ref()
142            .is_none_or(|event| matches!(event, Event::Done | Event::Error(_)))
143        {
144            self.terminal = true;
145        }
146        event
147    }
148
149    pub fn try_next_event(&mut self) -> Result<Option<Event>, Error> {
150        match self.events.try_recv() {
151            Ok(event) => {
152                if matches!(event, Event::Done | Event::Error(_)) {
153                    self.terminal = true;
154                }
155                Ok(Some(event))
156            }
157            Err(mpsc::error::TryRecvError::Empty) => Ok(None),
158            Err(mpsc::error::TryRecvError::Disconnected) => {
159                self.terminal = true;
160                Err(self.unavailable())
161            }
162        }
163    }
164
165    pub async fn respond(
166        &self,
167        call_id: impl Into<String>,
168        result: ToolResult,
169    ) -> Result<(), Error> {
170        let client = self.client.upgrade().ok_or_else(|| self.unavailable())?;
171        let (reply, answer) = oneshot::channel();
172        client
173            .commands
174            .send(Command::Respond {
175                key: self.key.clone(),
176                turn: self.serial,
177                call: call_id.into(),
178                result,
179                reply,
180            })
181            .map_err(|_| self.unavailable())?;
182        answer.await.map_err(|_| self.unavailable())?
183    }
184
185    fn unavailable(&self) -> Error {
186        self.diagnostics
187            .error(ErrorKind::Unavailable, "Codex app-server is unavailable")
188    }
189}
190
191impl Drop for Turn {
192    fn drop(&mut self) {
193        if !self.terminal
194            && let Some(client) = self.client.upgrade()
195        {
196            let _ = client.commands.send(Command::Abandon {
197                key: self.key.clone(),
198                turn: self.serial,
199            });
200        }
201    }
202}
203
204struct Actor {
205    transport: Transport,
206    commands: mpsc::UnboundedReceiver<Command>,
207    config: Config,
208    diagnostics: Diagnostics,
209    state: State<Events, CloseReply>,
210}
211
212async fn open(config: Config) -> Result<Adapter, Error> {
213    config.validate()?;
214    let mut wire = WireConfig::new(config.executable.clone(), config.working_directory.clone());
215    wire.reasoning_effort = config.reasoning_effort.clone();
216    wire.client = ClientInfo {
217        name: "kcode-k1-codex-adapter".into(),
218        title: "K1 Codex Adapter".into(),
219        version: "0.2.1".into(),
220    };
221    let transport = Transport::open(wire).await.map_err(map_wire)?;
222    let diagnostics = Diagnostics::new(transport.diagnostics());
223    let (commands, receiver) = mpsc::unbounded_channel();
224    let client = Arc::new(Client {
225        commands,
226        diagnostics: diagnostics.clone(),
227    });
228    let mut actor = Actor {
229        transport,
230        commands: receiver,
231        config,
232        diagnostics,
233        state: State::default(),
234    };
235    tokio::spawn(async move {
236        let outcome = actor.run().await;
237        actor.refresh();
238        let mut error = outcome
239            .err()
240            .unwrap_or_else(|| actor.error(ErrorKind::Unavailable, "Codex app-server was closed"));
241        let _ = actor.reject_all("adapter stopped").await;
242        let sinks = actor.state.take_sinks();
243        actor.refresh();
244        error.diagnostics = actor.diagnostics.snapshot();
245        for sink in sinks {
246            let _ = sink.send(Event::Error(error.clone()));
247        }
248        let _ = actor.transport.close().await;
249        actor.refresh();
250    });
251    Ok(Adapter { client })
252}
253
254impl Actor {
255    async fn run(&mut self) -> Result<(), Error> {
256        let mut tick = tokio::time::interval(Duration::from_millis(20));
257        loop {
258            tokio::select! {
259                command = self.commands.recv() => match command {
260                    Some(command) => self.command(command).await?,
261                    None => return Ok(()),
262                },
263                message = self.transport.receive() => self.message(message.map_err(map_wire)?).await?,
264                _ = tick.tick() => self.refresh(),
265            }
266        }
267    }
268
269    async fn command(&mut self, command: Command) -> Result<(), Error> {
270        self.refresh();
271        match command {
272            Command::Start {
273                key,
274                input,
275                events,
276                reply,
277            } => {
278                if self.start_pending(&key) {
279                    let _ = reply.send(Err(
280                        self.error(ErrorKind::Busy, "conversation already has an active turn")
281                    ));
282                    return Ok(());
283                }
284                let serial = match self.state.begin_turn(key.clone(), events) {
285                    Ok(serial) => serial,
286                    Err(error) => {
287                        let _ = reply.send(Err(self.stamp(error)));
288                        return Ok(());
289                    }
290                };
291                let request = if let Some(thread) = self.state.thread(&key).map(str::to_owned) {
292                    self.request(
293                        "turn/start",
294                        turn_start_params(&thread, input),
295                        Pending::Turn {
296                            key: key.clone(),
297                            serial,
298                        },
299                    )
300                    .await
301                } else {
302                    self.request(
303                        "thread/start",
304                        thread_start_params(&self.config),
305                        Pending::Thread {
306                            key: key.clone(),
307                            serial,
308                            input,
309                        },
310                    )
311                    .await
312                };
313                if let Err(error) = request {
314                    self.state.take_active(&key, serial);
315                    let _ = reply.send(Err(error.clone()));
316                    return Err(error);
317                }
318                if reply.send(Ok(serial)).is_err()
319                    && let Some(active) = self.state.take_active(&key, serial)
320                {
321                    self.reject_early(active.early, "turn start receiver was dropped")
322                        .await?;
323                }
324            }
325            Command::Respond {
326                key,
327                turn,
328                call,
329                result,
330                reply,
331            } => {
332                let token: ToolToken = (key, turn, call);
333                let Some(pending) = self.state.take_tool(&token) else {
334                    let _ = reply.send(Err(self.error(
335                        ErrorKind::InvalidToolResult,
336                        "tool call is not pending on this turn",
337                    )));
338                    return Ok(());
339                };
340                let message = json!({
341                    "id": pending.id,
342                    "result": {
343                        "contentItems": [{"type": "inputText", "text": result.output}],
344                        "success": result.success
345                    }
346                });
347                if let Err(error) = self.transport.send(&message).await.map_err(map_wire) {
348                    let _ = reply.send(Err(error.clone()));
349                    return Err(error);
350                }
351                let _ = reply.send(Ok(()));
352            }
353            Command::Abandon { key, turn } => self.abandon(&key, turn).await?,
354            Command::Close { key, reply } => {
355                if self.start_pending(&key) {
356                    let _ = reply.send(Err(
357                        self.error(ErrorKind::Busy, "conversation is active or already closing")
358                    ));
359                    return Ok(());
360                }
361                let thread = match self.state.begin_close(&key) {
362                    Ok(Some(thread)) => thread,
363                    Ok(None) => {
364                        let _ = reply.send(Ok(()));
365                        return Ok(());
366                    }
367                    Err(error) => {
368                        let _ = reply.send(Err(self.stamp(error)));
369                        return Ok(());
370                    }
371                };
372                if let Err(error) = self
373                    .request(
374                        "thread/unsubscribe",
375                        json!({"threadId": thread}),
376                        Pending::Close {
377                            key: key.clone(),
378                            thread,
379                            reply: Some(reply),
380                        },
381                    )
382                    .await
383                {
384                    self.state.cancel_close(&key);
385                    return Err(error);
386                }
387            }
388        }
389        Ok(())
390    }
391
392    async fn message(&mut self, message: Value) -> Result<(), Error> {
393        self.refresh();
394        match decode(message.clone()) {
395            Inbound::ClientResponse(response) => self.response(response.id, response.outcome).await,
396            Inbound::Scoped(event) => self.scoped(event, message).await,
397            Inbound::ResolvedRequest(id) => {
398                self.state
399                    .resolve_tool(&id.as_value())
400                    .map_err(|error| self.stamp(error))?;
401                Ok(())
402            }
403            Inbound::GlobalServerError(error) => Err(error.to_error(self.transport.diagnostics())),
404            Inbound::ModelReroute(reroute) => {
405                if let Some(id) = reroute.id {
406                    self.reject(&id, -32602, "model rerouting is unsupported")
407                        .await?;
408                }
409                Err(self.error(ErrorKind::Protocol, "model rerouting is unsupported"))
410            }
411            Inbound::UnsupportedServerRequest(request) => {
412                self.reject(&request.id, -32601, "unsupported server request")
413                    .await
414            }
415            Inbound::IgnoredNotification { .. } => Ok(()),
416            Inbound::Malformed(error) => self.malformed(&message, error).await,
417        }
418    }
419
420    async fn malformed(&mut self, message: &Value, error: Error) -> Result<(), Error> {
421        if message.get("method").and_then(Value::as_str) == Some("item/tool/call")
422            && let Some(id) = message.get("id")
423            && let Ok(id) = RpcId::try_from(id)
424        {
425            self.reject(&id, -32602, "malformed dynamic tool request")
426                .await?;
427        }
428        Err(self.stamp(error))
429    }
430
431    async fn response(&mut self, id: u64, outcome: ResponseOutcome) -> Result<(), Error> {
432        let pending = self
433            .state
434            .take_pending(id)
435            .map_err(|error| self.stamp(error))?;
436        let result = match outcome {
437            ResponseOutcome::Result(result) => result,
438            ResponseOutcome::Error(error) => {
439                return self
440                    .pending_failed(pending, error.to_error(self.transport.diagnostics()))
441                    .await;
442            }
443        };
444        match pending {
445            Pending::Thread { key, serial, input } => {
446                let Some(thread) = result
447                    .pointer("/thread/id")
448                    .and_then(Value::as_str)
449                    .map(str::to_owned)
450                else {
451                    let error = self.error(
452                        ErrorKind::Protocol,
453                        "thread/start response omitted thread.id",
454                    );
455                    return self.prestart_fail(&key, serial, error).await;
456                };
457                let live = match self
458                    .state
459                    .conversations
460                    .get(&key)
461                    .and_then(|conversation| conversation.active.as_ref())
462                {
463                    Some(active) if active.serial == serial => !active.cancelled,
464                    Some(_) => {
465                        return Err(self.error(
466                            ErrorKind::Protocol,
467                            "thread/start response crossed conversation turns",
468                        ));
469                    }
470                    None => false,
471                };
472                if let Err(error) = self.state.set_thread(&key, thread.clone()) {
473                    let error = self.stamp(error);
474                    return self.prestart_fail(&key, serial, error).await;
475                }
476                if !live {
477                    if let Some(active) = self.state.take_active(&key, serial) {
478                        self.reject_early(active.early, "turn was abandoned")
479                            .await?;
480                    }
481                    return self.retire_thread(&key).await;
482                }
483                self.request(
484                    "turn/start",
485                    turn_start_params(&thread, input),
486                    Pending::Turn { key, serial },
487                )
488                .await?;
489            }
490            Pending::Turn { key, serial } => {
491                let Some(turn) = result
492                    .pointer("/turn/id")
493                    .and_then(Value::as_str)
494                    .map(str::to_owned)
495                else {
496                    let error =
497                        self.error(ErrorKind::Protocol, "turn/start response omitted turn.id");
498                    return self.prestart_fail(&key, serial, error).await;
499                };
500                let Some(early) = self.state.set_native_turn(&key, serial, turn.clone()) else {
501                    return self.retire_turn(&key, &turn).await;
502                };
503                for message in early {
504                    match decode(message.clone()) {
505                        Inbound::Scoped(event) => self.scoped(event, message).await?,
506                        Inbound::Malformed(error) => return Err(self.stamp(error)),
507                        _ => {
508                            return Err(self.error(
509                                ErrorKind::Protocol,
510                                "buffered event changed classification",
511                            ));
512                        }
513                    }
514                }
515                self.interrupt(&key, serial).await?;
516            }
517            Pending::Close { key, thread, reply } => {
518                if let Err(error) = self.state.finish_close(&key, &thread) {
519                    let error = self.stamp(error);
520                    if let Some(reply) = reply {
521                        let _ = reply.send(Err(error.clone()));
522                    }
523                    return Err(error);
524                }
525                if let Some(reply) = reply {
526                    let _ = reply.send(Ok(()));
527                }
528            }
529            Pending::Interrupt => {}
530        }
531        Ok(())
532    }
533
534    async fn pending_failed(
535        &mut self,
536        pending: Pending<CloseReply>,
537        error: Error,
538    ) -> Result<(), Error> {
539        match pending {
540            Pending::Thread { key, serial, .. } | Pending::Turn { key, serial } => {
541                self.prestart_fail(&key, serial, error).await
542            }
543            Pending::Close { key, reply, .. } => {
544                self.state.cancel_close(&key);
545                if let Some(reply) = reply {
546                    let _ = reply.send(Err(error));
547                }
548                Ok(())
549            }
550            Pending::Interrupt => Err(error),
551        }
552    }
553
554    async fn scoped(&mut self, event: ScopedEvent, message: Value) -> Result<(), Error> {
555        let ScopedEvent { scope, kind } = event;
556        let Some(key) = self.state.owner(&scope.thread_id).map(str::to_owned) else {
557            self.reject_if_tool(&kind, "unknown thread").await?;
558            return Ok(());
559        };
560        let Some((serial, current, cancelled)) = self
561            .state
562            .conversations
563            .get(&key)
564            .and_then(|conversation| conversation.active.as_ref())
565            .map(|active| (active.serial, active.turn.clone(), active.cancelled))
566        else {
567            self.reject_if_tool(&kind, "no active turn").await?;
568            return Ok(());
569        };
570        let Some(current) = current else {
571            if cancelled && !matches!(&kind, ScopedKind::TurnCompleted(_) | ScopedKind::Error(_)) {
572                self.reject_if_tool(&kind, "turn was abandoned").await?;
573                return Ok(());
574            }
575            let active = self
576                .state
577                .conversations
578                .get_mut(&key)
579                .and_then(|conversation| conversation.active.as_mut())
580                .expect("active turn exists");
581            if active.early.len() >= EARLY_LIMIT {
582                self.reject_if_tool(&kind, "early event limit exceeded")
583                    .await?;
584                return Err(self.error(ErrorKind::Protocol, "early event limit exceeded"));
585            }
586            active.early.push(message);
587            return Ok(());
588        };
589        if current != scope.turn_id {
590            self.reject_if_tool(&kind, "request is outside the active turn")
591                .await?;
592            return Ok(());
593        }
594        if cancelled && !matches!(&kind, ScopedKind::TurnCompleted(_)) {
595            self.reject_if_tool(&kind, "turn was abandoned").await?;
596            return Ok(());
597        }
598        self.dispatch(key, serial, kind).await
599    }
600
601    async fn dispatch(&mut self, key: String, serial: u64, kind: ScopedKind) -> Result<(), Error> {
602        match kind {
603            ScopedKind::AgentTextDelta(delta) => {
604                if !self.emit(&key, serial, Event::TextDelta(delta)) {
605                    self.abandon(&key, serial).await?;
606                }
607            }
608            ScopedKind::DynamicToolCall(call) => self.tool_call(&key, serial, call).await?,
609            ScopedKind::TurnCompleted(completed) => self.complete(&key, serial, completed).await?,
610            ScopedKind::Error(error) => {
611                if let Some(active) = self
612                    .state
613                    .conversations
614                    .get_mut(&key)
615                    .and_then(|conversation| conversation.active.as_mut())
616                    .filter(|active| active.serial == serial)
617                {
618                    active.failure = Some(error.details);
619                }
620            }
621            ScopedKind::TurnStarted => {}
622        }
623        Ok(())
624    }
625
626    async fn tool_call(
627        &mut self,
628        key: &str,
629        serial: u64,
630        call: DecodedToolCall,
631    ) -> Result<(), Error> {
632        if !self.config.tools.iter().any(|tool| tool.name == call.name) {
633            self.reject(&call.rpc_id, -32602, "unconfigured dynamic tool")
634                .await?;
635            return self
636                .fail_turn(
637                    key,
638                    serial,
639                    self.error(ErrorKind::Protocol, "unconfigured dynamic tool requested"),
640                )
641                .await;
642        }
643        let id = call.rpc_id.as_value();
644        if let Err(error) = self
645            .state
646            .track_tool(key, serial, call.call_id.clone(), &id)
647        {
648            let error = self.stamp(error);
649            if error.message != "duplicate dynamic tool call id" {
650                return Err(error);
651            }
652            self.reject(&call.rpc_id, -32602, &error.message).await?;
653            return self.fail_turn(key, serial, error).await;
654        }
655        if !self.emit(
656            key,
657            serial,
658            Event::ToolCall(ToolCall {
659                call_id: call.call_id,
660                name: call.name,
661                arguments: call.arguments,
662            }),
663        ) {
664            self.abandon(key, serial).await?;
665        }
666        Ok(())
667    }
668
669    async fn complete(
670        &mut self,
671        key: &str,
672        serial: u64,
673        completed: TurnCompleted,
674    ) -> Result<(), Error> {
675        let pending = self
676            .state
677            .tools
678            .keys()
679            .filter(|(owner, turn, _)| owner == key && *turn == serial)
680            .count();
681        if pending != 0 {
682            self.reject_tools(key, serial, "turn completed before tool results")
683                .await?;
684        }
685        let Some(mut active) = self.state.take_active(key, serial) else {
686            return Ok(());
687        };
688        if active.cancelled {
689            return Ok(());
690        }
691        let failure = completed
692            .failure
693            .or_else(|| active.failure.take().map(ServerError::new));
694        let event = match completed.status {
695            TurnStatus::Completed if pending == 0 => Event::Done,
696            TurnStatus::Completed => Event::Error(self.error(
697                ErrorKind::Protocol,
698                "turn completed with pending tool calls",
699            )),
700            TurnStatus::Interrupted => {
701                Event::Error(self.error(ErrorKind::Interrupted, "Codex turn was interrupted"))
702            }
703            TurnStatus::Failed => Event::Error(
704                failure
705                    .unwrap_or_else(|| ServerError::new(Value::Null))
706                    .to_error(self.transport.diagnostics()),
707            ),
708        };
709        if let Some(events) = active.sink {
710            let _ = events.send(event);
711        }
712        Ok(())
713    }
714
715    async fn prestart_fail(&mut self, key: &str, serial: u64, error: Error) -> Result<(), Error> {
716        if let Some(active) = self.state.take_active(key, serial) {
717            if let Some(events) = active.sink {
718                let _ = events.send(Event::Error(error));
719            }
720            self.reject_early(active.early, "turn failed to start")
721                .await?;
722        }
723        Ok(())
724    }
725
726    async fn fail_turn(&mut self, key: &str, serial: u64, error: Error) -> Result<(), Error> {
727        if let Some(active) = self
728            .state
729            .conversations
730            .get_mut(key)
731            .and_then(|conversation| conversation.active.as_mut())
732            .filter(|active| active.serial == serial)
733        {
734            if let Some(events) = active.sink.take() {
735                let _ = events.send(Event::Error(error));
736            }
737            active.cancelled = true;
738        }
739        self.reject_tools(key, serial, "turn failed").await?;
740        self.interrupt(key, serial).await
741    }
742
743    async fn abandon(&mut self, key: &str, serial: u64) -> Result<(), Error> {
744        if let Some(active) = self
745            .state
746            .conversations
747            .get_mut(key)
748            .and_then(|conversation| conversation.active.as_mut())
749            .filter(|active| active.serial == serial)
750        {
751            active.sink.take();
752            active.cancelled = true;
753        }
754        self.reject_tools(key, serial, "turn was abandoned").await?;
755        self.interrupt(key, serial).await
756    }
757
758    async fn interrupt(&mut self, key: &str, serial: u64) -> Result<(), Error> {
759        let cancelled = self
760            .state
761            .conversations
762            .get(key)
763            .and_then(|conversation| conversation.active.as_ref())
764            .is_some_and(|active| active.serial == serial && active.cancelled);
765        if cancelled && let Some((thread, turn)) = self.state.interrupt_target(key, serial) {
766            self.request(
767                "turn/interrupt",
768                json!({"threadId": thread, "turnId": turn}),
769                Pending::Interrupt,
770            )
771            .await?;
772        }
773        Ok(())
774    }
775
776    async fn retire_turn(&mut self, key: &str, turn: &str) -> Result<(), Error> {
777        let thread = self
778            .state
779            .thread(key)
780            .ok_or_else(|| {
781                self.error(
782                    ErrorKind::Protocol,
783                    "orphan turn had no conversation thread",
784                )
785            })?
786            .to_owned();
787        self.request(
788            "turn/interrupt",
789            json!({"threadId": thread, "turnId": turn}),
790            Pending::Interrupt,
791        )
792        .await?;
793        self.retire_thread(key).await
794    }
795
796    async fn retire_thread(&mut self, key: &str) -> Result<(), Error> {
797        let thread = match self.state.begin_close(key) {
798            Ok(Some(thread)) => thread,
799            Ok(None) => return Ok(()),
800            Err(error) => return Err(self.stamp(error)),
801        };
802        self.request(
803            "thread/unsubscribe",
804            json!({"threadId": thread}),
805            Pending::Close {
806                key: key.to_owned(),
807                thread,
808                reply: None,
809            },
810        )
811        .await
812    }
813
814    async fn reject_early(&mut self, early: Vec<Value>, reason: &str) -> Result<(), Error> {
815        for message in early {
816            match decode(message) {
817                Inbound::Scoped(event) => self.reject_if_tool(&event.kind, reason).await?,
818                Inbound::Malformed(error) => return Err(self.stamp(error)),
819                _ => {
820                    return Err(
821                        self.error(ErrorKind::Protocol, "buffered event changed classification")
822                    );
823                }
824            }
825        }
826        Ok(())
827    }
828
829    async fn reject_tools(&mut self, key: &str, serial: u64, reason: &str) -> Result<(), Error> {
830        for (_, pending) in self.state.take_turn_tools(key, serial) {
831            let id = RpcId::try_from(&pending.id).map_err(|error| self.stamp(error))?;
832            self.reject(&id, -32800, reason).await?;
833        }
834        Ok(())
835    }
836
837    async fn reject_all(&mut self, reason: &str) -> Result<(), Error> {
838        let tokens: Vec<_> = self.state.tools.keys().cloned().collect();
839        for token in tokens {
840            if let Some(pending) = self.state.take_tool(&token) {
841                let id = RpcId::try_from(&pending.id).map_err(|error| self.stamp(error))?;
842                self.reject(&id, -32800, reason).await?;
843            }
844        }
845        Ok(())
846    }
847
848    async fn reject_if_tool(&mut self, kind: &ScopedKind, reason: &str) -> Result<(), Error> {
849        if let ScopedKind::DynamicToolCall(call) = kind {
850            self.reject(&call.rpc_id, -32602, reason).await?;
851        }
852        Ok(())
853    }
854
855    async fn reject(&mut self, id: &RpcId, code: i64, reason: &str) -> Result<(), Error> {
856        let (_, key) = parse_rpc_id(&id.as_value()).map_err(|error| self.stamp(error))?;
857        if self.state.rpc_ids.contains(&key) {
858            return Err(self.error(ErrorKind::Protocol, "duplicate app-server request id"));
859        }
860        self.transport
861            .send(&rejection_payload(id, code, reason))
862            .await
863            .map_err(map_wire)
864    }
865
866    async fn request(
867        &mut self,
868        method: &str,
869        params: Value,
870        pending: Pending<CloseReply>,
871    ) -> Result<(), Error> {
872        let id = self
873            .state
874            .allocate_request_id()
875            .map_err(|error| self.stamp(error))?;
876        self.transport
877            .send(&json!({"method": method, "id": id, "params": params}))
878            .await
879            .map_err(map_wire)?;
880        self.state
881            .insert_pending(id, pending)
882            .map_err(|error| self.stamp(error))
883    }
884
885    fn start_pending(&self, key: &str) -> bool {
886        self.state.pending.values().any(|pending| {
887            matches!(
888                pending,
889                Pending::Thread { key: owner, .. }
890                    | Pending::Turn { key: owner, .. }
891                    if owner == key
892            )
893        })
894    }
895
896    fn emit(&mut self, key: &str, serial: u64, event: Event) -> bool {
897        self.state
898            .conversations
899            .get_mut(key)
900            .and_then(|conversation| conversation.active.as_mut())
901            .filter(|active| active.serial == serial)
902            .and_then(|active| active.sink.as_ref())
903            .is_some_and(|events| events.send(event).is_ok())
904    }
905
906    fn error(&self, kind: ErrorKind, message: impl Into<String>) -> Error {
907        Error {
908            kind,
909            message: message.into(),
910            diagnostics: self.transport.diagnostics(),
911        }
912    }
913
914    fn stamp(&self, mut error: Error) -> Error {
915        error.diagnostics = self.transport.diagnostics();
916        error
917    }
918
919    fn refresh(&self) {
920        self.diagnostics.replace(self.transport.diagnostics());
921    }
922}
923
924fn map_wire(error: WireError) -> Error {
925    let kind = match error.kind {
926        WireKind::Protocol => ErrorKind::Protocol,
927        WireKind::Server => ErrorKind::Server,
928        WireKind::Spawn | WireKind::Timeout | WireKind::Io | WireKind::Closed => {
929            ErrorKind::Unavailable
930        }
931    };
932    Error {
933        kind,
934        message: error.message,
935        diagnostics: error.diagnostics,
936    }
937}