pub use kcode_k1_codex_conversations::{
Config, DynamicTool, Error, ErrorKind, Event, ToolCall, ToolResult,
};
use kcode_k1_codex_conversations::{
Diagnostics, Pending, State, ToolToken, parse_rpc_id, thread_start_params, turn_start_params,
};
use kcode_k1_codex_events::{
DynamicToolCall as DecodedToolCall, Inbound, ResponseOutcome, RpcId, ScopedEvent, ScopedKind,
ServerError, TurnCompleted, TurnStatus, decode, rejection_payload,
};
use kcode_k1_codex_transport::{
ClientInfo, Config as WireConfig, Error as WireError, ErrorKind as WireKind, Transport,
};
use serde_json::{Value, json};
use std::{
fmt,
sync::{Arc, Weak},
time::Duration,
};
use tokio::sync::{mpsc, oneshot};
const EARLY_LIMIT: usize = 1024;
type Events = mpsc::UnboundedSender<Event>;
type Reply = oneshot::Sender<Result<(), Error>>;
type CloseReply = Option<Reply>;
type StartReply = oneshot::Sender<Result<u64, Error>>;
enum Command {
Start {
key: String,
input: String,
events: Events,
reply: StartReply,
},
Respond {
key: String,
turn: u64,
call: String,
result: ToolResult,
reply: Reply,
},
Abandon {
key: String,
turn: u64,
},
Close {
key: String,
reply: Reply,
},
}
struct Client {
commands: mpsc::UnboundedSender<Command>,
diagnostics: Diagnostics,
}
#[derive(Clone)]
pub struct Adapter {
client: Arc<Client>,
}
impl Adapter {
pub async fn open(config: Config) -> Result<Self, Error> {
open(config).await
}
pub async fn start_turn(
&self,
conversation_key: impl Into<String>,
input: impl Into<String>,
) -> Result<Turn, Error> {
let key = conversation_key.into();
let (events, receiver) = mpsc::unbounded_channel();
let (reply, answer) = oneshot::channel();
self.client
.commands
.send(Command::Start {
key: key.clone(),
input: input.into(),
events,
reply,
})
.map_err(|_| self.unavailable())?;
let serial = answer.await.map_err(|_| self.unavailable())??;
Ok(Turn {
key,
serial,
client: Arc::downgrade(&self.client),
events: receiver,
diagnostics: self.client.diagnostics.clone(),
terminal: false,
})
}
pub async fn close_conversation(&self, key: impl Into<String>) -> Result<(), Error> {
let (reply, answer) = oneshot::channel();
self.client
.commands
.send(Command::Close {
key: key.into(),
reply,
})
.map_err(|_| self.unavailable())?;
answer.await.map_err(|_| self.unavailable())?
}
pub fn diagnostics(&self) -> Vec<u8> {
self.client.diagnostics.snapshot()
}
fn unavailable(&self) -> Error {
self.client
.diagnostics
.error(ErrorKind::Unavailable, "Codex app-server is unavailable")
}
}
pub struct Turn {
key: String,
serial: u64,
client: Weak<Client>,
events: mpsc::UnboundedReceiver<Event>,
diagnostics: Diagnostics,
terminal: bool,
}
impl fmt::Debug for Turn {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("Turn")
.field("key", &self.key)
.field("serial", &self.serial)
.finish_non_exhaustive()
}
}
impl Turn {
pub async fn next_event(&mut self) -> Option<Event> {
let event = self.events.recv().await;
if event
.as_ref()
.is_none_or(|event| matches!(event, Event::Done | Event::Error(_)))
{
self.terminal = true;
}
event
}
pub fn try_next_event(&mut self) -> Result<Option<Event>, Error> {
match self.events.try_recv() {
Ok(event) => {
if matches!(event, Event::Done | Event::Error(_)) {
self.terminal = true;
}
Ok(Some(event))
}
Err(mpsc::error::TryRecvError::Empty) => Ok(None),
Err(mpsc::error::TryRecvError::Disconnected) => {
self.terminal = true;
Err(self.unavailable())
}
}
}
pub async fn respond(
&self,
call_id: impl Into<String>,
result: ToolResult,
) -> Result<(), Error> {
let client = self.client.upgrade().ok_or_else(|| self.unavailable())?;
let (reply, answer) = oneshot::channel();
client
.commands
.send(Command::Respond {
key: self.key.clone(),
turn: self.serial,
call: call_id.into(),
result,
reply,
})
.map_err(|_| self.unavailable())?;
answer.await.map_err(|_| self.unavailable())?
}
fn unavailable(&self) -> Error {
self.diagnostics
.error(ErrorKind::Unavailable, "Codex app-server is unavailable")
}
}
impl Drop for Turn {
fn drop(&mut self) {
if !self.terminal
&& let Some(client) = self.client.upgrade()
{
let _ = client.commands.send(Command::Abandon {
key: self.key.clone(),
turn: self.serial,
});
}
}
}
struct Actor {
transport: Transport,
commands: mpsc::UnboundedReceiver<Command>,
config: Config,
diagnostics: Diagnostics,
state: State<Events, CloseReply>,
}
async fn open(config: Config) -> Result<Adapter, Error> {
config.validate()?;
let mut wire = WireConfig::new(config.executable.clone(), config.working_directory.clone());
wire.reasoning_effort = config.reasoning_effort.clone();
wire.client = ClientInfo {
name: "kcode-k1-codex-adapter".into(),
title: "K1 Codex Adapter".into(),
version: "0.2.1".into(),
};
let transport = Transport::open(wire).await.map_err(map_wire)?;
let diagnostics = Diagnostics::new(transport.diagnostics());
let (commands, receiver) = mpsc::unbounded_channel();
let client = Arc::new(Client {
commands,
diagnostics: diagnostics.clone(),
});
let mut actor = Actor {
transport,
commands: receiver,
config,
diagnostics,
state: State::default(),
};
tokio::spawn(async move {
let outcome = actor.run().await;
actor.refresh();
let mut error = outcome
.err()
.unwrap_or_else(|| actor.error(ErrorKind::Unavailable, "Codex app-server was closed"));
let _ = actor.reject_all("adapter stopped").await;
let sinks = actor.state.take_sinks();
actor.refresh();
error.diagnostics = actor.diagnostics.snapshot();
for sink in sinks {
let _ = sink.send(Event::Error(error.clone()));
}
let _ = actor.transport.close().await;
actor.refresh();
});
Ok(Adapter { client })
}
impl Actor {
async fn run(&mut self) -> Result<(), Error> {
let mut tick = tokio::time::interval(Duration::from_millis(20));
loop {
tokio::select! {
command = self.commands.recv() => match command {
Some(command) => self.command(command).await?,
None => return Ok(()),
},
message = self.transport.receive() => self.message(message.map_err(map_wire)?).await?,
_ = tick.tick() => self.refresh(),
}
}
}
async fn command(&mut self, command: Command) -> Result<(), Error> {
self.refresh();
match command {
Command::Start {
key,
input,
events,
reply,
} => {
if self.start_pending(&key) {
let _ = reply.send(Err(
self.error(ErrorKind::Busy, "conversation already has an active turn")
));
return Ok(());
}
let serial = match self.state.begin_turn(key.clone(), events) {
Ok(serial) => serial,
Err(error) => {
let _ = reply.send(Err(self.stamp(error)));
return Ok(());
}
};
let request = if let Some(thread) = self.state.thread(&key).map(str::to_owned) {
self.request(
"turn/start",
turn_start_params(&thread, input),
Pending::Turn {
key: key.clone(),
serial,
},
)
.await
} else {
self.request(
"thread/start",
thread_start_params(&self.config),
Pending::Thread {
key: key.clone(),
serial,
input,
},
)
.await
};
if let Err(error) = request {
self.state.take_active(&key, serial);
let _ = reply.send(Err(error.clone()));
return Err(error);
}
if reply.send(Ok(serial)).is_err()
&& let Some(active) = self.state.take_active(&key, serial)
{
self.reject_early(active.early, "turn start receiver was dropped")
.await?;
}
}
Command::Respond {
key,
turn,
call,
result,
reply,
} => {
let token: ToolToken = (key, turn, call);
let Some(pending) = self.state.take_tool(&token) else {
let _ = reply.send(Err(self.error(
ErrorKind::InvalidToolResult,
"tool call is not pending on this turn",
)));
return Ok(());
};
let message = json!({
"id": pending.id,
"result": {
"contentItems": [{"type": "inputText", "text": result.output}],
"success": result.success
}
});
if let Err(error) = self.transport.send(&message).await.map_err(map_wire) {
let _ = reply.send(Err(error.clone()));
return Err(error);
}
let _ = reply.send(Ok(()));
}
Command::Abandon { key, turn } => self.abandon(&key, turn).await?,
Command::Close { key, reply } => {
if self.start_pending(&key) {
let _ = reply.send(Err(
self.error(ErrorKind::Busy, "conversation is active or already closing")
));
return Ok(());
}
let thread = match self.state.begin_close(&key) {
Ok(Some(thread)) => thread,
Ok(None) => {
let _ = reply.send(Ok(()));
return Ok(());
}
Err(error) => {
let _ = reply.send(Err(self.stamp(error)));
return Ok(());
}
};
if let Err(error) = self
.request(
"thread/unsubscribe",
json!({"threadId": thread}),
Pending::Close {
key: key.clone(),
thread,
reply: Some(reply),
},
)
.await
{
self.state.cancel_close(&key);
return Err(error);
}
}
}
Ok(())
}
async fn message(&mut self, message: Value) -> Result<(), Error> {
self.refresh();
match decode(message.clone()) {
Inbound::ClientResponse(response) => self.response(response.id, response.outcome).await,
Inbound::Scoped(event) => self.scoped(event, message).await,
Inbound::ResolvedRequest(id) => {
self.state
.resolve_tool(&id.as_value())
.map_err(|error| self.stamp(error))?;
Ok(())
}
Inbound::GlobalServerError(error) => Err(error.to_error(self.transport.diagnostics())),
Inbound::ModelReroute(reroute) => {
if let Some(id) = reroute.id {
self.reject(&id, -32602, "model rerouting is unsupported")
.await?;
}
Err(self.error(ErrorKind::Protocol, "model rerouting is unsupported"))
}
Inbound::UnsupportedServerRequest(request) => {
self.reject(&request.id, -32601, "unsupported server request")
.await
}
Inbound::IgnoredNotification { .. } => Ok(()),
Inbound::Malformed(error) => self.malformed(&message, error).await,
}
}
async fn malformed(&mut self, message: &Value, error: Error) -> Result<(), Error> {
if message.get("method").and_then(Value::as_str) == Some("item/tool/call")
&& let Some(id) = message.get("id")
&& let Ok(id) = RpcId::try_from(id)
{
self.reject(&id, -32602, "malformed dynamic tool request")
.await?;
}
Err(self.stamp(error))
}
async fn response(&mut self, id: u64, outcome: ResponseOutcome) -> Result<(), Error> {
let pending = self
.state
.take_pending(id)
.map_err(|error| self.stamp(error))?;
let result = match outcome {
ResponseOutcome::Result(result) => result,
ResponseOutcome::Error(error) => {
return self
.pending_failed(pending, error.to_error(self.transport.diagnostics()))
.await;
}
};
match pending {
Pending::Thread { key, serial, input } => {
let Some(thread) = result
.pointer("/thread/id")
.and_then(Value::as_str)
.map(str::to_owned)
else {
let error = self.error(
ErrorKind::Protocol,
"thread/start response omitted thread.id",
);
return self.prestart_fail(&key, serial, error).await;
};
let live = match self
.state
.conversations
.get(&key)
.and_then(|conversation| conversation.active.as_ref())
{
Some(active) if active.serial == serial => !active.cancelled,
Some(_) => {
return Err(self.error(
ErrorKind::Protocol,
"thread/start response crossed conversation turns",
));
}
None => false,
};
if let Err(error) = self.state.set_thread(&key, thread.clone()) {
let error = self.stamp(error);
return self.prestart_fail(&key, serial, error).await;
}
if !live {
if let Some(active) = self.state.take_active(&key, serial) {
self.reject_early(active.early, "turn was abandoned")
.await?;
}
return self.retire_thread(&key).await;
}
self.request(
"turn/start",
turn_start_params(&thread, input),
Pending::Turn { key, serial },
)
.await?;
}
Pending::Turn { key, serial } => {
let Some(turn) = result
.pointer("/turn/id")
.and_then(Value::as_str)
.map(str::to_owned)
else {
let error =
self.error(ErrorKind::Protocol, "turn/start response omitted turn.id");
return self.prestart_fail(&key, serial, error).await;
};
let Some(early) = self.state.set_native_turn(&key, serial, turn.clone()) else {
return self.retire_turn(&key, &turn).await;
};
for message in early {
match decode(message.clone()) {
Inbound::Scoped(event) => self.scoped(event, message).await?,
Inbound::Malformed(error) => return Err(self.stamp(error)),
_ => {
return Err(self.error(
ErrorKind::Protocol,
"buffered event changed classification",
));
}
}
}
self.interrupt(&key, serial).await?;
}
Pending::Close { key, thread, reply } => {
if let Err(error) = self.state.finish_close(&key, &thread) {
let error = self.stamp(error);
if let Some(reply) = reply {
let _ = reply.send(Err(error.clone()));
}
return Err(error);
}
if let Some(reply) = reply {
let _ = reply.send(Ok(()));
}
}
Pending::Interrupt => {}
}
Ok(())
}
async fn pending_failed(
&mut self,
pending: Pending<CloseReply>,
error: Error,
) -> Result<(), Error> {
match pending {
Pending::Thread { key, serial, .. } | Pending::Turn { key, serial } => {
self.prestart_fail(&key, serial, error).await
}
Pending::Close { key, reply, .. } => {
self.state.cancel_close(&key);
if let Some(reply) = reply {
let _ = reply.send(Err(error));
}
Ok(())
}
Pending::Interrupt => Err(error),
}
}
async fn scoped(&mut self, event: ScopedEvent, message: Value) -> Result<(), Error> {
let ScopedEvent { scope, kind } = event;
let Some(key) = self.state.owner(&scope.thread_id).map(str::to_owned) else {
self.reject_if_tool(&kind, "unknown thread").await?;
return Ok(());
};
let Some((serial, current, cancelled)) = self
.state
.conversations
.get(&key)
.and_then(|conversation| conversation.active.as_ref())
.map(|active| (active.serial, active.turn.clone(), active.cancelled))
else {
self.reject_if_tool(&kind, "no active turn").await?;
return Ok(());
};
let Some(current) = current else {
if cancelled && !matches!(&kind, ScopedKind::TurnCompleted(_) | ScopedKind::Error(_)) {
self.reject_if_tool(&kind, "turn was abandoned").await?;
return Ok(());
}
let active = self
.state
.conversations
.get_mut(&key)
.and_then(|conversation| conversation.active.as_mut())
.expect("active turn exists");
if active.early.len() >= EARLY_LIMIT {
self.reject_if_tool(&kind, "early event limit exceeded")
.await?;
return Err(self.error(ErrorKind::Protocol, "early event limit exceeded"));
}
active.early.push(message);
return Ok(());
};
if current != scope.turn_id {
self.reject_if_tool(&kind, "request is outside the active turn")
.await?;
return Ok(());
}
if cancelled && !matches!(&kind, ScopedKind::TurnCompleted(_)) {
self.reject_if_tool(&kind, "turn was abandoned").await?;
return Ok(());
}
self.dispatch(key, serial, kind).await
}
async fn dispatch(&mut self, key: String, serial: u64, kind: ScopedKind) -> Result<(), Error> {
match kind {
ScopedKind::AgentTextDelta(delta) => {
if !self.emit(&key, serial, Event::TextDelta(delta)) {
self.abandon(&key, serial).await?;
}
}
ScopedKind::DynamicToolCall(call) => self.tool_call(&key, serial, call).await?,
ScopedKind::TurnCompleted(completed) => self.complete(&key, serial, completed).await?,
ScopedKind::Error(error) => {
if let Some(active) = self
.state
.conversations
.get_mut(&key)
.and_then(|conversation| conversation.active.as_mut())
.filter(|active| active.serial == serial)
{
active.failure = Some(error.details);
}
}
ScopedKind::TurnStarted => {}
}
Ok(())
}
async fn tool_call(
&mut self,
key: &str,
serial: u64,
call: DecodedToolCall,
) -> Result<(), Error> {
if !self.config.tools.iter().any(|tool| tool.name == call.name) {
self.reject(&call.rpc_id, -32602, "unconfigured dynamic tool")
.await?;
return self
.fail_turn(
key,
serial,
self.error(ErrorKind::Protocol, "unconfigured dynamic tool requested"),
)
.await;
}
let id = call.rpc_id.as_value();
if let Err(error) = self
.state
.track_tool(key, serial, call.call_id.clone(), &id)
{
let error = self.stamp(error);
if error.message != "duplicate dynamic tool call id" {
return Err(error);
}
self.reject(&call.rpc_id, -32602, &error.message).await?;
return self.fail_turn(key, serial, error).await;
}
if !self.emit(
key,
serial,
Event::ToolCall(ToolCall {
call_id: call.call_id,
name: call.name,
arguments: call.arguments,
}),
) {
self.abandon(key, serial).await?;
}
Ok(())
}
async fn complete(
&mut self,
key: &str,
serial: u64,
completed: TurnCompleted,
) -> Result<(), Error> {
let pending = self
.state
.tools
.keys()
.filter(|(owner, turn, _)| owner == key && *turn == serial)
.count();
if pending != 0 {
self.reject_tools(key, serial, "turn completed before tool results")
.await?;
}
let Some(mut active) = self.state.take_active(key, serial) else {
return Ok(());
};
if active.cancelled {
return Ok(());
}
let failure = completed
.failure
.or_else(|| active.failure.take().map(ServerError::new));
let event = match completed.status {
TurnStatus::Completed if pending == 0 => Event::Done,
TurnStatus::Completed => Event::Error(self.error(
ErrorKind::Protocol,
"turn completed with pending tool calls",
)),
TurnStatus::Interrupted => {
Event::Error(self.error(ErrorKind::Interrupted, "Codex turn was interrupted"))
}
TurnStatus::Failed => Event::Error(
failure
.unwrap_or_else(|| ServerError::new(Value::Null))
.to_error(self.transport.diagnostics()),
),
};
if let Some(events) = active.sink {
let _ = events.send(event);
}
Ok(())
}
async fn prestart_fail(&mut self, key: &str, serial: u64, error: Error) -> Result<(), Error> {
if let Some(active) = self.state.take_active(key, serial) {
if let Some(events) = active.sink {
let _ = events.send(Event::Error(error));
}
self.reject_early(active.early, "turn failed to start")
.await?;
}
Ok(())
}
async fn fail_turn(&mut self, key: &str, serial: u64, error: Error) -> Result<(), Error> {
if let Some(active) = self
.state
.conversations
.get_mut(key)
.and_then(|conversation| conversation.active.as_mut())
.filter(|active| active.serial == serial)
{
if let Some(events) = active.sink.take() {
let _ = events.send(Event::Error(error));
}
active.cancelled = true;
}
self.reject_tools(key, serial, "turn failed").await?;
self.interrupt(key, serial).await
}
async fn abandon(&mut self, key: &str, serial: u64) -> Result<(), Error> {
if let Some(active) = self
.state
.conversations
.get_mut(key)
.and_then(|conversation| conversation.active.as_mut())
.filter(|active| active.serial == serial)
{
active.sink.take();
active.cancelled = true;
}
self.reject_tools(key, serial, "turn was abandoned").await?;
self.interrupt(key, serial).await
}
async fn interrupt(&mut self, key: &str, serial: u64) -> Result<(), Error> {
let cancelled = self
.state
.conversations
.get(key)
.and_then(|conversation| conversation.active.as_ref())
.is_some_and(|active| active.serial == serial && active.cancelled);
if cancelled && let Some((thread, turn)) = self.state.interrupt_target(key, serial) {
self.request(
"turn/interrupt",
json!({"threadId": thread, "turnId": turn}),
Pending::Interrupt,
)
.await?;
}
Ok(())
}
async fn retire_turn(&mut self, key: &str, turn: &str) -> Result<(), Error> {
let thread = self
.state
.thread(key)
.ok_or_else(|| {
self.error(
ErrorKind::Protocol,
"orphan turn had no conversation thread",
)
})?
.to_owned();
self.request(
"turn/interrupt",
json!({"threadId": thread, "turnId": turn}),
Pending::Interrupt,
)
.await?;
self.retire_thread(key).await
}
async fn retire_thread(&mut self, key: &str) -> Result<(), Error> {
let thread = match self.state.begin_close(key) {
Ok(Some(thread)) => thread,
Ok(None) => return Ok(()),
Err(error) => return Err(self.stamp(error)),
};
self.request(
"thread/unsubscribe",
json!({"threadId": thread}),
Pending::Close {
key: key.to_owned(),
thread,
reply: None,
},
)
.await
}
async fn reject_early(&mut self, early: Vec<Value>, reason: &str) -> Result<(), Error> {
for message in early {
match decode(message) {
Inbound::Scoped(event) => self.reject_if_tool(&event.kind, reason).await?,
Inbound::Malformed(error) => return Err(self.stamp(error)),
_ => {
return Err(
self.error(ErrorKind::Protocol, "buffered event changed classification")
);
}
}
}
Ok(())
}
async fn reject_tools(&mut self, key: &str, serial: u64, reason: &str) -> Result<(), Error> {
for (_, pending) in self.state.take_turn_tools(key, serial) {
let id = RpcId::try_from(&pending.id).map_err(|error| self.stamp(error))?;
self.reject(&id, -32800, reason).await?;
}
Ok(())
}
async fn reject_all(&mut self, reason: &str) -> Result<(), Error> {
let tokens: Vec<_> = self.state.tools.keys().cloned().collect();
for token in tokens {
if let Some(pending) = self.state.take_tool(&token) {
let id = RpcId::try_from(&pending.id).map_err(|error| self.stamp(error))?;
self.reject(&id, -32800, reason).await?;
}
}
Ok(())
}
async fn reject_if_tool(&mut self, kind: &ScopedKind, reason: &str) -> Result<(), Error> {
if let ScopedKind::DynamicToolCall(call) = kind {
self.reject(&call.rpc_id, -32602, reason).await?;
}
Ok(())
}
async fn reject(&mut self, id: &RpcId, code: i64, reason: &str) -> Result<(), Error> {
let (_, key) = parse_rpc_id(&id.as_value()).map_err(|error| self.stamp(error))?;
if self.state.rpc_ids.contains(&key) {
return Err(self.error(ErrorKind::Protocol, "duplicate app-server request id"));
}
self.transport
.send(&rejection_payload(id, code, reason))
.await
.map_err(map_wire)
}
async fn request(
&mut self,
method: &str,
params: Value,
pending: Pending<CloseReply>,
) -> Result<(), Error> {
let id = self
.state
.allocate_request_id()
.map_err(|error| self.stamp(error))?;
self.transport
.send(&json!({"method": method, "id": id, "params": params}))
.await
.map_err(map_wire)?;
self.state
.insert_pending(id, pending)
.map_err(|error| self.stamp(error))
}
fn start_pending(&self, key: &str) -> bool {
self.state.pending.values().any(|pending| {
matches!(
pending,
Pending::Thread { key: owner, .. }
| Pending::Turn { key: owner, .. }
if owner == key
)
})
}
fn emit(&mut self, key: &str, serial: u64, event: Event) -> bool {
self.state
.conversations
.get_mut(key)
.and_then(|conversation| conversation.active.as_mut())
.filter(|active| active.serial == serial)
.and_then(|active| active.sink.as_ref())
.is_some_and(|events| events.send(event).is_ok())
}
fn error(&self, kind: ErrorKind, message: impl Into<String>) -> Error {
Error {
kind,
message: message.into(),
diagnostics: self.transport.diagnostics(),
}
}
fn stamp(&self, mut error: Error) -> Error {
error.diagnostics = self.transport.diagnostics();
error
}
fn refresh(&self) {
self.diagnostics.replace(self.transport.diagnostics());
}
}
fn map_wire(error: WireError) -> Error {
let kind = match error.kind {
WireKind::Protocol => ErrorKind::Protocol,
WireKind::Server => ErrorKind::Server,
WireKind::Spawn | WireKind::Timeout | WireKind::Io | WireKind::Closed => {
ErrorKind::Unavailable
}
};
Error {
kind,
message: error.message,
diagnostics: error.diagnostics,
}
}