ktxd 0.1.1

A local Responses API compatibility proxy for Codex CLI and Chat Completions backends.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use crate::domain::{Session, TurnOutcome, TurnRecord};
use crate::error::{ProxyError, Result};
use crate::ids::{ResponseId, TurnId};
use async_trait::async_trait;
use serde_json::Value;
use std::collections::BTreeMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[async_trait]
pub trait SessionStore: Send + Sync {
    async fn get(&self, response_id: &ResponseId) -> Result<Option<Session>>;
    async fn put(&self, session: Session) -> Result<()>;
}

#[async_trait]
pub trait TurnRecordStore: Send + Sync {
    async fn put(&self, record: TurnRecord) -> Result<()>;
    async fn get(&self, turn_id: &TurnId) -> Result<Option<TurnRecord>>;
    async fn count(&self) -> usize;
}

#[derive(Debug, Default)]
pub struct MemoryStore {
    sessions: RwLock<BTreeMap<ResponseId, Session>>,
    responses: RwLock<BTreeMap<ResponseId, Value>>,
    turns: RwLock<BTreeMap<TurnId, TurnRecord>>,
}

impl MemoryStore {
    pub fn shared() -> Arc<Self> {
        Arc::new(Self::default())
    }

    pub async fn commit_completed(&self, session: Session, record: TurnRecord) -> Result<()> {
        if session.response_id != record.response_id {
            return Err(ProxyError::Internal(
                "session and turn response IDs do not match".to_string(),
            ));
        }
        if session.parent_response_id != record.parent_response_id {
            return Err(ProxyError::Internal(
                "session and turn parent response IDs do not match".to_string(),
            ));
        }
        if record.outcome != TurnOutcome::Completed {
            return Err(ProxyError::Internal(
                "completed commit requires a completed turn record".to_string(),
            ));
        }

        let mut sessions = self.sessions.write().await;
        let mut responses = self.responses.write().await;
        let mut turns = self.turns.write().await;
        if turns.contains_key(&record.turn_id) {
            return Err(ProxyError::Internal("duplicate turn record".to_string()));
        }
        if sessions.contains_key(&session.response_id)
            || responses.contains_key(&session.response_id)
        {
            return Err(ProxyError::Internal(
                "duplicate response record".to_string(),
            ));
        }
        responses.insert(
            session.response_id.clone(),
            session.final_response_json.clone(),
        );
        sessions.insert(session.response_id.clone(), session);
        turns.insert(record.turn_id.clone(), record);
        Ok(())
    }

    pub async fn commit_terminal(
        &self,
        response_id: ResponseId,
        response: Value,
        record: TurnRecord,
    ) -> Result<()> {
        if response_id != record.response_id {
            return Err(ProxyError::Internal(
                "response JSON and turn response IDs do not match".to_string(),
            ));
        }
        if record.outcome == TurnOutcome::Completed {
            return Err(ProxyError::Internal(
                "terminal commit cannot store a completed turn record".to_string(),
            ));
        }

        let sessions = self.sessions.read().await;
        let mut responses = self.responses.write().await;
        let mut turns = self.turns.write().await;
        if turns.contains_key(&record.turn_id) {
            return Err(ProxyError::Internal("duplicate turn record".to_string()));
        }
        if sessions.contains_key(&response_id) || responses.contains_key(&response_id) {
            return Err(ProxyError::Internal(
                "duplicate response record".to_string(),
            ));
        }
        responses.insert(response_id, response);
        turns.insert(record.turn_id.clone(), record);
        Ok(())
    }

    pub async fn get_response_json(&self, response_id: &ResponseId) -> Result<Option<Value>> {
        Ok(self.responses.read().await.get(response_id).cloned())
    }

    pub async fn put_response_json(&self, response_id: ResponseId, response: Value) -> Result<()> {
        self.responses.write().await.insert(response_id, response);
        Ok(())
    }
}

#[async_trait]
impl SessionStore for MemoryStore {
    async fn get(&self, response_id: &ResponseId) -> Result<Option<Session>> {
        Ok(self.sessions.read().await.get(response_id).cloned())
    }

    async fn put(&self, session: Session) -> Result<()> {
        self.sessions
            .write()
            .await
            .insert(session.response_id.clone(), session);
        Ok(())
    }
}

#[async_trait]
impl TurnRecordStore for MemoryStore {
    async fn put(&self, record: TurnRecord) -> Result<()> {
        let mut turns = self.turns.write().await;
        if turns.contains_key(&record.turn_id) {
            return Err(ProxyError::Internal("duplicate turn record".to_string()));
        }
        turns.insert(record.turn_id.clone(), record);
        Ok(())
    }

    async fn get(&self, turn_id: &TurnId) -> Result<Option<TurnRecord>> {
        Ok(self.turns.read().await.get(turn_id).cloned())
    }

    async fn count(&self) -> usize {
        self.turns.read().await.len()
    }
}