dcp/context/
mod.rs

1//! Memory-mapped context for zero-copy sharing.
2//!
3//! Provides shared memory context for DCP protocol with atomic operations
4//! and memory fencing for thread safety.
5
6use std::sync::atomic::{AtomicU64, Ordering};
7
8use crate::DCPError;
9
10/// Magic number for context validation
11pub const CONTEXT_MAGIC: u64 = 0x4443505F_43545831; // "DCP_CTX1"
12
13/// Maximum number of tool states in context
14pub const MAX_TOOL_STATES: usize = 25;
15
16/// Context memory layout header offset
17pub const HEADER_OFFSET: usize = 0;
18/// Conversation ID offset
19pub const CONVERSATION_ID_OFFSET: usize = 8;
20/// Message count offset
21pub const MESSAGE_COUNT_OFFSET: usize = 16;
22/// Tool states offset
23pub const TOOL_STATES_OFFSET: usize = 24;
24/// Dynamic content offset (after tool states)
25pub const DYNAMIC_CONTENT_OFFSET: usize = TOOL_STATES_OFFSET + (MAX_TOOL_STATES * ToolState::SIZE);
26
27/// Tool state structure
28#[repr(C)]
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct ToolState {
31    /// Tool ID
32    pub tool_id: u32,
33    /// State flags
34    pub flags: u32,
35    /// Last invocation timestamp
36    pub last_invoked: u64,
37    /// State data (24 bytes)
38    pub data: [u8; 24],
39}
40
41impl ToolState {
42    /// Size of ToolState in bytes
43    pub const SIZE: usize = 40; // 4 + 4 + 8 + 24
44
45    /// Create a new empty tool state
46    pub fn new(tool_id: u32) -> Self {
47        Self {
48            tool_id,
49            flags: 0,
50            last_invoked: 0,
51            data: [0u8; 24],
52        }
53    }
54
55    /// Parse from bytes
56    #[inline(always)]
57    pub fn from_bytes(bytes: &[u8]) -> Result<&Self, DCPError> {
58        if bytes.len() < Self::SIZE {
59            return Err(DCPError::InsufficientData);
60        }
61        Ok(unsafe { &*(bytes.as_ptr() as *const Self) })
62    }
63
64    /// Serialize to bytes
65    #[inline(always)]
66    pub fn as_bytes(&self) -> &[u8] {
67        unsafe { std::slice::from_raw_parts(self as *const Self as *const u8, Self::SIZE) }
68    }
69}
70
71/// Context memory layout
72#[repr(C)]
73#[derive(Debug, Clone, Copy)]
74pub struct ContextLayout {
75    /// Magic + version (8 bytes)
76    pub header: u64,
77    /// Conversation ID (8 bytes)
78    pub conversation_id: u64,
79    /// Message count (8 bytes)
80    pub message_count: u64,
81    /// Tool states: 25 tools × 40 bytes = 1000 bytes
82    pub tool_states: [ToolState; MAX_TOOL_STATES],
83}
84
85impl ContextLayout {
86    /// Size of the layout in bytes
87    pub const SIZE: usize = 8 + 8 + 8 + (MAX_TOOL_STATES * ToolState::SIZE); // 1024 bytes
88
89    /// Create a new context layout
90    pub fn new(conversation_id: u64) -> Self {
91        Self {
92            header: CONTEXT_MAGIC,
93            conversation_id,
94            message_count: 0,
95            tool_states: [ToolState::new(0); MAX_TOOL_STATES],
96        }
97    }
98
99    /// Parse from bytes
100    #[inline(always)]
101    pub fn from_bytes(bytes: &[u8]) -> Result<&Self, DCPError> {
102        if bytes.len() < Self::SIZE {
103            return Err(DCPError::InsufficientData);
104        }
105        let layout = unsafe { &*(bytes.as_ptr() as *const Self) };
106        if layout.header != CONTEXT_MAGIC {
107            return Err(DCPError::InvalidMagic);
108        }
109        Ok(layout)
110    }
111
112    /// Serialize to bytes
113    #[inline(always)]
114    pub fn as_bytes(&self) -> &[u8] {
115        unsafe { std::slice::from_raw_parts(self as *const Self as *const u8, Self::SIZE) }
116    }
117}
118
119/// Memory-mapped context for zero-copy sharing
120pub struct DcpContext {
121    /// Shared memory buffer
122    buffer: Box<[u8]>,
123    /// Memory layout version
124    version: u32,
125}
126
127impl DcpContext {
128    /// Minimum buffer size
129    pub const MIN_SIZE: usize = ContextLayout::SIZE;
130
131    /// Create a new context with the given conversation ID
132    pub fn new(conversation_id: u64) -> Self {
133        let mut buffer = vec![0u8; Self::MIN_SIZE].into_boxed_slice();
134
135        // Initialize the layout
136        let layout = ContextLayout::new(conversation_id);
137        buffer[..ContextLayout::SIZE].copy_from_slice(layout.as_bytes());
138
139        Self { buffer, version: 1 }
140    }
141
142    /// Create a context from existing shared memory
143    pub fn from_shared(buffer: Box<[u8]>) -> Result<Self, DCPError> {
144        if buffer.len() < Self::MIN_SIZE {
145            return Err(DCPError::InsufficientData);
146        }
147
148        // Validate magic
149        let header = u64::from_le_bytes(buffer[0..8].try_into().unwrap());
150        if header != CONTEXT_MAGIC {
151            return Err(DCPError::InvalidMagic);
152        }
153
154        Ok(Self { buffer, version: 1 })
155    }
156
157    /// Get the raw buffer for sharing
158    pub fn as_bytes(&self) -> &[u8] {
159        &self.buffer
160    }
161
162    /// Get the buffer size
163    pub fn size(&self) -> usize {
164        self.buffer.len()
165    }
166
167    /// Get the context layout
168    #[inline(always)]
169    fn as_layout(&self) -> &ContextLayout {
170        unsafe { &*(self.buffer.as_ptr() as *const ContextLayout) }
171    }
172
173    /// Get the conversation ID
174    pub fn conversation_id(&self) -> u64 {
175        self.as_layout().conversation_id
176    }
177
178    /// Get the message count (atomic read)
179    pub fn message_count(&self) -> u64 {
180        let ptr = unsafe { self.buffer.as_ptr().add(MESSAGE_COUNT_OFFSET) as *const AtomicU64 };
181        unsafe { (*ptr).load(Ordering::Acquire) }
182    }
183
184    /// Increment the message count atomically
185    pub fn increment_message_count(&self) -> u64 {
186        let ptr = unsafe { self.buffer.as_ptr().add(MESSAGE_COUNT_OFFSET) as *const AtomicU64 };
187        unsafe { (*ptr).fetch_add(1, Ordering::AcqRel) + 1 }
188    }
189
190    /// Get the offset for a tool state by index
191    #[inline(always)]
192    fn tool_state_offset(&self, index: usize) -> usize {
193        TOOL_STATES_OFFSET + (index * ToolState::SIZE)
194    }
195
196    /// Zero-copy access to tool state by index
197    #[inline(always)]
198    pub fn get_tool_state(&self, index: usize) -> Option<&ToolState> {
199        if index >= MAX_TOOL_STATES {
200            return None;
201        }
202        let offset = self.tool_state_offset(index);
203        Some(unsafe { &*(self.buffer.as_ptr().add(offset) as *const ToolState) })
204    }
205
206    /// Find tool state by tool_id
207    pub fn find_tool_state(&self, tool_id: u32) -> Option<&ToolState> {
208        for i in 0..MAX_TOOL_STATES {
209            if let Some(state) = self.get_tool_state(i) {
210                if state.tool_id == tool_id {
211                    return Some(state);
212                }
213            }
214        }
215        None
216    }
217
218    /// Atomic update with memory fence
219    ///
220    /// # Safety
221    /// This function uses unsafe pointer operations but ensures memory safety
222    /// through proper bounds checking and memory fencing.
223    pub fn update_tool_state(&mut self, index: usize, state: &ToolState) -> Result<(), DCPError> {
224        if index >= MAX_TOOL_STATES {
225            return Err(DCPError::OutOfBounds);
226        }
227
228        let offset = self.tool_state_offset(index);
229
230        // Copy the state data
231        unsafe {
232            std::ptr::copy_nonoverlapping(
233                state as *const ToolState as *const u8,
234                self.buffer.as_mut_ptr().add(offset),
235                ToolState::SIZE,
236            );
237        }
238
239        // Memory fence for cross-thread visibility
240        std::sync::atomic::fence(Ordering::Release);
241
242        Ok(())
243    }
244
245    /// Update tool state by tool_id, finding or allocating a slot
246    pub fn set_tool_state(&mut self, state: &ToolState) -> Result<usize, DCPError> {
247        // First, try to find existing slot
248        for i in 0..MAX_TOOL_STATES {
249            if let Some(existing) = self.get_tool_state(i) {
250                if existing.tool_id == state.tool_id {
251                    self.update_tool_state(i, state)?;
252                    return Ok(i);
253                }
254            }
255        }
256
257        // Find empty slot (tool_id == 0)
258        for i in 0..MAX_TOOL_STATES {
259            if let Some(existing) = self.get_tool_state(i) {
260                if existing.tool_id == 0 {
261                    self.update_tool_state(i, state)?;
262                    return Ok(i);
263                }
264            }
265        }
266
267        Err(DCPError::OutOfBounds)
268    }
269
270    /// Clear a tool state slot
271    pub fn clear_tool_state(&mut self, index: usize) -> Result<(), DCPError> {
272        let empty = ToolState::new(0);
273        self.update_tool_state(index, &empty)
274    }
275
276    /// Get the version
277    pub fn version(&self) -> u32 {
278        self.version
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_tool_state_size() {
288        assert_eq!(std::mem::size_of::<ToolState>(), ToolState::SIZE);
289    }
290
291    #[test]
292    fn test_context_layout_size() {
293        assert_eq!(std::mem::size_of::<ContextLayout>(), ContextLayout::SIZE);
294    }
295
296    #[test]
297    fn test_context_creation() {
298        let ctx = DcpContext::new(12345);
299        assert_eq!(ctx.conversation_id(), 12345);
300        assert_eq!(ctx.message_count(), 0);
301    }
302
303    #[test]
304    fn test_message_count_increment() {
305        let ctx = DcpContext::new(1);
306        assert_eq!(ctx.message_count(), 0);
307        assert_eq!(ctx.increment_message_count(), 1);
308        assert_eq!(ctx.message_count(), 1);
309        assert_eq!(ctx.increment_message_count(), 2);
310        assert_eq!(ctx.message_count(), 2);
311    }
312
313    #[test]
314    fn test_tool_state_operations() {
315        let mut ctx = DcpContext::new(1);
316
317        let state = ToolState {
318            tool_id: 42,
319            flags: 0x1234,
320            last_invoked: 1234567890,
321            data: [0xAB; 24],
322        };
323
324        // Update tool state
325        ctx.update_tool_state(0, &state).unwrap();
326
327        // Read it back
328        let read_state = ctx.get_tool_state(0).unwrap();
329        assert_eq!(read_state.tool_id, 42);
330        assert_eq!(read_state.flags, 0x1234);
331        assert_eq!(read_state.last_invoked, 1234567890);
332        assert_eq!(read_state.data, [0xAB; 24]);
333    }
334
335    #[test]
336    fn test_find_tool_state() {
337        let mut ctx = DcpContext::new(1);
338
339        let state1 = ToolState {
340            tool_id: 100,
341            flags: 1,
342            last_invoked: 0,
343            data: [0; 24],
344        };
345        let state2 = ToolState {
346            tool_id: 200,
347            flags: 2,
348            last_invoked: 0,
349            data: [0; 24],
350        };
351
352        ctx.update_tool_state(0, &state1).unwrap();
353        ctx.update_tool_state(1, &state2).unwrap();
354
355        let found = ctx.find_tool_state(200).unwrap();
356        assert_eq!(found.tool_id, 200);
357        assert_eq!(found.flags, 2);
358
359        assert!(ctx.find_tool_state(999).is_none());
360    }
361
362    #[test]
363    fn test_set_tool_state() {
364        let mut ctx = DcpContext::new(1);
365
366        let state = ToolState {
367            tool_id: 42,
368            flags: 1,
369            last_invoked: 100,
370            data: [0; 24],
371        };
372
373        // First set should allocate slot 0
374        let idx = ctx.set_tool_state(&state).unwrap();
375        assert_eq!(idx, 0);
376
377        // Update same tool_id should use same slot
378        let updated = ToolState {
379            tool_id: 42,
380            flags: 2,
381            last_invoked: 200,
382            data: [0; 24],
383        };
384        let idx2 = ctx.set_tool_state(&updated).unwrap();
385        assert_eq!(idx2, 0);
386
387        // Verify update
388        let read = ctx.get_tool_state(0).unwrap();
389        assert_eq!(read.flags, 2);
390        assert_eq!(read.last_invoked, 200);
391    }
392
393    #[test]
394    fn test_out_of_bounds() {
395        let mut ctx = DcpContext::new(1);
396        let state = ToolState::new(1);
397
398        assert!(ctx.get_tool_state(MAX_TOOL_STATES).is_none());
399        assert_eq!(
400            ctx.update_tool_state(MAX_TOOL_STATES, &state),
401            Err(DCPError::OutOfBounds)
402        );
403    }
404
405    #[test]
406    fn test_from_shared() {
407        let ctx1 = DcpContext::new(99999);
408        let bytes = ctx1.as_bytes().to_vec().into_boxed_slice();
409
410        let ctx2 = DcpContext::from_shared(bytes).unwrap();
411        assert_eq!(ctx2.conversation_id(), 99999);
412    }
413
414    #[test]
415    fn test_invalid_magic() {
416        let mut buffer = vec![0u8; DcpContext::MIN_SIZE].into_boxed_slice();
417        buffer[0..8].copy_from_slice(&[0xFF; 8]);
418
419        let result = DcpContext::from_shared(buffer);
420        assert!(matches!(result, Err(DCPError::InvalidMagic)));
421    }
422
423    #[test]
424    fn test_tool_state_round_trip() {
425        let state = ToolState {
426            tool_id: 123,
427            flags: 0xDEAD,
428            last_invoked: 0xCAFEBABE,
429            data: [0x42; 24],
430        };
431
432        let bytes = state.as_bytes();
433        let parsed = ToolState::from_bytes(bytes).unwrap();
434
435        assert_eq!(parsed.tool_id, 123);
436        assert_eq!(parsed.flags, 0xDEAD);
437        assert_eq!(parsed.last_invoked, 0xCAFEBABE);
438        assert_eq!(parsed.data, [0x42; 24]);
439    }
440}
dcp/context/mod.rs

dcp/context/
mod.rs
