Skip to main content

driven/dcp/
mod.rs

1//! DCP (Development Context Protocol) Integration
2//!
3//! This module provides integration with the DCP crate for high-performance
4//! AI tool communication. DCP is our MCP competitor with 10-1000x performance
5//! improvements while maintaining backward compatibility.
6//!
7//! ## Features
8//!
9//! - Binary Message Envelope (BME) protocol support
10//! - Zero-Copy Tool Invocation (ZCTI)
11//! - MCP compatibility via McpAdapter
12//! - Capability manifest support
13//! - Ed25519 signing for secure tool definitions
14//!
15//! ## Usage
16//!
17//! ```rust,ignore
18//! use driven::dcp::{DcpClient, DcpConfig};
19//!
20//! let config = DcpConfig::default();
21//! let mut client = DcpClient::new(config);
22//!
23//! // Connect to a DCP server
24//! client.connect("tcp://localhost:9000")?;
25//!
26//! // Invoke a tool
27//! let result = client.invoke_tool(tool_id, &args)?;
28//!
29//! // Or use MCP fallback
30//! let result = client.invoke_tool_mcp(&json_rpc_request)?;
31//! ```
32
33use crate::{DrivenError, Result};
34use std::collections::HashMap;
35
36// Re-export DCP types for convenience
37pub use dcp::binary::{
38    ArgType, BinaryMessageEnvelope, ChunkFlags, Flags, MessageType, SignedInvocation,
39    SignedToolDef, StreamChunk, ToolInvocation,
40};
41pub use dcp::capability::CapabilityManifest;
42pub use dcp::compat::adapter::AdapterError;
43pub use dcp::compat::json_rpc::RequestId;
44pub use dcp::compat::{JsonRpcError, JsonRpcParser, JsonRpcRequest, JsonRpcResponse, McpAdapter};
45pub use dcp::security::{Signer, Verifier};
46
47/// Zero-Copy Tool Invocation (ZCTI) builder
48///
49/// Provides a high-level API for building tool invocations with typed arguments.
50/// Arguments are encoded in a compact binary format for zero-copy passing.
51#[derive(Debug, Clone)]
52pub struct ZctiBuilder {
53    /// Tool ID
54    tool_id: u32,
55    /// Argument layout (types encoded in bitfield)
56    arg_layout: u64,
57    /// Argument data buffer
58    args_data: Vec<u8>,
59    /// Current argument index
60    arg_index: usize,
61}
62
63impl ZctiBuilder {
64    /// Create a new ZCTI builder for the given tool
65    pub fn new(tool_id: u32) -> Self {
66        Self {
67            tool_id,
68            arg_layout: 0,
69            args_data: Vec::new(),
70            arg_index: 0,
71        }
72    }
73
74    /// Add a null argument
75    pub fn add_null(mut self) -> Self {
76        self.set_arg_type(ArgType::Null);
77        self
78    }
79
80    /// Add a boolean argument
81    pub fn add_bool(mut self, value: bool) -> Self {
82        self.set_arg_type(ArgType::Bool);
83        self.args_data.push(if value { 1 } else { 0 });
84        self
85    }
86
87    /// Add an i32 argument
88    pub fn add_i32(mut self, value: i32) -> Self {
89        self.set_arg_type(ArgType::I32);
90        self.args_data.extend_from_slice(&value.to_le_bytes());
91        self
92    }
93
94    /// Add an i64 argument
95    pub fn add_i64(mut self, value: i64) -> Self {
96        self.set_arg_type(ArgType::I64);
97        self.args_data.extend_from_slice(&value.to_le_bytes());
98        self
99    }
100
101    /// Add an f64 argument
102    pub fn add_f64(mut self, value: f64) -> Self {
103        self.set_arg_type(ArgType::F64);
104        self.args_data.extend_from_slice(&value.to_le_bytes());
105        self
106    }
107
108    /// Add a string argument
109    pub fn add_string(mut self, value: &str) -> Self {
110        self.set_arg_type(ArgType::String);
111        // Length-prefixed string (4 bytes length + data)
112        let len = value.len() as u32;
113        self.args_data.extend_from_slice(&len.to_le_bytes());
114        self.args_data.extend_from_slice(value.as_bytes());
115        self
116    }
117
118    /// Add a bytes argument
119    pub fn add_bytes(mut self, value: &[u8]) -> Self {
120        self.set_arg_type(ArgType::Bytes);
121        // Length-prefixed bytes (4 bytes length + data)
122        let len = value.len() as u32;
123        self.args_data.extend_from_slice(&len.to_le_bytes());
124        self.args_data.extend_from_slice(value);
125        self
126    }
127
128    /// Add raw bytes without type encoding (for pre-encoded data)
129    pub fn add_raw(mut self, data: &[u8]) -> Self {
130        self.args_data.extend_from_slice(data);
131        self
132    }
133
134    /// Set the argument type at the current index
135    fn set_arg_type(&mut self, arg_type: ArgType) {
136        if self.arg_index < ToolInvocation::MAX_ARGS {
137            let shift = self.arg_index * 4;
138            self.arg_layout |= (arg_type as u64) << shift;
139            self.arg_index += 1;
140        }
141    }
142
143    /// Build the tool invocation
144    pub fn build(self) -> ToolInvocation {
145        ToolInvocation::new(
146            self.tool_id,
147            self.arg_layout,
148            0, // Offset is set when placed in shared memory
149            self.args_data.len() as u32,
150        )
151    }
152
153    /// Build and return both the invocation and the argument data
154    pub fn build_with_data(self) -> (ToolInvocation, Vec<u8>) {
155        let invocation = ToolInvocation::new(
156            self.tool_id,
157            self.arg_layout,
158            0,
159            self.args_data.len() as u32,
160        );
161        (invocation, self.args_data)
162    }
163
164    /// Get the current argument count
165    pub fn arg_count(&self) -> usize {
166        self.arg_index
167    }
168
169    /// Get the current data size
170    pub fn data_size(&self) -> usize {
171        self.args_data.len()
172    }
173}
174
175/// Zero-Copy Tool Invocation (ZCTI) reader
176///
177/// Provides a high-level API for reading arguments from a tool invocation.
178#[derive(Debug)]
179pub struct ZctiReader<'a> {
180    /// The tool invocation header
181    invocation: ToolInvocation,
182    /// The argument data
183    data: &'a [u8],
184    /// Current read offset
185    offset: usize,
186    /// Current argument index
187    arg_index: usize,
188}
189
190impl<'a> ZctiReader<'a> {
191    /// Create a new ZCTI reader
192    pub fn new(invocation: &ToolInvocation, data: &'a [u8]) -> Self {
193        Self {
194            invocation: *invocation,
195            data,
196            offset: 0,
197            arg_index: 0,
198        }
199    }
200
201    /// Get the tool ID
202    pub fn tool_id(&self) -> u32 {
203        self.invocation.tool_id
204    }
205
206    /// Get the total argument count
207    pub fn arg_count(&self) -> usize {
208        self.invocation.arg_count()
209    }
210
211    /// Get the type of the next argument
212    pub fn peek_type(&self) -> Option<ArgType> {
213        self.invocation.get_arg_type(self.arg_index)
214    }
215
216    /// Read a boolean argument
217    pub fn read_bool(&mut self) -> Result<bool> {
218        self.check_type(ArgType::Bool)?;
219        if self.offset >= self.data.len() {
220            return Err(DrivenError::InvalidBinary(
221                "Insufficient data for bool".to_string(),
222            ));
223        }
224        let value = self.data[self.offset] != 0;
225        self.offset += 1;
226        self.arg_index += 1;
227        Ok(value)
228    }
229
230    /// Read an i32 argument
231    pub fn read_i32(&mut self) -> Result<i32> {
232        self.check_type(ArgType::I32)?;
233        if self.offset + 4 > self.data.len() {
234            return Err(DrivenError::InvalidBinary(
235                "Insufficient data for i32".to_string(),
236            ));
237        }
238        let bytes: [u8; 4] = self.data[self.offset..self.offset + 4].try_into().unwrap();
239        let value = i32::from_le_bytes(bytes);
240        self.offset += 4;
241        self.arg_index += 1;
242        Ok(value)
243    }
244
245    /// Read an i64 argument
246    pub fn read_i64(&mut self) -> Result<i64> {
247        self.check_type(ArgType::I64)?;
248        if self.offset + 8 > self.data.len() {
249            return Err(DrivenError::InvalidBinary(
250                "Insufficient data for i64".to_string(),
251            ));
252        }
253        let bytes: [u8; 8] = self.data[self.offset..self.offset + 8].try_into().unwrap();
254        let value = i64::from_le_bytes(bytes);
255        self.offset += 8;
256        self.arg_index += 1;
257        Ok(value)
258    }
259
260    /// Read an f64 argument
261    pub fn read_f64(&mut self) -> Result<f64> {
262        self.check_type(ArgType::F64)?;
263        if self.offset + 8 > self.data.len() {
264            return Err(DrivenError::InvalidBinary(
265                "Insufficient data for f64".to_string(),
266            ));
267        }
268        let bytes: [u8; 8] = self.data[self.offset..self.offset + 8].try_into().unwrap();
269        let value = f64::from_le_bytes(bytes);
270        self.offset += 8;
271        self.arg_index += 1;
272        Ok(value)
273    }
274
275    /// Read a string argument
276    pub fn read_string(&mut self) -> Result<&'a str> {
277        self.check_type(ArgType::String)?;
278        if self.offset + 4 > self.data.len() {
279            return Err(DrivenError::InvalidBinary(
280                "Insufficient data for string length".to_string(),
281            ));
282        }
283        let len_bytes: [u8; 4] = self.data[self.offset..self.offset + 4].try_into().unwrap();
284        let len = u32::from_le_bytes(len_bytes) as usize;
285        self.offset += 4;
286
287        if self.offset + len > self.data.len() {
288            return Err(DrivenError::InvalidBinary(
289                "Insufficient data for string content".to_string(),
290            ));
291        }
292        let value = std::str::from_utf8(&self.data[self.offset..self.offset + len])
293            .map_err(|e| DrivenError::InvalidBinary(format!("Invalid UTF-8: {}", e)))?;
294        self.offset += len;
295        self.arg_index += 1;
296        Ok(value)
297    }
298
299    /// Read a bytes argument
300    pub fn read_bytes(&mut self) -> Result<&'a [u8]> {
301        self.check_type(ArgType::Bytes)?;
302        if self.offset + 4 > self.data.len() {
303            return Err(DrivenError::InvalidBinary(
304                "Insufficient data for bytes length".to_string(),
305            ));
306        }
307        let len_bytes: [u8; 4] = self.data[self.offset..self.offset + 4].try_into().unwrap();
308        let len = u32::from_le_bytes(len_bytes) as usize;
309        self.offset += 4;
310
311        if self.offset + len > self.data.len() {
312            return Err(DrivenError::InvalidBinary(
313                "Insufficient data for bytes content".to_string(),
314            ));
315        }
316        let value = &self.data[self.offset..self.offset + len];
317        self.offset += len;
318        self.arg_index += 1;
319        Ok(value)
320    }
321
322    /// Skip the current argument
323    pub fn skip(&mut self) -> Result<()> {
324        let arg_type = self
325            .peek_type()
326            .ok_or_else(|| DrivenError::InvalidBinary("No more arguments".to_string()))?;
327
328        match arg_type {
329            ArgType::Null => {
330                self.arg_index += 1;
331            }
332            ArgType::Bool => {
333                self.offset += 1;
334                self.arg_index += 1;
335            }
336            ArgType::I32 => {
337                self.offset += 4;
338                self.arg_index += 1;
339            }
340            ArgType::I64 | ArgType::F64 => {
341                self.offset += 8;
342                self.arg_index += 1;
343            }
344            ArgType::String | ArgType::Bytes => {
345                if self.offset + 4 > self.data.len() {
346                    return Err(DrivenError::InvalidBinary(
347                        "Insufficient data for length".to_string(),
348                    ));
349                }
350                let len_bytes: [u8; 4] =
351                    self.data[self.offset..self.offset + 4].try_into().unwrap();
352                let len = u32::from_le_bytes(len_bytes) as usize;
353                self.offset += 4 + len;
354                self.arg_index += 1;
355            }
356            ArgType::Array | ArgType::Object => {
357                // For complex types, we'd need additional metadata
358                return Err(DrivenError::InvalidBinary(
359                    "Complex types not yet supported".to_string(),
360                ));
361            }
362        }
363        Ok(())
364    }
365
366    /// Check if there are more arguments to read
367    pub fn has_more(&self) -> bool {
368        self.arg_index < self.arg_count()
369    }
370
371    /// Get the remaining data slice
372    pub fn remaining_data(&self) -> &'a [u8] {
373        &self.data[self.offset..]
374    }
375
376    /// Check that the next argument has the expected type
377    fn check_type(&self, expected: ArgType) -> Result<()> {
378        let actual = self
379            .peek_type()
380            .ok_or_else(|| DrivenError::InvalidBinary("No more arguments".to_string()))?;
381        if actual != expected {
382            return Err(DrivenError::InvalidBinary(format!(
383                "Type mismatch: expected {:?}, got {:?}",
384                expected, actual
385            )));
386        }
387        Ok(())
388    }
389}
390
391/// Shared memory buffer for zero-copy argument passing
392///
393/// This is a simple implementation that uses a Vec<u8> as the backing store.
394/// In a real implementation, this would use SharedArrayBuffer or mmap.
395#[derive(Debug)]
396pub struct SharedArgBuffer {
397    /// The backing buffer
398    data: Vec<u8>,
399    /// Current write offset
400    write_offset: usize,
401}
402
403impl SharedArgBuffer {
404    /// Create a new shared argument buffer with the given capacity
405    pub fn new(capacity: usize) -> Self {
406        Self {
407            data: vec![0u8; capacity],
408            write_offset: 0,
409        }
410    }
411
412    /// Create with default capacity (64KB)
413    pub fn with_default_capacity() -> Self {
414        Self::new(65536)
415    }
416
417    /// Write data to the buffer and return the offset
418    pub fn write(&mut self, data: &[u8]) -> Result<u32> {
419        if self.write_offset + data.len() > self.data.len() {
420            return Err(DrivenError::InvalidBinary("Buffer overflow".to_string()));
421        }
422        let offset = self.write_offset as u32;
423        self.data[self.write_offset..self.write_offset + data.len()].copy_from_slice(data);
424        self.write_offset += data.len();
425        Ok(offset)
426    }
427
428    /// Read data from the buffer at the given offset
429    pub fn read(&self, offset: u32, len: u32) -> Result<&[u8]> {
430        let start = offset as usize;
431        let end = start + len as usize;
432        if end > self.data.len() {
433            return Err(DrivenError::InvalidBinary("Read out of bounds".to_string()));
434        }
435        Ok(&self.data[start..end])
436    }
437
438    /// Get the current write offset
439    pub fn offset(&self) -> u32 {
440        self.write_offset as u32
441    }
442
443    /// Reset the buffer
444    pub fn reset(&mut self) {
445        self.write_offset = 0;
446    }
447
448    /// Get the underlying data
449    pub fn data(&self) -> &[u8] {
450        &self.data[..self.write_offset]
451    }
452
453    /// Get the capacity
454    pub fn capacity(&self) -> usize {
455        self.data.len()
456    }
457
458    /// Get the remaining capacity
459    pub fn remaining(&self) -> usize {
460        self.data.len() - self.write_offset
461    }
462}
463
464/// DCP client configuration
465#[derive(Debug, Clone)]
466pub struct DcpConfig {
467    /// Whether DCP is enabled
468    pub enabled: bool,
469    /// Whether to prefer DCP over MCP when available
470    pub prefer_dcp: bool,
471    /// DCP server endpoint (e.g., "tcp://localhost:9000")
472    pub endpoint: Option<String>,
473    /// Connection timeout in milliseconds
474    pub timeout_ms: u64,
475    /// Whether to enable signing
476    pub signing_enabled: bool,
477    /// Seed for Ed25519 signing (32 bytes)
478    pub signing_seed: Option<[u8; 32]>,
479}
480
481impl Default for DcpConfig {
482    fn default() -> Self {
483        Self {
484            enabled: true,
485            prefer_dcp: true,
486            endpoint: None,
487            timeout_ms: 5000,
488            signing_enabled: false,
489            signing_seed: None,
490        }
491    }
492}
493
494/// Connection state
495#[derive(Debug, Clone, Copy, PartialEq, Eq)]
496pub enum ConnectionState {
497    /// Not connected
498    Disconnected,
499    /// Connecting to server
500    Connecting,
501    /// Connected via DCP protocol
502    ConnectedDcp,
503    /// Connected via MCP fallback
504    ConnectedMcp,
505    /// Connection failed
506    Failed,
507}
508
509/// Protocol type
510#[derive(Debug, Clone, Copy, PartialEq, Eq)]
511pub enum Protocol {
512    /// DCP binary protocol
513    Dcp,
514    /// MCP JSON-RPC protocol
515    Mcp,
516}
517
518/// Result of a tool invocation
519#[derive(Debug)]
520pub enum InvocationResult {
521    /// DCP binary response
522    Dcp(Vec<u8>),
523    /// MCP JSON-RPC response
524    Mcp(String),
525}
526
527impl InvocationResult {
528    /// Check if this is a DCP result
529    pub fn is_dcp(&self) -> bool {
530        matches!(self, InvocationResult::Dcp(_))
531    }
532
533    /// Check if this is an MCP result
534    pub fn is_mcp(&self) -> bool {
535        matches!(self, InvocationResult::Mcp(_))
536    }
537
538    /// Get the DCP bytes if this is a DCP result
539    pub fn as_dcp(&self) -> Option<&[u8]> {
540        match self {
541            InvocationResult::Dcp(bytes) => Some(bytes),
542            _ => None,
543        }
544    }
545
546    /// Get the MCP JSON string if this is an MCP result
547    pub fn as_mcp(&self) -> Option<&str> {
548        match self {
549            InvocationResult::Mcp(json) => Some(json),
550            _ => None,
551        }
552    }
553}
554
555/// DCP client for AI tool communication
556///
557/// Provides high-performance binary protocol communication with DCP servers,
558/// with automatic fallback to MCP when DCP is unavailable.
559pub struct DcpClient {
560    /// Client configuration
561    config: DcpConfig,
562    /// Current connection state
563    state: ConnectionState,
564    /// MCP adapter for fallback
565    mcp_adapter: Option<McpAdapter>,
566    /// Capability manifest
567    capabilities: CapabilityManifest,
568    /// Ed25519 signer for secure tool definitions
569    signer: Option<Signer>,
570    /// Registered tools
571    tools: HashMap<u32, ToolDefinition>,
572    /// Next tool ID
573    next_tool_id: u32,
574}
575
576/// Tool definition
577#[derive(Debug, Clone)]
578pub struct ToolDefinition {
579    /// Tool ID
580    pub id: u32,
581    /// Tool name
582    pub name: String,
583    /// Tool description
584    pub description: String,
585    /// Schema hash (Blake3)
586    pub schema_hash: [u8; 32],
587    /// Required capabilities
588    pub capabilities: u64,
589}
590
591/// Statistics about the current capability manifest
592#[derive(Debug, Clone, Copy)]
593pub struct CapabilityStats {
594    /// Number of registered tools
595    pub tool_count: u32,
596    /// Number of registered resources
597    pub resource_count: u32,
598    /// Number of registered prompts
599    pub prompt_count: u32,
600    /// Number of enabled extensions
601    pub extension_count: u32,
602    /// Protocol version
603    pub version: u16,
604}
605
606/// DCP message representing a complete binary message with envelope and payload
607#[derive(Debug, Clone)]
608pub struct DcpMessage {
609    /// Message type
610    pub message_type: MessageType,
611    /// Flags (streaming, compressed, signed)
612    pub flags: u8,
613    /// Payload data
614    pub payload: Vec<u8>,
615}
616
617impl DcpMessage {
618    /// Create a new DCP message
619    pub fn new(message_type: MessageType, flags: u8, payload: Vec<u8>) -> Self {
620        Self {
621            message_type,
622            flags,
623            payload,
624        }
625    }
626
627    /// Create a tool invocation message
628    pub fn tool(payload: Vec<u8>) -> Self {
629        Self::new(MessageType::Tool, 0, payload)
630    }
631
632    /// Create a resource request message
633    pub fn resource(payload: Vec<u8>) -> Self {
634        Self::new(MessageType::Resource, 0, payload)
635    }
636
637    /// Create a prompt message
638    pub fn prompt(payload: Vec<u8>) -> Self {
639        Self::new(MessageType::Prompt, 0, payload)
640    }
641
642    /// Create a response message
643    pub fn response(payload: Vec<u8>) -> Self {
644        Self::new(MessageType::Response, 0, payload)
645    }
646
647    /// Create an error message
648    pub fn error(payload: Vec<u8>) -> Self {
649        Self::new(MessageType::Error, 0, payload)
650    }
651
652    /// Create a streaming message
653    pub fn stream(payload: Vec<u8>) -> Self {
654        Self::new(MessageType::Stream, Flags::STREAMING, payload)
655    }
656
657    /// Check if this is a streaming message
658    pub fn is_streaming(&self) -> bool {
659        self.flags & Flags::STREAMING != 0
660    }
661
662    /// Check if this message is compressed
663    pub fn is_compressed(&self) -> bool {
664        self.flags & Flags::COMPRESSED != 0
665    }
666
667    /// Check if this message is signed
668    pub fn is_signed(&self) -> bool {
669        self.flags & Flags::SIGNED != 0
670    }
671
672    /// Set the streaming flag
673    pub fn with_streaming(mut self) -> Self {
674        self.flags |= Flags::STREAMING;
675        self
676    }
677
678    /// Set the compressed flag
679    pub fn with_compressed(mut self) -> Self {
680        self.flags |= Flags::COMPRESSED;
681        self
682    }
683
684    /// Set the signed flag
685    pub fn with_signed(mut self) -> Self {
686        self.flags |= Flags::SIGNED;
687        self
688    }
689
690    /// Encode the message to bytes (envelope + payload)
691    pub fn encode(&self) -> Vec<u8> {
692        let envelope =
693            BinaryMessageEnvelope::new(self.message_type, self.flags, self.payload.len() as u32);
694
695        let mut result = Vec::with_capacity(BinaryMessageEnvelope::SIZE + self.payload.len());
696        result.extend_from_slice(envelope.as_bytes());
697        result.extend_from_slice(&self.payload);
698        result
699    }
700
701    /// Decode a message from bytes
702    pub fn decode(bytes: &[u8]) -> Result<Self> {
703        if bytes.len() < BinaryMessageEnvelope::SIZE {
704            return Err(DrivenError::InvalidBinary(
705                "Insufficient data for BME header".to_string(),
706            ));
707        }
708
709        let envelope = BinaryMessageEnvelope::from_bytes(bytes)
710            .map_err(|e| DrivenError::InvalidBinary(format!("Invalid BME: {:?}", e)))?;
711
712        let message_type = envelope.get_message_type().ok_or_else(|| {
713            DrivenError::InvalidBinary(format!("Unknown message type: {}", envelope.message_type))
714        })?;
715
716        let payload_len = envelope.payload_len as usize;
717        let total_len = BinaryMessageEnvelope::SIZE + payload_len;
718
719        if bytes.len() < total_len {
720            return Err(DrivenError::InvalidBinary(format!(
721                "Insufficient data: expected {} bytes, got {}",
722                total_len,
723                bytes.len()
724            )));
725        }
726
727        let payload = bytes[BinaryMessageEnvelope::SIZE..total_len].to_vec();
728
729        Ok(Self {
730            message_type,
731            flags: envelope.flags,
732            payload,
733        })
734    }
735}
736
737/// Stream assembler for handling chunked streaming messages
738#[derive(Debug)]
739pub struct StreamAssembler {
740    /// Accumulated chunks
741    chunks: Vec<StreamChunk>,
742    /// Accumulated payload data
743    data: Vec<u8>,
744    /// Expected next sequence number
745    next_sequence: u32,
746    /// Whether the stream is complete
747    complete: bool,
748    /// Whether an error occurred
749    error: bool,
750}
751
752impl StreamAssembler {
753    /// Create a new stream assembler
754    pub fn new() -> Self {
755        Self {
756            chunks: Vec::new(),
757            data: Vec::new(),
758            next_sequence: 0,
759            complete: false,
760            error: false,
761        }
762    }
763
764    /// Add a chunk to the stream
765    pub fn add_chunk(&mut self, chunk_bytes: &[u8], payload: &[u8]) -> Result<()> {
766        if self.complete {
767            return Err(DrivenError::InvalidBinary(
768                "Stream already complete".to_string(),
769            ));
770        }
771
772        let chunk = StreamChunk::from_bytes(chunk_bytes)
773            .map_err(|e| DrivenError::InvalidBinary(format!("Invalid chunk: {:?}", e)))?;
774
775        // Copy values from packed struct
776        let sequence = chunk.sequence;
777        let flags = chunk.flags;
778        let len = chunk.len;
779
780        // Validate sequence
781        if sequence != self.next_sequence {
782            return Err(DrivenError::InvalidBinary(format!(
783                "Out of order chunk: expected {}, got {}",
784                self.next_sequence, sequence
785            )));
786        }
787
788        // Check for error
789        if flags & ChunkFlags::ERROR != 0 {
790            self.error = true;
791            self.complete = true;
792            return Err(DrivenError::InvalidBinary("Stream error".to_string()));
793        }
794
795        // Validate payload length
796        if payload.len() != len as usize {
797            return Err(DrivenError::InvalidBinary(format!(
798                "Payload length mismatch: header says {}, got {}",
799                len,
800                payload.len()
801            )));
802        }
803
804        // Accumulate data
805        self.data.extend_from_slice(payload);
806        self.chunks.push(*chunk);
807        self.next_sequence += 1;
808
809        // Check if complete
810        if flags & ChunkFlags::LAST != 0 {
811            self.complete = true;
812        }
813
814        Ok(())
815    }
816
817    /// Check if the stream is complete
818    pub fn is_complete(&self) -> bool {
819        self.complete
820    }
821
822    /// Check if an error occurred
823    pub fn has_error(&self) -> bool {
824        self.error
825    }
826
827    /// Get the assembled data (only valid when complete)
828    pub fn data(&self) -> Option<&[u8]> {
829        if self.complete && !self.error {
830            Some(&self.data)
831        } else {
832            None
833        }
834    }
835
836    /// Take the assembled data (consumes the assembler)
837    pub fn take_data(self) -> Option<Vec<u8>> {
838        if self.complete && !self.error {
839            Some(self.data)
840        } else {
841            None
842        }
843    }
844
845    /// Get the number of chunks received
846    pub fn chunk_count(&self) -> usize {
847        self.chunks.len()
848    }
849}
850
851impl Default for StreamAssembler {
852    fn default() -> Self {
853        Self::new()
854    }
855}
856
857/// Stream builder for creating chunked streaming messages
858#[derive(Debug)]
859pub struct StreamBuilder {
860    /// Chunk size
861    chunk_size: usize,
862    /// Current sequence number
863    sequence: u32,
864}
865
866impl StreamBuilder {
867    /// Create a new stream builder with the given chunk size
868    pub fn new(chunk_size: usize) -> Self {
869        Self {
870            chunk_size,
871            sequence: 0,
872        }
873    }
874
875    /// Create a stream builder with default chunk size (64KB)
876    pub fn with_default_chunk_size() -> Self {
877        Self::new(65536)
878    }
879
880    /// Build chunks from data
881    pub fn build_chunks(&mut self, data: &[u8]) -> Vec<(StreamChunk, Vec<u8>)> {
882        let mut chunks = Vec::new();
883        let mut offset = 0;
884
885        while offset < data.len() {
886            let remaining = data.len() - offset;
887            let chunk_len = remaining.min(self.chunk_size);
888            let is_first = offset == 0;
889            let is_last = offset + chunk_len >= data.len();
890
891            let flags = if is_first && is_last {
892                ChunkFlags::FIRST | ChunkFlags::LAST
893            } else if is_first {
894                ChunkFlags::FIRST
895            } else if is_last {
896                ChunkFlags::LAST
897            } else {
898                ChunkFlags::CONTINUE
899            };
900
901            let chunk = StreamChunk::new(self.sequence, flags, chunk_len as u16);
902            let payload = data[offset..offset + chunk_len].to_vec();
903
904            chunks.push((chunk, payload));
905            self.sequence += 1;
906            offset += chunk_len;
907        }
908
909        chunks
910    }
911
912    /// Reset the sequence counter
913    pub fn reset(&mut self) {
914        self.sequence = 0;
915    }
916}
917
918impl DcpClient {
919    /// Create a new DCP client with the given configuration
920    pub fn new(config: DcpConfig) -> Self {
921        let signer = if config.signing_enabled {
922            config.signing_seed.map(|seed| Signer::from_seed(&seed))
923        } else {
924            None
925        };
926
927        Self {
928            config,
929            state: ConnectionState::Disconnected,
930            mcp_adapter: None,
931            capabilities: CapabilityManifest::new(1),
932            signer,
933            tools: HashMap::new(),
934            next_tool_id: 1,
935        }
936    }
937
938    /// Create a new DCP client with default configuration
939    pub fn with_defaults() -> Self {
940        Self::new(DcpConfig::default())
941    }
942
943    /// Get the current connection state
944    pub fn state(&self) -> ConnectionState {
945        self.state
946    }
947
948    /// Check if connected (either DCP or MCP)
949    pub fn is_connected(&self) -> bool {
950        matches!(
951            self.state,
952            ConnectionState::ConnectedDcp | ConnectionState::ConnectedMcp
953        )
954    }
955
956    /// Check if connected via DCP
957    pub fn is_dcp_connected(&self) -> bool {
958        self.state == ConnectionState::ConnectedDcp
959    }
960
961    /// Check if connected via MCP
962    pub fn is_mcp_connected(&self) -> bool {
963        self.state == ConnectionState::ConnectedMcp
964    }
965
966    /// Check if DCP should be preferred over MCP
967    pub fn prefer_dcp(&self) -> bool {
968        self.config.prefer_dcp && self.config.enabled
969    }
970
971    /// Set DCP preference
972    pub fn set_prefer_dcp(&mut self, prefer: bool) {
973        self.config.prefer_dcp = prefer;
974    }
975
976    /// Enable DCP protocol
977    pub fn enable_dcp(&mut self) {
978        self.config.enabled = true;
979    }
980
981    /// Disable DCP protocol (will use MCP only)
982    pub fn disable_dcp(&mut self) {
983        self.config.enabled = false;
984    }
985
986    /// Check if DCP is enabled
987    pub fn is_dcp_enabled(&self) -> bool {
988        self.config.enabled
989    }
990
991    /// Get the current protocol in use
992    pub fn current_protocol(&self) -> Option<Protocol> {
993        match self.state {
994            ConnectionState::ConnectedDcp => Some(Protocol::Dcp),
995            ConnectionState::ConnectedMcp => Some(Protocol::Mcp),
996            _ => None,
997        }
998    }
999
1000    /// Select the best protocol for a given operation
1001    ///
1002    /// Returns DCP if connected via DCP and DCP is preferred,
1003    /// otherwise returns MCP if available.
1004    pub fn select_protocol(&self) -> Option<Protocol> {
1005        if self.is_dcp_connected() && self.prefer_dcp() {
1006            Some(Protocol::Dcp)
1007        } else if self.is_mcp_connected() || self.has_mcp_fallback() {
1008            Some(Protocol::Mcp)
1009        } else if self.is_dcp_connected() {
1010            Some(Protocol::Dcp)
1011        } else {
1012            None
1013        }
1014    }
1015
1016    /// Invoke a tool using the best available protocol
1017    ///
1018    /// Automatically selects DCP or MCP based on connection state and preferences.
1019    pub fn invoke_tool_auto(&self, tool_id: u32, args: &[u8]) -> Result<InvocationResult> {
1020        match self.select_protocol() {
1021            Some(Protocol::Dcp) => {
1022                let encoded = self.invoke_tool(tool_id, args)?;
1023                Ok(InvocationResult::Dcp(encoded))
1024            }
1025            Some(Protocol::Mcp) => {
1026                // Convert to MCP format and invoke
1027                let tool = self
1028                    .tools
1029                    .get(&tool_id)
1030                    .ok_or_else(|| DrivenError::Config(format!("Tool {} not found", tool_id)))?;
1031
1032                let request = serde_json::json!({
1033                    "jsonrpc": "2.0",
1034                    "method": "tools/call",
1035                    "params": {
1036                        "name": tool.name,
1037                        "arguments": serde_json::from_slice::<serde_json::Value>(args)
1038                            .unwrap_or(serde_json::Value::Null)
1039                    },
1040                    "id": 1
1041                });
1042
1043                let response = self.handle_mcp_request(&request.to_string())?;
1044                Ok(InvocationResult::Mcp(response))
1045            }
1046            None => Err(DrivenError::Config("No protocol available".to_string())),
1047        }
1048    }
1049
1050    /// Get the capability manifest
1051    pub fn capabilities(&self) -> &CapabilityManifest {
1052        &self.capabilities
1053    }
1054
1055    /// Get a mutable reference to the capability manifest
1056    pub fn capabilities_mut(&mut self) -> &mut CapabilityManifest {
1057        &mut self.capabilities
1058    }
1059
1060    /// Connect to a DCP server
1061    ///
1062    /// If DCP connection fails and MCP fallback is available, will fall back to MCP.
1063    pub fn connect(&mut self, endpoint: &str) -> Result<()> {
1064        self.state = ConnectionState::Connecting;
1065
1066        // Try DCP connection first if enabled
1067        if self.config.enabled {
1068            match self.try_dcp_connect(endpoint) {
1069                Ok(()) => {
1070                    self.state = ConnectionState::ConnectedDcp;
1071                    return Ok(());
1072                }
1073                Err(e) => {
1074                    // Log the error but continue to MCP fallback
1075                    tracing::warn!("DCP connection failed, falling back to MCP: {}", e);
1076                }
1077            }
1078        }
1079
1080        // Fall back to MCP
1081        match self.try_mcp_connect(endpoint) {
1082            Ok(()) => {
1083                self.state = ConnectionState::ConnectedMcp;
1084                Ok(())
1085            }
1086            Err(e) => {
1087                self.state = ConnectionState::Failed;
1088                Err(e)
1089            }
1090        }
1091    }
1092
1093    /// Try to connect via DCP protocol
1094    fn try_dcp_connect(&mut self, _endpoint: &str) -> Result<()> {
1095        // TODO: Implement actual DCP connection
1096        // For now, this is a placeholder that will be implemented
1097        // when we have the full DCP server infrastructure
1098        Err(DrivenError::Config(
1099            "DCP connection not yet implemented".to_string(),
1100        ))
1101    }
1102
1103    /// Try to connect via MCP protocol
1104    fn try_mcp_connect(&mut self, _endpoint: &str) -> Result<()> {
1105        // Create MCP adapter for fallback
1106        self.mcp_adapter = Some(McpAdapter::new());
1107        Ok(())
1108    }
1109
1110    /// Disconnect from the server
1111    pub fn disconnect(&mut self) {
1112        self.state = ConnectionState::Disconnected;
1113        self.mcp_adapter = None;
1114    }
1115
1116    /// Register a tool
1117    pub fn register_tool(
1118        &mut self,
1119        name: &str,
1120        description: &str,
1121        schema_hash: [u8; 32],
1122        capabilities: u64,
1123    ) -> u32 {
1124        let id = self.next_tool_id;
1125        self.next_tool_id += 1;
1126
1127        let tool = ToolDefinition {
1128            id,
1129            name: name.to_string(),
1130            description: description.to_string(),
1131            schema_hash,
1132            capabilities,
1133        };
1134
1135        self.tools.insert(id, tool);
1136
1137        // Update capability manifest
1138        if id < CapabilityManifest::MAX_TOOLS as u32 {
1139            self.capabilities.set_tool(id as u16);
1140        }
1141
1142        id
1143    }
1144
1145    /// Get a registered tool by ID
1146    pub fn get_tool(&self, id: u32) -> Option<&ToolDefinition> {
1147        self.tools.get(&id)
1148    }
1149
1150    /// List all registered tools
1151    pub fn list_tools(&self) -> Vec<&ToolDefinition> {
1152        self.tools.values().collect()
1153    }
1154
1155    /// Invoke a tool via DCP binary protocol
1156    ///
1157    /// Uses Zero-Copy Tool Invocation (ZCTI) for maximum performance.
1158    pub fn invoke_tool(&self, tool_id: u32, args: &[u8]) -> Result<Vec<u8>> {
1159        if !self.is_connected() {
1160            return Err(DrivenError::Config("Not connected".to_string()));
1161        }
1162
1163        // Check if tool exists
1164        if !self.tools.contains_key(&tool_id) {
1165            return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1166        }
1167
1168        // Create tool invocation
1169        let invocation = ToolInvocation::new(tool_id, 0, 0, args.len() as u32);
1170
1171        // Build payload: invocation header + args
1172        let mut payload = Vec::with_capacity(ToolInvocation::SIZE + args.len());
1173        payload.extend_from_slice(&invocation.as_bytes());
1174        payload.extend_from_slice(args);
1175
1176        // Create BME message
1177        let flags = if self.signer.is_some() {
1178            Flags::SIGNED
1179        } else {
1180            0
1181        };
1182        let message = DcpMessage::new(MessageType::Tool, flags, payload);
1183
1184        // Encode the message
1185        let encoded = message.encode();
1186
1187        // TODO: Actually send the message when DCP transport is implemented
1188        // For now, return the encoded message as a placeholder
1189        Ok(encoded)
1190    }
1191
1192    /// Encode a DCP message with BME envelope
1193    pub fn encode_message(&self, message_type: MessageType, payload: &[u8]) -> Vec<u8> {
1194        let flags = if self.signer.is_some() {
1195            Flags::SIGNED
1196        } else {
1197            0
1198        };
1199        let message = DcpMessage::new(message_type, flags, payload.to_vec());
1200        message.encode()
1201    }
1202
1203    /// Decode a DCP message from bytes
1204    pub fn decode_message(&self, bytes: &[u8]) -> Result<DcpMessage> {
1205        DcpMessage::decode(bytes)
1206    }
1207
1208    /// Create a tool invocation message
1209    pub fn create_tool_message(&self, tool_id: u32, args: &[u8]) -> Result<DcpMessage> {
1210        if !self.tools.contains_key(&tool_id) {
1211            return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1212        }
1213
1214        let invocation = ToolInvocation::new(tool_id, 0, 0, args.len() as u32);
1215
1216        let mut payload = Vec::with_capacity(ToolInvocation::SIZE + args.len());
1217        payload.extend_from_slice(&invocation.as_bytes());
1218        payload.extend_from_slice(args);
1219
1220        let flags = if self.signer.is_some() {
1221            Flags::SIGNED
1222        } else {
1223            0
1224        };
1225        Ok(DcpMessage::new(MessageType::Tool, flags, payload))
1226    }
1227
1228    /// Create a resource request message
1229    pub fn create_resource_message(&self, resource_uri: &str) -> DcpMessage {
1230        let payload = resource_uri.as_bytes().to_vec();
1231        let flags = if self.signer.is_some() {
1232            Flags::SIGNED
1233        } else {
1234            0
1235        };
1236        DcpMessage::new(MessageType::Resource, flags, payload)
1237    }
1238
1239    /// Create a prompt message
1240    pub fn create_prompt_message(&self, prompt: &str) -> DcpMessage {
1241        let payload = prompt.as_bytes().to_vec();
1242        let flags = if self.signer.is_some() {
1243            Flags::SIGNED
1244        } else {
1245            0
1246        };
1247        DcpMessage::new(MessageType::Prompt, flags, payload)
1248    }
1249
1250    /// Create a response message
1251    pub fn create_response_message(&self, response: &[u8]) -> DcpMessage {
1252        let flags = if self.signer.is_some() {
1253            Flags::SIGNED
1254        } else {
1255            0
1256        };
1257        DcpMessage::new(MessageType::Response, flags, response.to_vec())
1258    }
1259
1260    /// Create an error message
1261    pub fn create_error_message(&self, error_code: u32, error_msg: &str) -> DcpMessage {
1262        let mut payload = Vec::with_capacity(4 + error_msg.len());
1263        payload.extend_from_slice(&error_code.to_le_bytes());
1264        payload.extend_from_slice(error_msg.as_bytes());
1265        DcpMessage::new(MessageType::Error, 0, payload)
1266    }
1267
1268    /// Create a streaming message builder
1269    pub fn create_stream_builder(&self, chunk_size: usize) -> StreamBuilder {
1270        StreamBuilder::new(chunk_size)
1271    }
1272
1273    /// Create a stream assembler for receiving chunked data
1274    pub fn create_stream_assembler(&self) -> StreamAssembler {
1275        StreamAssembler::new()
1276    }
1277
1278    /// Create a ZCTI builder for the given tool
1279    pub fn create_zcti_builder(&self, tool_id: u32) -> Result<ZctiBuilder> {
1280        if !self.tools.contains_key(&tool_id) {
1281            return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1282        }
1283        Ok(ZctiBuilder::new(tool_id))
1284    }
1285
1286    /// Invoke a tool using ZCTI (Zero-Copy Tool Invocation)
1287    pub fn invoke_tool_zcti(&self, builder: ZctiBuilder) -> Result<Vec<u8>> {
1288        if !self.is_connected() {
1289            return Err(DrivenError::Config("Not connected".to_string()));
1290        }
1291
1292        let tool_id = builder.tool_id;
1293        if !self.tools.contains_key(&tool_id) {
1294            return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1295        }
1296
1297        let (invocation, args_data) = builder.build_with_data();
1298
1299        // Build payload: invocation header + args
1300        let mut payload = Vec::with_capacity(ToolInvocation::SIZE + args_data.len());
1301        payload.extend_from_slice(&invocation.as_bytes());
1302        payload.extend_from_slice(&args_data);
1303
1304        // Create BME message
1305        let flags = if self.signer.is_some() {
1306            Flags::SIGNED
1307        } else {
1308            0
1309        };
1310        let message = DcpMessage::new(MessageType::Tool, flags, payload);
1311
1312        // Encode and return
1313        Ok(message.encode())
1314    }
1315
1316    /// Parse a tool invocation from a received message
1317    pub fn parse_tool_invocation<'a>(&self, message: &'a DcpMessage) -> Result<ZctiReader<'a>> {
1318        if message.message_type != MessageType::Tool {
1319            return Err(DrivenError::InvalidBinary(
1320                "Message is not a tool invocation".to_string(),
1321            ));
1322        }
1323
1324        if message.payload.len() < ToolInvocation::SIZE {
1325            return Err(DrivenError::InvalidBinary(
1326                "Payload too small for tool invocation".to_string(),
1327            ));
1328        }
1329
1330        let invocation = ToolInvocation::from_bytes(&message.payload)
1331            .map_err(|e| DrivenError::InvalidBinary(format!("Invalid invocation: {:?}", e)))?;
1332
1333        let args_data = &message.payload[ToolInvocation::SIZE..];
1334        Ok(ZctiReader::new(&invocation, args_data))
1335    }
1336
1337    /// Create a shared argument buffer
1338    pub fn create_shared_buffer(&self, capacity: usize) -> SharedArgBuffer {
1339        SharedArgBuffer::new(capacity)
1340    }
1341
1342    /// Get the MCP adapter for JSON-RPC compatibility
1343    pub fn mcp_adapter(&self) -> Option<&McpAdapter> {
1344        self.mcp_adapter.as_ref()
1345    }
1346
1347    /// Get a mutable reference to the MCP adapter
1348    pub fn mcp_adapter_mut(&mut self) -> Option<&mut McpAdapter> {
1349        self.mcp_adapter.as_mut()
1350    }
1351
1352    /// Register a tool with the MCP adapter for JSON-RPC compatibility
1353    pub fn register_tool_mcp(&mut self, name: &str, tool_id: u32) {
1354        if let Some(ref mut adapter) = self.mcp_adapter {
1355            adapter.register_tool(name, tool_id as u16);
1356        }
1357    }
1358
1359    /// Handle an MCP JSON-RPC request and return a JSON-RPC response
1360    ///
1361    /// This provides seamless fallback to MCP when DCP is unavailable.
1362    pub fn handle_mcp_request(&self, json: &str) -> Result<String> {
1363        let adapter = self
1364            .mcp_adapter
1365            .as_ref()
1366            .ok_or_else(|| DrivenError::Config("MCP adapter not available".to_string()))?;
1367
1368        // Parse the request
1369        let request = adapter
1370            .parse_request(json)
1371            .map_err(|e| DrivenError::Parse(format!("Invalid JSON-RPC request: {:?}", e)))?;
1372
1373        // Handle based on method
1374        match request.method.as_str() {
1375            "initialize" => adapter
1376                .handle_initialize(&request)
1377                .map_err(|e| DrivenError::Format(format!("Failed to handle initialize: {:?}", e))),
1378            "tools/list" => adapter
1379                .handle_tools_list(&request)
1380                .map_err(|e| DrivenError::Format(format!("Failed to handle tools/list: {:?}", e))),
1381            "tools/call" => {
1382                // For tools/call, we need to handle it ourselves since we don't have a router
1383                self.handle_mcp_tools_call(&request)
1384            }
1385            "resources/list" => {
1386                // Return empty resources list
1387                let result = serde_json::json!({ "resources": [] });
1388                adapter
1389                    .format_success_response(request.id, result)
1390                    .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))
1391            }
1392            "prompts/list" => {
1393                // Return empty prompts list
1394                let result = serde_json::json!({ "prompts": [] });
1395                adapter
1396                    .format_success_response(request.id, result)
1397                    .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))
1398            }
1399            _ => {
1400                // Unknown method
1401                adapter
1402                    .format_error_response(request.id, JsonRpcError::method_not_found())
1403                    .map_err(|e| DrivenError::Format(format!("Failed to format error: {:?}", e)))
1404            }
1405        }
1406    }
1407
1408    /// Handle an MCP tools/call request
1409    fn handle_mcp_tools_call(&self, request: &JsonRpcRequest) -> Result<String> {
1410        let adapter = self
1411            .mcp_adapter
1412            .as_ref()
1413            .ok_or_else(|| DrivenError::Config("MCP adapter not available".to_string()))?;
1414
1415        // Extract tool name and arguments from params
1416        let params = request
1417            .params
1418            .as_ref()
1419            .ok_or_else(|| DrivenError::Parse("Missing params".to_string()))?;
1420
1421        let tool_name = params
1422            .get("name")
1423            .and_then(|v| v.as_str())
1424            .ok_or_else(|| DrivenError::Parse("Missing tool name".to_string()))?;
1425
1426        let arguments = params.get("arguments").cloned();
1427
1428        // Resolve tool name to ID
1429        let tool_id = adapter
1430            .resolve_tool_name(tool_name)
1431            .ok_or_else(|| DrivenError::Config(format!("Unknown tool: {}", tool_name)))?;
1432
1433        // Check if tool exists in our registry
1434        if !self.tools.contains_key(&(tool_id as u32)) {
1435            return adapter
1436                .format_error_response(
1437                    request.id.clone(),
1438                    JsonRpcError::new(-32602, format!("Tool not found: {}", tool_name)),
1439                )
1440                .map_err(|e| DrivenError::Format(format!("Failed to format error: {:?}", e)));
1441        }
1442
1443        // Translate params to bytes
1444        let args_bytes = adapter.translate_params(&arguments);
1445
1446        // For now, return a placeholder response since we don't have actual tool execution
1447        // In a full implementation, this would execute the tool and return the result
1448        let response_result = serde_json::json!({
1449            "content": [{
1450                "type": "text",
1451                "text": format!("Tool {} called with {} bytes of arguments", tool_name, args_bytes.len())
1452            }]
1453        });
1454
1455        adapter
1456            .format_success_response(request.id.clone(), response_result)
1457            .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))
1458    }
1459
1460    /// Check if MCP fallback is available
1461    pub fn has_mcp_fallback(&self) -> bool {
1462        self.mcp_adapter.is_some()
1463    }
1464
1465    /// Convert a DCP message to MCP JSON-RPC format
1466    pub fn dcp_to_mcp(&self, message: &DcpMessage) -> Result<String> {
1467        match message.message_type {
1468            MessageType::Response => {
1469                // Convert response payload to JSON
1470                let result: serde_json::Value = serde_json::from_slice(&message.payload)
1471                    .unwrap_or_else(|_| {
1472                        serde_json::Value::String(
1473                            String::from_utf8_lossy(&message.payload).to_string(),
1474                        )
1475                    });
1476
1477                let response = serde_json::json!({
1478                    "jsonrpc": "2.0",
1479                    "result": result,
1480                    "id": null
1481                });
1482
1483                serde_json::to_string(&response)
1484                    .map_err(|e| DrivenError::Format(format!("Failed to serialize: {}", e)))
1485            }
1486            MessageType::Error => {
1487                // Parse error code and message from payload
1488                let (code, msg) = if message.payload.len() >= 4 {
1489                    let code_bytes: [u8; 4] = message.payload[0..4].try_into().unwrap();
1490                    let code = i32::from_le_bytes(code_bytes);
1491                    let msg = String::from_utf8_lossy(&message.payload[4..]).to_string();
1492                    (code, msg)
1493                } else {
1494                    (-32000, "Unknown error".to_string())
1495                };
1496
1497                let response = serde_json::json!({
1498                    "jsonrpc": "2.0",
1499                    "error": {
1500                        "code": code,
1501                        "message": msg
1502                    },
1503                    "id": null
1504                });
1505
1506                serde_json::to_string(&response)
1507                    .map_err(|e| DrivenError::Format(format!("Failed to serialize: {}", e)))
1508            }
1509            _ => Err(DrivenError::Format(format!(
1510                "Cannot convert {:?} message to MCP format",
1511                message.message_type
1512            ))),
1513        }
1514    }
1515
1516    /// Convert an MCP JSON-RPC request to DCP message
1517    pub fn mcp_to_dcp(&self, json: &str) -> Result<DcpMessage> {
1518        let adapter = self
1519            .mcp_adapter
1520            .as_ref()
1521            .ok_or_else(|| DrivenError::Config("MCP adapter not available".to_string()))?;
1522
1523        let request = adapter
1524            .parse_request(json)
1525            .map_err(|e| DrivenError::Parse(format!("Invalid JSON-RPC: {:?}", e)))?;
1526
1527        // Determine message type based on method
1528        let message_type = match request.method.as_str() {
1529            "tools/call" => MessageType::Tool,
1530            "resources/read" => MessageType::Resource,
1531            "prompts/get" => MessageType::Prompt,
1532            _ => MessageType::Tool, // Default to tool
1533        };
1534
1535        // Serialize the request as payload
1536        let payload = serde_json::to_vec(&request)
1537            .map_err(|e| DrivenError::Format(format!("Failed to serialize: {}", e)))?;
1538
1539        Ok(DcpMessage::new(message_type, 0, payload))
1540    }
1541
1542    /// Invoke a tool via MCP JSON-RPC protocol
1543    ///
1544    /// Used as fallback when DCP is unavailable.
1545    pub fn invoke_tool_mcp(&self, request: &str) -> Result<String> {
1546        if !self.is_connected() {
1547            return Err(DrivenError::Config("Not connected".to_string()));
1548        }
1549
1550        // Parse the JSON-RPC request
1551        let _parsed = JsonRpcParser::parse_request(request)
1552            .map_err(|e| DrivenError::Parse(format!("Invalid JSON-RPC request: {:?}", e)))?;
1553
1554        // Use MCP adapter if available
1555        if let Some(ref adapter) = self.mcp_adapter {
1556            // For now, return a method not found error since we don't have a router
1557            // In a full implementation, this would route to the appropriate handler
1558            let response = adapter
1559                .format_error_response(RequestId::Null, JsonRpcError::method_not_found())
1560                .map_err(|e| DrivenError::Format(format!("Failed to format response: {:?}", e)))?;
1561            Ok(response)
1562        } else {
1563            Err(DrivenError::Config("MCP adapter not available".to_string()))
1564        }
1565    }
1566
1567    /// Create a signed tool definition
1568    pub fn sign_tool_def(&self, tool_id: u32) -> Result<dcp::binary::SignedToolDef> {
1569        let signer = self
1570            .signer
1571            .as_ref()
1572            .ok_or_else(|| DrivenError::Security("Signing not enabled".to_string()))?;
1573
1574        let tool = self
1575            .tools
1576            .get(&tool_id)
1577            .ok_or_else(|| DrivenError::Config(format!("Tool {} not found", tool_id)))?;
1578
1579        Ok(signer.sign_tool_def(tool_id, tool.schema_hash, tool.capabilities))
1580    }
1581
1582    /// Verify a signed tool definition
1583    pub fn verify_tool_def(&self, def: &dcp::binary::SignedToolDef) -> Result<()> {
1584        Verifier::verify_tool_def(def)
1585            .map_err(|e| DrivenError::Security(format!("Signature verification failed: {:?}", e)))
1586    }
1587
1588    /// Sign a tool invocation
1589    ///
1590    /// Creates a signed invocation that can be verified by the server.
1591    /// Requires signing to be enabled.
1592    pub fn sign_invocation(
1593        &self,
1594        tool_id: u32,
1595        nonce: u64,
1596        timestamp: u64,
1597        args: &[u8],
1598    ) -> Result<SignedInvocation> {
1599        let signer = self
1600            .signer
1601            .as_ref()
1602            .ok_or_else(|| DrivenError::Security("Signing not enabled".to_string()))?;
1603
1604        if !self.tools.contains_key(&tool_id) {
1605            return Err(DrivenError::Config(format!("Tool {} not found", tool_id)));
1606        }
1607
1608        Ok(signer.sign_invocation(tool_id, nonce, timestamp, args))
1609    }
1610
1611    /// Verify a signed invocation
1612    ///
1613    /// Verifies that the invocation was signed by the holder of the given public key.
1614    pub fn verify_invocation(&self, inv: &SignedInvocation, public_key: &[u8; 32]) -> Result<()> {
1615        Verifier::verify_invocation(inv, public_key)
1616            .map_err(|e| DrivenError::Security(format!("Invocation verification failed: {:?}", e)))
1617    }
1618
1619    /// Verify that the args hash in a signed invocation matches the provided arguments
1620    pub fn verify_args_hash(&self, inv: &SignedInvocation, args: &[u8]) -> bool {
1621        Verifier::verify_args_hash(inv, args)
1622    }
1623
1624    /// Generate a random signer (for testing or ephemeral keys)
1625    pub fn generate_signer() -> Signer {
1626        Signer::generate()
1627    }
1628
1629    // ==================== Capability Manifest Support ====================
1630
1631    /// Parse a capability manifest from bytes
1632    ///
1633    /// This is used when receiving a capability manifest from a remote server
1634    /// during capability negotiation.
1635    pub fn parse_capability_manifest(bytes: &[u8]) -> Result<&CapabilityManifest> {
1636        CapabilityManifest::from_bytes(bytes).map_err(|e| {
1637            DrivenError::InvalidBinary(format!("Failed to parse capability manifest: {:?}", e))
1638        })
1639    }
1640
1641    /// Serialize the current capability manifest to bytes
1642    pub fn serialize_capabilities(&self) -> Vec<u8> {
1643        self.capabilities.as_bytes().to_vec()
1644    }
1645
1646    /// Register a resource in the capability manifest
1647    pub fn register_resource(&mut self, resource_id: u16) {
1648        self.capabilities.set_resource(resource_id);
1649    }
1650
1651    /// Unregister a resource from the capability manifest
1652    pub fn unregister_resource(&mut self, resource_id: u16) {
1653        self.capabilities.clear_resource(resource_id);
1654    }
1655
1656    /// Check if a resource is registered
1657    pub fn has_resource(&self, resource_id: u16) -> bool {
1658        self.capabilities.has_resource(resource_id)
1659    }
1660
1661    /// Register a prompt in the capability manifest
1662    pub fn register_prompt(&mut self, prompt_id: u16) {
1663        self.capabilities.set_prompt(prompt_id);
1664    }
1665
1666    /// Unregister a prompt from the capability manifest
1667    pub fn unregister_prompt(&mut self, prompt_id: u16) {
1668        self.capabilities.clear_prompt(prompt_id);
1669    }
1670
1671    /// Check if a prompt is registered
1672    pub fn has_prompt(&self, prompt_id: u16) -> bool {
1673        self.capabilities.has_prompt(prompt_id)
1674    }
1675
1676    /// Set an extension flag in the capability manifest
1677    pub fn set_extension(&mut self, bit: u8) {
1678        self.capabilities.set_extension(bit);
1679    }
1680
1681    /// Clear an extension flag from the capability manifest
1682    pub fn clear_extension(&mut self, bit: u8) {
1683        self.capabilities.clear_extension(bit);
1684    }
1685
1686    /// Check if an extension is enabled
1687    pub fn has_extension(&self, bit: u8) -> bool {
1688        self.capabilities.has_extension(bit)
1689    }
1690
1691    /// Enforce that a tool is available before invocation
1692    ///
1693    /// Returns an error if the tool is not registered in the capability manifest.
1694    pub fn enforce_tool(&self, tool_id: u32) -> Result<()> {
1695        if tool_id >= CapabilityManifest::MAX_TOOLS as u32 {
1696            return Err(DrivenError::Config(format!(
1697                "Tool ID {} exceeds maximum ({})",
1698                tool_id,
1699                CapabilityManifest::MAX_TOOLS
1700            )));
1701        }
1702        if !self.capabilities.has_tool(tool_id as u16) {
1703            return Err(DrivenError::Config(format!(
1704                "Tool {} is not available in capability manifest",
1705                tool_id
1706            )));
1707        }
1708        Ok(())
1709    }
1710
1711    /// Enforce that a resource is available before access
1712    ///
1713    /// Returns an error if the resource is not registered in the capability manifest.
1714    pub fn enforce_resource(&self, resource_id: u16) -> Result<()> {
1715        if !self.capabilities.has_resource(resource_id) {
1716            return Err(DrivenError::Config(format!(
1717                "Resource {} is not available in capability manifest",
1718                resource_id
1719            )));
1720        }
1721        Ok(())
1722    }
1723
1724    /// Enforce that a prompt is available before use
1725    ///
1726    /// Returns an error if the prompt is not registered in the capability manifest.
1727    pub fn enforce_prompt(&self, prompt_id: u16) -> Result<()> {
1728        if !self.capabilities.has_prompt(prompt_id) {
1729            return Err(DrivenError::Config(format!(
1730                "Prompt {} is not available in capability manifest",
1731                prompt_id
1732            )));
1733        }
1734        Ok(())
1735    }
1736
1737    /// Negotiate capabilities with a remote server
1738    ///
1739    /// Takes the server's capability manifest and computes the intersection,
1740    /// updating the local capabilities to only include mutually supported features.
1741    pub fn negotiate_capabilities(
1742        &mut self,
1743        server_manifest: &CapabilityManifest,
1744    ) -> CapabilityManifest {
1745        let negotiated = self.capabilities.intersect(server_manifest);
1746        self.capabilities = negotiated.clone();
1747        negotiated
1748    }
1749
1750    /// Get statistics about the current capabilities
1751    pub fn capability_stats(&self) -> CapabilityStats {
1752        CapabilityStats {
1753            tool_count: self.capabilities.tool_count(),
1754            resource_count: self.capabilities.resource_count(),
1755            prompt_count: self.capabilities.prompt_count(),
1756            extension_count: self.capabilities.extension_count(),
1757            version: self.capabilities.version,
1758        }
1759    }
1760
1761    /// Get an iterator over all registered tool IDs
1762    pub fn tool_ids(&self) -> impl Iterator<Item = u16> + '_ {
1763        self.capabilities.tool_ids()
1764    }
1765
1766    /// Get an iterator over all registered resource IDs
1767    pub fn resource_ids(&self) -> impl Iterator<Item = u16> + '_ {
1768        self.capabilities.resource_ids()
1769    }
1770
1771    /// Get an iterator over all registered prompt IDs
1772    pub fn prompt_ids(&self) -> impl Iterator<Item = u16> + '_ {
1773        self.capabilities.prompt_ids()
1774    }
1775
1776    /// Compute capability intersection with another manifest
1777    pub fn intersect_capabilities(&self, other: &CapabilityManifest) -> CapabilityManifest {
1778        self.capabilities.intersect(other)
1779    }
1780
1781    /// Enable signing with a seed
1782    pub fn enable_signing(&mut self, seed: [u8; 32]) {
1783        self.signer = Some(Signer::from_seed(&seed));
1784        self.config.signing_enabled = true;
1785        self.config.signing_seed = Some(seed);
1786    }
1787
1788    /// Disable signing
1789    pub fn disable_signing(&mut self) {
1790        self.signer = None;
1791        self.config.signing_enabled = false;
1792        self.config.signing_seed = None;
1793    }
1794
1795    /// Get the public key if signing is enabled
1796    pub fn public_key(&self) -> Option<[u8; 32]> {
1797        self.signer.as_ref().map(|s| s.public_key_bytes())
1798    }
1799}
1800
1801/// Create a Binary Message Envelope
1802pub fn create_envelope(
1803    message_type: MessageType,
1804    flags: u8,
1805    payload_len: u32,
1806) -> BinaryMessageEnvelope {
1807    BinaryMessageEnvelope::new(message_type, flags, payload_len)
1808}
1809
1810/// Parse a Binary Message Envelope from bytes
1811pub fn parse_envelope(bytes: &[u8]) -> Result<&BinaryMessageEnvelope> {
1812    BinaryMessageEnvelope::from_bytes(bytes)
1813        .map_err(|e| DrivenError::InvalidBinary(format!("Failed to parse BME: {:?}", e)))
1814}
1815
1816#[cfg(test)]
1817mod tests {
1818    use super::*;
1819    use proptest::prelude::*;
1820
1821    #[test]
1822    fn test_dcp_client_creation() {
1823        let client = DcpClient::with_defaults();
1824        assert_eq!(client.state(), ConnectionState::Disconnected);
1825        assert!(!client.is_connected());
1826        assert!(client.prefer_dcp());
1827    }
1828
1829    #[test]
1830    fn test_dcp_client_config() {
1831        let config = DcpConfig {
1832            enabled: false,
1833            prefer_dcp: false,
1834            endpoint: Some("tcp://localhost:9000".to_string()),
1835            timeout_ms: 10000,
1836            signing_enabled: true,
1837            signing_seed: Some([42u8; 32]),
1838        };
1839
1840        let client = DcpClient::new(config);
1841        assert!(!client.prefer_dcp());
1842        assert!(client.signer.is_some());
1843    }
1844
1845    #[test]
1846    fn test_tool_registration() {
1847        let mut client = DcpClient::with_defaults();
1848
1849        let tool_id = client.register_tool("test_tool", "A test tool", [0xAB; 32], 0x1234);
1850
1851        assert_eq!(tool_id, 1);
1852        assert!(client.get_tool(tool_id).is_some());
1853        assert!(client.capabilities().has_tool(tool_id as u16));
1854
1855        let tool = client.get_tool(tool_id).unwrap();
1856        assert_eq!(tool.name, "test_tool");
1857        assert_eq!(tool.description, "A test tool");
1858    }
1859
1860    #[test]
1861    fn test_signing() {
1862        let mut client = DcpClient::with_defaults();
1863        client.enable_signing([42u8; 32]);
1864
1865        let tool_id = client.register_tool("signed_tool", "A signed tool", [0xCD; 32], 0x5678);
1866
1867        let signed_def = client.sign_tool_def(tool_id).unwrap();
1868        assert!(client.verify_tool_def(&signed_def).is_ok());
1869    }
1870
1871    #[test]
1872    fn test_sign_invocation() {
1873        let mut client = DcpClient::with_defaults();
1874        client.enable_signing([42u8; 32]);
1875
1876        let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
1877        let args = b"test arguments";
1878
1879        let signed_inv = client
1880            .sign_invocation(tool_id, 12345, 1234567890, args)
1881            .unwrap();
1882
1883        let public_key = client.public_key().unwrap();
1884        assert!(client.verify_invocation(&signed_inv, &public_key).is_ok());
1885        assert!(client.verify_args_hash(&signed_inv, args));
1886        assert!(!client.verify_args_hash(&signed_inv, b"wrong args"));
1887    }
1888
1889    #[test]
1890    fn test_sign_invocation_without_signing_enabled() {
1891        let mut client = DcpClient::with_defaults();
1892        let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
1893
1894        // Should fail because signing is not enabled
1895        let result = client.sign_invocation(tool_id, 12345, 1234567890, b"args");
1896        assert!(result.is_err());
1897    }
1898
1899    #[test]
1900    fn test_sign_invocation_unknown_tool() {
1901        let mut client = DcpClient::with_defaults();
1902        client.enable_signing([42u8; 32]);
1903
1904        // Should fail because tool doesn't exist
1905        let result = client.sign_invocation(999, 12345, 1234567890, b"args");
1906        assert!(result.is_err());
1907    }
1908
1909    #[test]
1910    fn test_generate_signer() {
1911        let signer1 = DcpClient::generate_signer();
1912        let signer2 = DcpClient::generate_signer();
1913
1914        // Different signers should have different public keys
1915        assert_ne!(signer1.public_key_bytes(), signer2.public_key_bytes());
1916    }
1917
1918    #[test]
1919    fn test_capability_intersection() {
1920        let mut client = DcpClient::with_defaults();
1921        client.register_tool("tool1", "Tool 1", [0; 32], 0);
1922        client.register_tool("tool2", "Tool 2", [0; 32], 0);
1923        client.register_tool("tool3", "Tool 3", [0; 32], 0);
1924
1925        let mut other = CapabilityManifest::new(1);
1926        other.set_tool(2);
1927        other.set_tool(3);
1928        other.set_tool(4);
1929
1930        let intersection = client.intersect_capabilities(&other);
1931        assert!(!intersection.has_tool(1));
1932        assert!(intersection.has_tool(2));
1933        assert!(intersection.has_tool(3));
1934        assert!(!intersection.has_tool(4));
1935    }
1936
1937    #[test]
1938    fn test_envelope_creation() {
1939        let envelope = create_envelope(MessageType::Tool, Flags::STREAMING, 1024);
1940        assert!(envelope.is_streaming());
1941        assert!(!envelope.is_compressed());
1942        assert!(!envelope.is_signed());
1943
1944        let bytes = envelope.as_bytes();
1945        let parsed = parse_envelope(bytes).unwrap();
1946        assert_eq!(parsed.get_message_type(), Some(MessageType::Tool));
1947    }
1948
1949    #[test]
1950    fn test_mcp_fallback() {
1951        let mut client = DcpClient::with_defaults();
1952
1953        // Connect should fall back to MCP since DCP is not implemented
1954        let result = client.connect("tcp://localhost:9000");
1955        assert!(result.is_ok());
1956        assert_eq!(client.state(), ConnectionState::ConnectedMcp);
1957    }
1958
1959    #[test]
1960    fn test_protocol_preference() {
1961        let mut client = DcpClient::with_defaults();
1962
1963        // Default should prefer DCP
1964        assert!(client.prefer_dcp());
1965        assert!(client.is_dcp_enabled());
1966
1967        // Disable DCP preference
1968        client.set_prefer_dcp(false);
1969        assert!(!client.prefer_dcp());
1970
1971        // Re-enable
1972        client.set_prefer_dcp(true);
1973        assert!(client.prefer_dcp());
1974    }
1975
1976    #[test]
1977    fn test_enable_disable_dcp() {
1978        let mut client = DcpClient::with_defaults();
1979
1980        assert!(client.is_dcp_enabled());
1981
1982        client.disable_dcp();
1983        assert!(!client.is_dcp_enabled());
1984        assert!(!client.prefer_dcp()); // prefer_dcp requires enabled
1985
1986        client.enable_dcp();
1987        assert!(client.is_dcp_enabled());
1988    }
1989
1990    #[test]
1991    fn test_current_protocol() {
1992        let mut client = DcpClient::with_defaults();
1993
1994        // Not connected yet
1995        assert!(client.current_protocol().is_none());
1996
1997        // Connect (falls back to MCP)
1998        client.connect("tcp://localhost:9000").unwrap();
1999        assert_eq!(client.current_protocol(), Some(Protocol::Mcp));
2000    }
2001
2002    #[test]
2003    fn test_select_protocol() {
2004        let mut client = DcpClient::with_defaults();
2005
2006        // Not connected - no protocol available
2007        assert!(client.select_protocol().is_none());
2008
2009        // Connect (falls back to MCP)
2010        client.connect("tcp://localhost:9000").unwrap();
2011
2012        // Should select MCP since that's what we're connected with
2013        assert_eq!(client.select_protocol(), Some(Protocol::Mcp));
2014    }
2015
2016    #[test]
2017    fn test_is_mcp_connected() {
2018        let mut client = DcpClient::with_defaults();
2019
2020        assert!(!client.is_mcp_connected());
2021
2022        client.connect("tcp://localhost:9000").unwrap();
2023        assert!(client.is_mcp_connected());
2024        assert!(!client.is_dcp_connected());
2025    }
2026
2027    #[test]
2028    fn test_invoke_tool_auto() {
2029        let mut client = DcpClient::with_defaults();
2030        client.connect("tcp://localhost:9000").unwrap();
2031
2032        let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
2033        client.register_tool_mcp("test_tool", tool_id);
2034
2035        let args = b"{}";
2036        let result = client.invoke_tool_auto(tool_id, args);
2037
2038        // Should succeed and return MCP result (since we're connected via MCP)
2039        assert!(result.is_ok());
2040        let result = result.unwrap();
2041        assert!(result.is_mcp());
2042    }
2043
2044    #[test]
2045    fn test_invocation_result() {
2046        let dcp_result = InvocationResult::Dcp(vec![1, 2, 3]);
2047        assert!(dcp_result.is_dcp());
2048        assert!(!dcp_result.is_mcp());
2049        assert_eq!(dcp_result.as_dcp(), Some(&[1, 2, 3][..]));
2050        assert!(dcp_result.as_mcp().is_none());
2051
2052        let mcp_result = InvocationResult::Mcp("{}".to_string());
2053        assert!(!mcp_result.is_dcp());
2054        assert!(mcp_result.is_mcp());
2055        assert!(mcp_result.as_dcp().is_none());
2056        assert_eq!(mcp_result.as_mcp(), Some("{}"));
2057    }
2058
2059    // BME Message Tests
2060
2061    #[test]
2062    fn test_dcp_message_tool() {
2063        let payload = vec![1, 2, 3, 4, 5];
2064        let message = DcpMessage::tool(payload.clone());
2065
2066        assert_eq!(message.message_type, MessageType::Tool);
2067        assert_eq!(message.flags, 0);
2068        assert_eq!(message.payload, payload);
2069        assert!(!message.is_streaming());
2070        assert!(!message.is_compressed());
2071        assert!(!message.is_signed());
2072    }
2073
2074    #[test]
2075    fn test_dcp_message_resource() {
2076        let payload = b"resource://test".to_vec();
2077        let message = DcpMessage::resource(payload.clone());
2078
2079        assert_eq!(message.message_type, MessageType::Resource);
2080        assert_eq!(message.payload, payload);
2081    }
2082
2083    #[test]
2084    fn test_dcp_message_prompt() {
2085        let payload = b"Hello, AI!".to_vec();
2086        let message = DcpMessage::prompt(payload.clone());
2087
2088        assert_eq!(message.message_type, MessageType::Prompt);
2089        assert_eq!(message.payload, payload);
2090    }
2091
2092    #[test]
2093    fn test_dcp_message_response() {
2094        let payload = b"Response data".to_vec();
2095        let message = DcpMessage::response(payload.clone());
2096
2097        assert_eq!(message.message_type, MessageType::Response);
2098        assert_eq!(message.payload, payload);
2099    }
2100
2101    #[test]
2102    fn test_dcp_message_error() {
2103        let payload = b"Error occurred".to_vec();
2104        let message = DcpMessage::error(payload.clone());
2105
2106        assert_eq!(message.message_type, MessageType::Error);
2107        assert_eq!(message.payload, payload);
2108    }
2109
2110    #[test]
2111    fn test_dcp_message_stream() {
2112        let payload = b"Streaming data".to_vec();
2113        let message = DcpMessage::stream(payload.clone());
2114
2115        assert_eq!(message.message_type, MessageType::Stream);
2116        assert!(message.is_streaming());
2117        assert_eq!(message.payload, payload);
2118    }
2119
2120    #[test]
2121    fn test_dcp_message_flags() {
2122        let message = DcpMessage::tool(vec![])
2123            .with_streaming()
2124            .with_compressed()
2125            .with_signed();
2126
2127        assert!(message.is_streaming());
2128        assert!(message.is_compressed());
2129        assert!(message.is_signed());
2130    }
2131
2132    #[test]
2133    fn test_dcp_message_encode_decode_roundtrip() {
2134        let original = DcpMessage::new(
2135            MessageType::Tool,
2136            Flags::STREAMING | Flags::SIGNED,
2137            vec![0xDE, 0xAD, 0xBE, 0xEF],
2138        );
2139
2140        let encoded = original.encode();
2141        let decoded = DcpMessage::decode(&encoded).unwrap();
2142
2143        assert_eq!(decoded.message_type, original.message_type);
2144        assert_eq!(decoded.flags, original.flags);
2145        assert_eq!(decoded.payload, original.payload);
2146    }
2147
2148    #[test]
2149    fn test_dcp_message_decode_insufficient_data() {
2150        let bytes = vec![0u8; 4]; // Less than BME header size
2151        let result = DcpMessage::decode(&bytes);
2152        assert!(result.is_err());
2153    }
2154
2155    #[test]
2156    fn test_dcp_message_decode_invalid_magic() {
2157        let mut bytes = vec![0u8; 16];
2158        bytes[0] = 0xFF;
2159        bytes[1] = 0xFF;
2160        let result = DcpMessage::decode(&bytes);
2161        assert!(result.is_err());
2162    }
2163
2164    #[test]
2165    fn test_dcp_message_all_types_roundtrip() {
2166        let test_cases = vec![
2167            DcpMessage::tool(vec![1, 2, 3]),
2168            DcpMessage::resource(b"uri://test".to_vec()),
2169            DcpMessage::prompt(b"test prompt".to_vec()),
2170            DcpMessage::response(b"test response".to_vec()),
2171            DcpMessage::error(b"test error".to_vec()),
2172            DcpMessage::stream(b"streaming".to_vec()),
2173        ];
2174
2175        for original in test_cases {
2176            let encoded = original.encode();
2177            let decoded = DcpMessage::decode(&encoded).unwrap();
2178            assert_eq!(decoded.message_type, original.message_type);
2179            assert_eq!(decoded.payload, original.payload);
2180        }
2181    }
2182
2183    // Stream Assembler Tests
2184
2185    #[test]
2186    fn test_stream_assembler_single_chunk() {
2187        let mut assembler = StreamAssembler::new();
2188
2189        let chunk = StreamChunk::new(0, ChunkFlags::FIRST | ChunkFlags::LAST, 5);
2190        let payload = vec![1, 2, 3, 4, 5];
2191
2192        assembler.add_chunk(chunk.as_bytes(), &payload).unwrap();
2193
2194        assert!(assembler.is_complete());
2195        assert!(!assembler.has_error());
2196        assert_eq!(assembler.data(), Some(payload.as_slice()));
2197        assert_eq!(assembler.chunk_count(), 1);
2198    }
2199
2200    #[test]
2201    fn test_stream_assembler_multiple_chunks() {
2202        let mut assembler = StreamAssembler::new();
2203
2204        // First chunk
2205        let chunk1 = StreamChunk::first(0, 3);
2206        assembler.add_chunk(chunk1.as_bytes(), &[1, 2, 3]).unwrap();
2207        assert!(!assembler.is_complete());
2208
2209        // Continuation chunk
2210        let chunk2 = StreamChunk::continuation(1, 3);
2211        assembler.add_chunk(chunk2.as_bytes(), &[4, 5, 6]).unwrap();
2212        assert!(!assembler.is_complete());
2213
2214        // Last chunk
2215        let chunk3 = StreamChunk::last(2, 2);
2216        assembler.add_chunk(chunk3.as_bytes(), &[7, 8]).unwrap();
2217        assert!(assembler.is_complete());
2218
2219        assert_eq!(assembler.data(), Some(&[1, 2, 3, 4, 5, 6, 7, 8][..]));
2220        assert_eq!(assembler.chunk_count(), 3);
2221    }
2222
2223    #[test]
2224    fn test_stream_assembler_out_of_order() {
2225        let mut assembler = StreamAssembler::new();
2226
2227        let chunk1 = StreamChunk::first(0, 3);
2228        assembler.add_chunk(chunk1.as_bytes(), &[1, 2, 3]).unwrap();
2229
2230        // Try to add chunk with wrong sequence
2231        let chunk_wrong = StreamChunk::continuation(5, 3);
2232        let result = assembler.add_chunk(chunk_wrong.as_bytes(), &[4, 5, 6]);
2233        assert!(result.is_err());
2234    }
2235
2236    #[test]
2237    fn test_stream_assembler_error_chunk() {
2238        let mut assembler = StreamAssembler::new();
2239
2240        let chunk = StreamChunk::error(0, 0);
2241        let result = assembler.add_chunk(chunk.as_bytes(), &[]);
2242
2243        assert!(result.is_err());
2244        assert!(assembler.has_error());
2245        assert!(assembler.is_complete());
2246        assert!(assembler.data().is_none());
2247    }
2248
2249    #[test]
2250    fn test_stream_assembler_take_data() {
2251        let mut assembler = StreamAssembler::new();
2252
2253        let chunk = StreamChunk::new(0, ChunkFlags::FIRST | ChunkFlags::LAST, 3);
2254        assembler.add_chunk(chunk.as_bytes(), &[1, 2, 3]).unwrap();
2255
2256        let data = assembler.take_data();
2257        assert_eq!(data, Some(vec![1, 2, 3]));
2258    }
2259
2260    // Stream Builder Tests
2261
2262    #[test]
2263    fn test_stream_builder_single_chunk() {
2264        let mut builder = StreamBuilder::new(100);
2265        let data = vec![1, 2, 3, 4, 5];
2266
2267        let chunks = builder.build_chunks(&data);
2268
2269        assert_eq!(chunks.len(), 1);
2270        let (chunk, payload) = &chunks[0];
2271        assert!(chunk.is_first());
2272        assert!(chunk.is_last());
2273        assert_eq!(payload, &data);
2274    }
2275
2276    #[test]
2277    fn test_stream_builder_multiple_chunks() {
2278        let mut builder = StreamBuilder::new(3);
2279        let data = vec![1, 2, 3, 4, 5, 6, 7, 8];
2280
2281        let chunks = builder.build_chunks(&data);
2282
2283        assert_eq!(chunks.len(), 3);
2284
2285        // First chunk
2286        assert!(chunks[0].0.is_first());
2287        assert!(!chunks[0].0.is_last());
2288        assert_eq!(chunks[0].1, vec![1, 2, 3]);
2289
2290        // Middle chunk
2291        assert!(chunks[1].0.is_continuation());
2292        assert!(!chunks[1].0.is_last());
2293        assert_eq!(chunks[1].1, vec![4, 5, 6]);
2294
2295        // Last chunk
2296        assert!(!chunks[2].0.is_first());
2297        assert!(chunks[2].0.is_last());
2298        assert_eq!(chunks[2].1, vec![7, 8]);
2299    }
2300
2301    #[test]
2302    fn test_stream_builder_reset() {
2303        let mut builder = StreamBuilder::new(10);
2304
2305        let _ = builder.build_chunks(&[1, 2, 3]);
2306        builder.reset();
2307
2308        let chunks = builder.build_chunks(&[4, 5, 6]);
2309        // Copy sequence from packed struct
2310        let sequence = chunks[0].0.sequence;
2311        assert_eq!(sequence, 0);
2312    }
2313
2314    #[test]
2315    fn test_stream_builder_assembler_roundtrip() {
2316        let mut builder = StreamBuilder::new(5);
2317        let original_data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
2318
2319        let chunks = builder.build_chunks(&original_data);
2320
2321        let mut assembler = StreamAssembler::new();
2322        for (chunk, payload) in chunks {
2323            assembler.add_chunk(chunk.as_bytes(), &payload).unwrap();
2324        }
2325
2326        assert!(assembler.is_complete());
2327        assert_eq!(assembler.take_data(), Some(original_data));
2328    }
2329
2330    // Client Message Creation Tests
2331
2332    #[test]
2333    fn test_client_create_tool_message() {
2334        let mut client = DcpClient::with_defaults();
2335        client.connect("tcp://localhost:9000").unwrap();
2336
2337        let tool_id = client.register_tool("test", "Test tool", [0; 32], 0);
2338        let args = vec![1, 2, 3, 4];
2339
2340        let message = client.create_tool_message(tool_id, &args).unwrap();
2341
2342        assert_eq!(message.message_type, MessageType::Tool);
2343        assert!(message.payload.len() >= ToolInvocation::SIZE + args.len());
2344    }
2345
2346    #[test]
2347    fn test_client_create_resource_message() {
2348        let client = DcpClient::with_defaults();
2349        let message = client.create_resource_message("file:///test.txt");
2350
2351        assert_eq!(message.message_type, MessageType::Resource);
2352        assert_eq!(message.payload, b"file:///test.txt");
2353    }
2354
2355    #[test]
2356    fn test_client_create_prompt_message() {
2357        let client = DcpClient::with_defaults();
2358        let message = client.create_prompt_message("Hello, AI!");
2359
2360        assert_eq!(message.message_type, MessageType::Prompt);
2361        assert_eq!(message.payload, b"Hello, AI!");
2362    }
2363
2364    #[test]
2365    fn test_client_create_error_message() {
2366        let client = DcpClient::with_defaults();
2367        let message = client.create_error_message(404, "Not found");
2368
2369        assert_eq!(message.message_type, MessageType::Error);
2370        // First 4 bytes are error code
2371        assert_eq!(&message.payload[0..4], &404u32.to_le_bytes());
2372        assert_eq!(&message.payload[4..], b"Not found");
2373    }
2374
2375    #[test]
2376    fn test_client_encode_decode_message() {
2377        let client = DcpClient::with_defaults();
2378
2379        let payload = b"test payload".to_vec();
2380        let encoded = client.encode_message(MessageType::Response, &payload);
2381        let decoded = client.decode_message(&encoded).unwrap();
2382
2383        assert_eq!(decoded.message_type, MessageType::Response);
2384        assert_eq!(decoded.payload, payload);
2385    }
2386
2387    #[test]
2388    fn test_client_stream_builder_assembler() {
2389        let client = DcpClient::with_defaults();
2390
2391        let mut builder = client.create_stream_builder(10);
2392        let assembler = client.create_stream_assembler();
2393
2394        // Verify they work
2395        let chunks = builder.build_chunks(&[1, 2, 3]);
2396        assert_eq!(chunks.len(), 1);
2397        assert!(!assembler.is_complete());
2398    }
2399
2400    // ZCTI (Zero-Copy Tool Invocation) Tests
2401
2402    #[test]
2403    fn test_zcti_builder_basic() {
2404        let builder = ZctiBuilder::new(42)
2405            .add_bool(true)
2406            .add_i32(123)
2407            .add_string("hello");
2408
2409        assert_eq!(builder.arg_count(), 3);
2410
2411        let (invocation, data) = builder.build_with_data();
2412        assert_eq!(invocation.tool_id, 42);
2413        assert_eq!(invocation.arg_count(), 3);
2414        assert!(!data.is_empty());
2415    }
2416
2417    #[test]
2418    fn test_zcti_builder_all_types() {
2419        let builder = ZctiBuilder::new(1)
2420            .add_null()
2421            .add_bool(false)
2422            .add_i32(-42)
2423            .add_i64(i64::MAX)
2424            .add_f64(3.14159)
2425            .add_string("test")
2426            .add_bytes(&[0xDE, 0xAD, 0xBE, 0xEF]);
2427
2428        assert_eq!(builder.arg_count(), 7);
2429    }
2430
2431    #[test]
2432    fn test_zcti_reader_basic() {
2433        let builder = ZctiBuilder::new(42)
2434            .add_bool(true)
2435            .add_i32(123)
2436            .add_string("hello");
2437
2438        let (invocation, data) = builder.build_with_data();
2439        let mut reader = ZctiReader::new(&invocation, &data);
2440
2441        assert_eq!(reader.tool_id(), 42);
2442        assert_eq!(reader.arg_count(), 3);
2443        assert!(reader.has_more());
2444
2445        assert_eq!(reader.read_bool().unwrap(), true);
2446        assert_eq!(reader.read_i32().unwrap(), 123);
2447        assert_eq!(reader.read_string().unwrap(), "hello");
2448        assert!(!reader.has_more());
2449    }
2450
2451    #[test]
2452    fn test_zcti_reader_all_types() {
2453        let builder = ZctiBuilder::new(1)
2454            .add_bool(false)
2455            .add_i32(-42)
2456            .add_i64(i64::MAX)
2457            .add_f64(3.14159)
2458            .add_string("test")
2459            .add_bytes(&[0xDE, 0xAD, 0xBE, 0xEF]);
2460
2461        let (invocation, data) = builder.build_with_data();
2462        let mut reader = ZctiReader::new(&invocation, &data);
2463
2464        assert_eq!(reader.read_bool().unwrap(), false);
2465        assert_eq!(reader.read_i32().unwrap(), -42);
2466        assert_eq!(reader.read_i64().unwrap(), i64::MAX);
2467        assert!((reader.read_f64().unwrap() - 3.14159).abs() < 0.00001);
2468        assert_eq!(reader.read_string().unwrap(), "test");
2469        assert_eq!(reader.read_bytes().unwrap(), &[0xDE, 0xAD, 0xBE, 0xEF]);
2470    }
2471
2472    #[test]
2473    fn test_zcti_reader_skip() {
2474        let builder = ZctiBuilder::new(1)
2475            .add_bool(true)
2476            .add_i32(123)
2477            .add_string("skip me")
2478            .add_i64(456);
2479
2480        let (invocation, data) = builder.build_with_data();
2481        let mut reader = ZctiReader::new(&invocation, &data);
2482
2483        reader.skip().unwrap(); // Skip bool
2484        reader.skip().unwrap(); // Skip i32
2485        reader.skip().unwrap(); // Skip string
2486        assert_eq!(reader.read_i64().unwrap(), 456);
2487    }
2488
2489    #[test]
2490    fn test_zcti_type_mismatch() {
2491        let builder = ZctiBuilder::new(1).add_bool(true);
2492        let (invocation, data) = builder.build_with_data();
2493        let mut reader = ZctiReader::new(&invocation, &data);
2494
2495        // Try to read as wrong type
2496        let result = reader.read_i32();
2497        assert!(result.is_err());
2498    }
2499
2500    #[test]
2501    fn test_shared_arg_buffer() {
2502        let mut buffer = SharedArgBuffer::new(1024);
2503
2504        let data1 = b"hello";
2505        let offset1 = buffer.write(data1).unwrap();
2506        assert_eq!(offset1, 0);
2507
2508        let data2 = b"world";
2509        let offset2 = buffer.write(data2).unwrap();
2510        assert_eq!(offset2, 5);
2511
2512        assert_eq!(buffer.read(offset1, 5).unwrap(), b"hello");
2513        assert_eq!(buffer.read(offset2, 5).unwrap(), b"world");
2514    }
2515
2516    #[test]
2517    fn test_shared_arg_buffer_overflow() {
2518        let mut buffer = SharedArgBuffer::new(10);
2519
2520        let result = buffer.write(&[0u8; 20]);
2521        assert!(result.is_err());
2522    }
2523
2524    #[test]
2525    fn test_shared_arg_buffer_reset() {
2526        let mut buffer = SharedArgBuffer::new(100);
2527
2528        buffer.write(b"test").unwrap();
2529        assert_eq!(buffer.offset(), 4);
2530
2531        buffer.reset();
2532        assert_eq!(buffer.offset(), 0);
2533    }
2534
2535    #[test]
2536    fn test_client_zcti_integration() {
2537        let mut client = DcpClient::with_defaults();
2538        client.connect("tcp://localhost:9000").unwrap();
2539
2540        let tool_id = client.register_tool("test_tool", "Test", [0; 32], 0);
2541
2542        let builder = client
2543            .create_zcti_builder(tool_id)
2544            .unwrap()
2545            .add_string("arg1")
2546            .add_i32(42);
2547
2548        let encoded = client.invoke_tool_zcti(builder).unwrap();
2549
2550        // Decode and verify
2551        let message = client.decode_message(&encoded).unwrap();
2552        assert_eq!(message.message_type, MessageType::Tool);
2553
2554        let mut reader = client.parse_tool_invocation(&message).unwrap();
2555        assert_eq!(reader.tool_id(), tool_id);
2556        assert_eq!(reader.read_string().unwrap(), "arg1");
2557        assert_eq!(reader.read_i32().unwrap(), 42);
2558    }
2559
2560    // MCP Compatibility Tests
2561
2562    #[test]
2563    fn test_mcp_adapter_available() {
2564        let mut client = DcpClient::with_defaults();
2565        client.connect("tcp://localhost:9000").unwrap();
2566
2567        assert!(client.has_mcp_fallback());
2568        assert!(client.mcp_adapter().is_some());
2569    }
2570
2571    #[test]
2572    fn test_mcp_register_tool() {
2573        let mut client = DcpClient::with_defaults();
2574        client.connect("tcp://localhost:9000").unwrap();
2575
2576        let tool_id = client.register_tool("read_file", "Read a file", [0; 32], 0);
2577        client.register_tool_mcp("read_file", tool_id);
2578
2579        let adapter = client.mcp_adapter().unwrap();
2580        assert_eq!(adapter.resolve_tool_name("read_file"), Some(tool_id as u16));
2581    }
2582
2583    #[test]
2584    fn test_mcp_handle_initialize() {
2585        let mut client = DcpClient::with_defaults();
2586        client.connect("tcp://localhost:9000").unwrap();
2587
2588        let request = r#"{"jsonrpc":"2.0","method":"initialize","id":1}"#;
2589        let response = client.handle_mcp_request(request).unwrap();
2590
2591        assert!(response.contains("protocolVersion"));
2592        assert!(response.contains("capabilities"));
2593    }
2594
2595    #[test]
2596    fn test_mcp_handle_tools_list() {
2597        let mut client = DcpClient::with_defaults();
2598        client.connect("tcp://localhost:9000").unwrap();
2599
2600        let tool_id = client.register_tool("test_tool", "Test tool", [0; 32], 0);
2601        client.register_tool_mcp("test_tool", tool_id);
2602
2603        let request = r#"{"jsonrpc":"2.0","method":"tools/list","id":2}"#;
2604        let response = client.handle_mcp_request(request).unwrap();
2605
2606        assert!(response.contains("tools"));
2607        assert!(response.contains("test_tool"));
2608    }
2609
2610    #[test]
2611    fn test_mcp_handle_tools_call() {
2612        let mut client = DcpClient::with_defaults();
2613        client.connect("tcp://localhost:9000").unwrap();
2614
2615        let tool_id = client.register_tool("echo", "Echo tool", [0; 32], 0);
2616        client.register_tool_mcp("echo", tool_id);
2617
2618        let request = r#"{"jsonrpc":"2.0","method":"tools/call","params":{"name":"echo","arguments":{"text":"hello"}},"id":3}"#;
2619        let response = client.handle_mcp_request(request).unwrap();
2620
2621        assert!(response.contains("content"));
2622        assert!(response.contains("echo"));
2623    }
2624
2625    #[test]
2626    fn test_mcp_handle_unknown_method() {
2627        let mut client = DcpClient::with_defaults();
2628        client.connect("tcp://localhost:9000").unwrap();
2629
2630        let request = r#"{"jsonrpc":"2.0","method":"unknown/method","id":4}"#;
2631        let response = client.handle_mcp_request(request).unwrap();
2632
2633        assert!(response.contains("error"));
2634        assert!(response.contains("-32601")); // Method not found
2635    }
2636
2637    #[test]
2638    fn test_mcp_handle_resources_list() {
2639        let mut client = DcpClient::with_defaults();
2640        client.connect("tcp://localhost:9000").unwrap();
2641
2642        let request = r#"{"jsonrpc":"2.0","method":"resources/list","id":5}"#;
2643        let response = client.handle_mcp_request(request).unwrap();
2644
2645        assert!(response.contains("resources"));
2646        assert!(response.contains("[]")); // Empty list
2647    }
2648
2649    #[test]
2650    fn test_mcp_handle_prompts_list() {
2651        let mut client = DcpClient::with_defaults();
2652        client.connect("tcp://localhost:9000").unwrap();
2653
2654        let request = r#"{"jsonrpc":"2.0","method":"prompts/list","id":6}"#;
2655        let response = client.handle_mcp_request(request).unwrap();
2656
2657        assert!(response.contains("prompts"));
2658        assert!(response.contains("[]")); // Empty list
2659    }
2660
2661    #[test]
2662    fn test_dcp_to_mcp_response() {
2663        let client = DcpClient::with_defaults();
2664
2665        let payload = serde_json::to_vec(&serde_json::json!({"result": "success"})).unwrap();
2666        let message = DcpMessage::response(payload);
2667
2668        let mcp_json = client.dcp_to_mcp(&message).unwrap();
2669        assert!(mcp_json.contains("jsonrpc"));
2670        assert!(mcp_json.contains("result"));
2671    }
2672
2673    #[test]
2674    fn test_dcp_to_mcp_error() {
2675        let client = DcpClient::with_defaults();
2676
2677        let message = client.create_error_message(404, "Not found");
2678        let mcp_json = client.dcp_to_mcp(&message).unwrap();
2679
2680        assert!(mcp_json.contains("error"));
2681        assert!(mcp_json.contains("404"));
2682        assert!(mcp_json.contains("Not found"));
2683    }
2684
2685    #[test]
2686    fn test_mcp_to_dcp() {
2687        let mut client = DcpClient::with_defaults();
2688        client.connect("tcp://localhost:9000").unwrap();
2689
2690        let mcp_request =
2691            r#"{"jsonrpc":"2.0","method":"tools/call","params":{"name":"test"},"id":1}"#;
2692        let dcp_message = client.mcp_to_dcp(mcp_request).unwrap();
2693
2694        assert_eq!(dcp_message.message_type, MessageType::Tool);
2695        assert!(!dcp_message.payload.is_empty());
2696    }
2697
2698    // Property-Based Test Generators
2699
2700    /// Generate an arbitrary MessageType
2701    fn arb_message_type() -> impl Strategy<Value = MessageType> {
2702        prop_oneof![
2703            Just(MessageType::Tool),
2704            Just(MessageType::Resource),
2705            Just(MessageType::Prompt),
2706            Just(MessageType::Response),
2707            Just(MessageType::Error),
2708            Just(MessageType::Stream),
2709        ]
2710    }
2711
2712    /// Generate arbitrary flags (valid combinations)
2713    fn arb_flags() -> impl Strategy<Value = u8> {
2714        prop_oneof![
2715            Just(0u8),
2716            Just(Flags::STREAMING),
2717            Just(Flags::COMPRESSED),
2718            Just(Flags::SIGNED),
2719            Just(Flags::STREAMING | Flags::COMPRESSED),
2720            Just(Flags::STREAMING | Flags::SIGNED),
2721            Just(Flags::COMPRESSED | Flags::SIGNED),
2722            Just(Flags::STREAMING | Flags::COMPRESSED | Flags::SIGNED),
2723        ]
2724    }
2725
2726    /// Generate arbitrary payload (limited size for testing)
2727    fn arb_payload() -> impl Strategy<Value = Vec<u8>> {
2728        prop::collection::vec(any::<u8>(), 0..1024)
2729    }
2730
2731    /// Generate an arbitrary DcpMessage
2732    fn arb_dcp_message() -> impl Strategy<Value = DcpMessage> {
2733        (arb_message_type(), arb_flags(), arb_payload()).prop_map(
2734            |(message_type, flags, payload)| DcpMessage::new(message_type, flags, payload),
2735        )
2736    }
2737
2738    /// Generate arbitrary chunk size (reasonable range)
2739    fn arb_chunk_size() -> impl Strategy<Value = usize> {
2740        1usize..256
2741    }
2742
2743    /// Generate arbitrary data for streaming
2744    fn arb_stream_data() -> impl Strategy<Value = Vec<u8>> {
2745        prop::collection::vec(any::<u8>(), 1..2048)
2746    }
2747
2748    proptest! {
2749        #![proptest_config(ProptestConfig::with_cases(100))]
2750
2751        /// Property 4: DCP Binary Message Envelope Round-Trip
2752        /// *For any* valid DCP message, encoding to BME format and decoding back
2753        /// SHALL produce an equivalent message.
2754        /// **Validates: Requirements 2.2**
2755        #[test]
2756        fn prop_bme_roundtrip(message in arb_dcp_message()) {
2757            let encoded = message.encode();
2758            let decoded = DcpMessage::decode(&encoded).expect("Decoding should succeed");
2759
2760            prop_assert_eq!(decoded.message_type, message.message_type);
2761            prop_assert_eq!(decoded.flags, message.flags);
2762            prop_assert_eq!(decoded.payload, message.payload);
2763        }
2764
2765        /// Property 4b: BME envelope header is correctly preserved
2766        /// *For any* valid message type, flags, and payload length, the envelope
2767        /// header SHALL be correctly encoded and decoded.
2768        #[test]
2769        fn prop_bme_envelope_header(
2770            message_type in arb_message_type(),
2771            flags in arb_flags(),
2772            payload_len in 0u32..65536
2773        ) {
2774            let envelope = BinaryMessageEnvelope::new(message_type, flags, payload_len);
2775            let bytes = envelope.as_bytes();
2776            let parsed = BinaryMessageEnvelope::from_bytes(bytes).expect("Parsing should succeed");
2777
2778            // Copy values from packed struct
2779            let parsed_type = parsed.message_type;
2780            let parsed_flags = parsed.flags;
2781            let parsed_len = parsed.payload_len;
2782
2783            prop_assert_eq!(parsed_type, message_type as u8);
2784            prop_assert_eq!(parsed_flags, flags);
2785            prop_assert_eq!(parsed_len, payload_len);
2786        }
2787
2788        /// Property 4c: Stream chunking and reassembly round-trip
2789        /// *For any* data and chunk size, chunking and reassembling SHALL produce
2790        /// the original data.
2791        #[test]
2792        fn prop_stream_roundtrip(
2793            data in arb_stream_data(),
2794            chunk_size in arb_chunk_size()
2795        ) {
2796            let mut builder = StreamBuilder::new(chunk_size);
2797            let chunks = builder.build_chunks(&data);
2798
2799            let mut assembler = StreamAssembler::new();
2800            for (chunk, payload) in chunks {
2801                assembler.add_chunk(chunk.as_bytes(), &payload).expect("Adding chunk should succeed");
2802            }
2803
2804            prop_assert!(assembler.is_complete());
2805            prop_assert!(!assembler.has_error());
2806            prop_assert_eq!(assembler.take_data(), Some(data));
2807        }
2808
2809        /// Property 4d: All message types round-trip correctly
2810        /// *For any* message type and payload, encoding and decoding SHALL preserve
2811        /// the message type and payload.
2812        #[test]
2813        fn prop_all_message_types_roundtrip(
2814            message_type in arb_message_type(),
2815            payload in arb_payload()
2816        ) {
2817            let message = DcpMessage::new(message_type, 0, payload.clone());
2818            let encoded = message.encode();
2819            let decoded = DcpMessage::decode(&encoded).expect("Decoding should succeed");
2820
2821            prop_assert_eq!(decoded.message_type, message_type);
2822            prop_assert_eq!(decoded.payload, payload);
2823        }
2824
2825        /// Property 4e: Flags are correctly preserved through round-trip
2826        /// *For any* valid flag combination, encoding and decoding SHALL preserve
2827        /// all flag bits.
2828        #[test]
2829        fn prop_flags_preserved(flags in arb_flags()) {
2830            let message = DcpMessage::new(MessageType::Tool, flags, vec![1, 2, 3]);
2831            let encoded = message.encode();
2832            let decoded = DcpMessage::decode(&encoded).expect("Decoding should succeed");
2833
2834            prop_assert_eq!(decoded.flags, flags);
2835            prop_assert_eq!(decoded.is_streaming(), flags & Flags::STREAMING != 0);
2836            prop_assert_eq!(decoded.is_compressed(), flags & Flags::COMPRESSED != 0);
2837            prop_assert_eq!(decoded.is_signed(), flags & Flags::SIGNED != 0);
2838        }
2839
2840        /// Property 5: DCP MCP Compatibility
2841        /// *For any* valid MCP JSON-RPC request, the McpAdapter SHALL produce a valid
2842        /// response that matches MCP protocol semantics.
2843        /// **Validates: Requirements 2.4**
2844        #[test]
2845        fn prop_mcp_initialize_response_valid(id in 1i64..1000) {
2846            let request = format!(r#"{{"jsonrpc":"2.0","method":"initialize","id":{}}}"#, id);
2847
2848            let mut client = DcpClient::with_defaults();
2849            client.connect("tcp://localhost:9000").expect("Connection should succeed");
2850
2851            let response = client.handle_mcp_request(&request).expect("Should handle initialize");
2852
2853            // Verify response is valid JSON-RPC
2854            let parsed: serde_json::Value = serde_json::from_str(&response).expect("Should be valid JSON");
2855            prop_assert_eq!(parsed["jsonrpc"].as_str(), Some("2.0"));
2856            prop_assert!(parsed.get("result").is_some() || parsed.get("error").is_some());
2857
2858            // If success, verify capabilities
2859            if let Some(result) = parsed.get("result") {
2860                prop_assert!(result.get("capabilities").is_some());
2861                prop_assert!(result.get("protocolVersion").is_some());
2862            }
2863        }
2864
2865        /// Property 5b: MCP tools/list returns registered tools
2866        /// *For any* set of registered tools, tools/list SHALL return all of them.
2867        #[test]
2868        fn prop_mcp_tools_list_contains_registered(
2869            tool_names in prop::collection::vec("[a-z_]+", 1..5)
2870        ) {
2871            let mut client = DcpClient::with_defaults();
2872            client.connect("tcp://localhost:9000").expect("Connection should succeed");
2873
2874            // Register tools
2875            for name in &tool_names {
2876                let tool_id = client.register_tool(name, "Test tool", [0; 32], 0);
2877                client.register_tool_mcp(name, tool_id);
2878            }
2879
2880            let request = r#"{"jsonrpc":"2.0","method":"tools/list","id":1}"#;
2881            let response = client.handle_mcp_request(request).expect("Should handle tools/list");
2882
2883            let parsed: serde_json::Value = serde_json::from_str(&response).expect("Should be valid JSON");
2884            let tools = parsed["result"]["tools"].as_array().expect("Should have tools array");
2885
2886            // Verify all registered tools are in the response
2887            for name in &tool_names {
2888                let found = tools.iter().any(|t| t["name"].as_str() == Some(name.as_str()));
2889                prop_assert!(found, "Tool {} should be in response", name);
2890            }
2891        }
2892
2893        /// Property 6: DCP Capability Manifest Intersection
2894        /// *For any* two capability manifests, their intersection SHALL contain only
2895        /// capabilities present in both.
2896        /// **Validates: Requirements 2.6**
2897        #[test]
2898        fn prop_capability_intersection(
2899            tools1 in prop::collection::vec(0u16..1000, 0..20),
2900            tools2 in prop::collection::vec(0u16..1000, 0..20),
2901            resources1 in prop::collection::vec(0u16..500, 0..10),
2902            resources2 in prop::collection::vec(0u16..500, 0..10),
2903            prompts1 in prop::collection::vec(0u16..200, 0..5),
2904            prompts2 in prop::collection::vec(0u16..200, 0..5),
2905        ) {
2906            let mut m1 = CapabilityManifest::new(1);
2907            let mut m2 = CapabilityManifest::new(1);
2908
2909            // Set capabilities in m1
2910            for &t in &tools1 { m1.set_tool(t); }
2911            for &r in &resources1 { m1.set_resource(r); }
2912            for &p in &prompts1 { m1.set_prompt(p); }
2913
2914            // Set capabilities in m2
2915            for &t in &tools2 { m2.set_tool(t); }
2916            for &r in &resources2 { m2.set_resource(r); }
2917            for &p in &prompts2 { m2.set_prompt(p); }
2918
2919            // Compute intersection
2920            let intersection = m1.intersect(&m2);
2921
2922            // Verify: intersection only contains capabilities in BOTH manifests
2923            for t in 0u16..1000 {
2924                let in_both = m1.has_tool(t) && m2.has_tool(t);
2925                prop_assert_eq!(
2926                    intersection.has_tool(t), in_both,
2927                    "Tool {} should be in intersection iff in both manifests", t
2928                );
2929            }
2930
2931            for r in 0u16..500 {
2932                let in_both = m1.has_resource(r) && m2.has_resource(r);
2933                prop_assert_eq!(
2934                    intersection.has_resource(r), in_both,
2935                    "Resource {} should be in intersection iff in both manifests", r
2936                );
2937            }
2938
2939            for p in 0u16..200 {
2940                let in_both = m1.has_prompt(p) && m2.has_prompt(p);
2941                prop_assert_eq!(
2942                    intersection.has_prompt(p), in_both,
2943                    "Prompt {} should be in intersection iff in both manifests", p
2944                );
2945            }
2946        }
2947
2948        /// Property 6b: Capability manifest round-trip through bytes
2949        /// *For any* capability manifest, serializing and parsing SHALL produce
2950        /// an equivalent manifest.
2951        #[test]
2952        fn prop_capability_manifest_roundtrip(
2953            tools in prop::collection::vec(0u16..8000, 0..50),
2954            resources in prop::collection::vec(0u16..1000, 0..20),
2955            prompts in prop::collection::vec(0u16..500, 0..10),
2956            extensions in 0u64..u64::MAX,
2957        ) {
2958            let mut manifest = CapabilityManifest::new(1);
2959
2960            for &t in &tools { manifest.set_tool(t); }
2961            for &r in &resources { manifest.set_resource(r); }
2962            for &p in &prompts { manifest.set_prompt(p); }
2963            manifest.extensions = extensions;
2964
2965            // Serialize and parse
2966            let bytes = manifest.as_bytes();
2967            let parsed = CapabilityManifest::from_bytes(bytes).expect("Should parse");
2968
2969            // Verify all capabilities preserved
2970            for &t in &tools {
2971                prop_assert!(parsed.has_tool(t), "Tool {} should be preserved", t);
2972            }
2973            for &r in &resources {
2974                prop_assert!(parsed.has_resource(r), "Resource {} should be preserved", r);
2975            }
2976            for &p in &prompts {
2977                prop_assert!(parsed.has_prompt(p), "Prompt {} should be preserved", p);
2978            }
2979            prop_assert_eq!(parsed.extensions, extensions);
2980        }
2981
2982        /// Property 7: Ed25519 Signature Verification
2983        /// *For any* signed tool definition, verification with the correct public key
2984        /// SHALL succeed, and verification with any other key SHALL fail.
2985        /// **Validates: Requirements 2.7**
2986        #[test]
2987        fn prop_signature_verification_tool_def(
2988            seed in prop::array::uniform32(any::<u8>()),
2989            tool_id in 1u32..1000,
2990            schema_hash in prop::array::uniform32(any::<u8>()),
2991            capabilities in any::<u64>(),
2992        ) {
2993            let signer = Signer::from_seed(&seed);
2994            let signed_def = signer.sign_tool_def(tool_id, schema_hash, capabilities);
2995
2996            // Verification with correct key should succeed
2997            prop_assert!(Verifier::verify_tool_def(&signed_def).is_ok());
2998
2999            // Verification should fail if any field is tampered
3000            let mut tampered = signed_def.clone();
3001            tampered.tool_id = tool_id.wrapping_add(1);
3002            prop_assert!(Verifier::verify_tool_def(&tampered).is_err());
3003
3004            let mut tampered = signed_def.clone();
3005            tampered.schema_hash[0] ^= 0xFF;
3006            prop_assert!(Verifier::verify_tool_def(&tampered).is_err());
3007
3008            let mut tampered = signed_def.clone();
3009            tampered.capabilities ^= 0xFFFFFFFF;
3010            prop_assert!(Verifier::verify_tool_def(&tampered).is_err());
3011        }
3012
3013        /// Property 7b: Ed25519 Signature Verification for Invocations
3014        /// *For any* signed invocation, verification with the correct public key
3015        /// SHALL succeed, and verification with any other key SHALL fail.
3016        #[test]
3017        fn prop_signature_verification_invocation(
3018            seed in prop::array::uniform32(any::<u8>()),
3019            wrong_seed in prop::array::uniform32(any::<u8>()),
3020            tool_id in 1u32..1000,
3021            nonce in any::<u64>(),
3022            timestamp in any::<u64>(),
3023            args in prop::collection::vec(any::<u8>(), 0..256),
3024        ) {
3025            let signer = Signer::from_seed(&seed);
3026            let public_key = signer.public_key_bytes();
3027            let signed_inv = signer.sign_invocation(tool_id, nonce, timestamp, &args);
3028
3029            // Verification with correct key should succeed
3030            prop_assert!(Verifier::verify_invocation(&signed_inv, &public_key).is_ok());
3031
3032            // Args hash should match
3033            prop_assert!(Verifier::verify_args_hash(&signed_inv, &args));
3034
3035            // Verification with wrong key should fail (unless seeds happen to be equal)
3036            if seed != wrong_seed {
3037                let wrong_signer = Signer::from_seed(&wrong_seed);
3038                let wrong_key = wrong_signer.public_key_bytes();
3039                prop_assert!(Verifier::verify_invocation(&signed_inv, &wrong_key).is_err());
3040            }
3041
3042            // Verification should fail if any field is tampered
3043            let mut tampered = signed_inv.clone();
3044            tampered.tool_id = tool_id.wrapping_add(1);
3045            prop_assert!(Verifier::verify_invocation(&tampered, &public_key).is_err());
3046
3047            let mut tampered = signed_inv.clone();
3048            tampered.nonce ^= 0xFFFFFFFF;
3049            prop_assert!(Verifier::verify_invocation(&tampered, &public_key).is_err());
3050        }
3051    }
3052
3053    // Capability Manifest Unit Tests
3054
3055    #[test]
3056    fn test_capability_manifest_parsing() {
3057        let mut manifest = CapabilityManifest::new(1);
3058        manifest.set_tool(42);
3059        manifest.set_resource(10);
3060        manifest.set_prompt(5);
3061
3062        let bytes = manifest.as_bytes();
3063        let parsed = DcpClient::parse_capability_manifest(bytes).unwrap();
3064
3065        assert!(parsed.has_tool(42));
3066        assert!(parsed.has_resource(10));
3067        assert!(parsed.has_prompt(5));
3068    }
3069
3070    #[test]
3071    fn test_capability_manifest_serialization() {
3072        let mut client = DcpClient::with_defaults();
3073        client.register_tool("tool1", "Tool 1", [0; 32], 0);
3074        client.register_resource(10);
3075        client.register_prompt(5);
3076
3077        let bytes = client.serialize_capabilities();
3078        let parsed = CapabilityManifest::from_bytes(&bytes).unwrap();
3079
3080        assert!(parsed.has_tool(1)); // First tool gets ID 1
3081        assert!(parsed.has_resource(10));
3082        assert!(parsed.has_prompt(5));
3083    }
3084
3085    #[test]
3086    fn test_resource_registration() {
3087        let mut client = DcpClient::with_defaults();
3088
3089        assert!(!client.has_resource(42));
3090        client.register_resource(42);
3091        assert!(client.has_resource(42));
3092        client.unregister_resource(42);
3093        assert!(!client.has_resource(42));
3094    }
3095
3096    #[test]
3097    fn test_prompt_registration() {
3098        let mut client = DcpClient::with_defaults();
3099
3100        assert!(!client.has_prompt(10));
3101        client.register_prompt(10);
3102        assert!(client.has_prompt(10));
3103        client.unregister_prompt(10);
3104        assert!(!client.has_prompt(10));
3105    }
3106
3107    #[test]
3108    fn test_extension_flags() {
3109        let mut client = DcpClient::with_defaults();
3110
3111        assert!(!client.has_extension(5));
3112        client.set_extension(5);
3113        assert!(client.has_extension(5));
3114        client.clear_extension(5);
3115        assert!(!client.has_extension(5));
3116    }
3117
3118    #[test]
3119    fn test_enforce_tool_success() {
3120        let mut client = DcpClient::with_defaults();
3121        let tool_id = client.register_tool("test", "Test", [0; 32], 0);
3122
3123        assert!(client.enforce_tool(tool_id).is_ok());
3124    }
3125
3126    #[test]
3127    fn test_enforce_tool_failure() {
3128        let client = DcpClient::with_defaults();
3129
3130        // Tool 999 is not registered
3131        assert!(client.enforce_tool(999).is_err());
3132    }
3133
3134    #[test]
3135    fn test_enforce_resource_success() {
3136        let mut client = DcpClient::with_defaults();
3137        client.register_resource(42);
3138
3139        assert!(client.enforce_resource(42).is_ok());
3140    }
3141
3142    #[test]
3143    fn test_enforce_resource_failure() {
3144        let client = DcpClient::with_defaults();
3145
3146        assert!(client.enforce_resource(999).is_err());
3147    }
3148
3149    #[test]
3150    fn test_enforce_prompt_success() {
3151        let mut client = DcpClient::with_defaults();
3152        client.register_prompt(10);
3153
3154        assert!(client.enforce_prompt(10).is_ok());
3155    }
3156
3157    #[test]
3158    fn test_enforce_prompt_failure() {
3159        let client = DcpClient::with_defaults();
3160
3161        assert!(client.enforce_prompt(999).is_err());
3162    }
3163
3164    #[test]
3165    fn test_negotiate_capabilities() {
3166        let mut client = DcpClient::with_defaults();
3167        client.register_tool("tool1", "Tool 1", [0; 32], 0);
3168        client.register_tool("tool2", "Tool 2", [0; 32], 0);
3169        client.register_tool("tool3", "Tool 3", [0; 32], 0);
3170        client.register_resource(1);
3171        client.register_resource(2);
3172
3173        // Server only supports tools 2, 3, 4 and resource 2, 3
3174        let mut server_manifest = CapabilityManifest::new(1);
3175        server_manifest.set_tool(2);
3176        server_manifest.set_tool(3);
3177        server_manifest.set_tool(4);
3178        server_manifest.set_resource(2);
3179        server_manifest.set_resource(3);
3180
3181        let negotiated = client.negotiate_capabilities(&server_manifest);
3182
3183        // After negotiation, only common capabilities remain
3184        assert!(!negotiated.has_tool(1)); // Only client had this
3185        assert!(negotiated.has_tool(2)); // Both had this
3186        assert!(negotiated.has_tool(3)); // Both had this
3187        assert!(!negotiated.has_tool(4)); // Only server had this
3188
3189        assert!(!negotiated.has_resource(1)); // Only client had this
3190        assert!(negotiated.has_resource(2)); // Both had this
3191        assert!(!negotiated.has_resource(3)); // Only server had this
3192    }
3193
3194    #[test]
3195    fn test_capability_stats() {
3196        let mut client = DcpClient::with_defaults();
3197        client.register_tool("tool1", "Tool 1", [0; 32], 0);
3198        client.register_tool("tool2", "Tool 2", [0; 32], 0);
3199        client.register_resource(1);
3200        client.register_prompt(1);
3201        client.register_prompt(2);
3202        client.set_extension(0);
3203
3204        let stats = client.capability_stats();
3205
3206        assert_eq!(stats.tool_count, 2);
3207        assert_eq!(stats.resource_count, 1);
3208        assert_eq!(stats.prompt_count, 2);
3209        assert_eq!(stats.extension_count, 1);
3210        assert_eq!(stats.version, 1);
3211    }
3212
3213    #[test]
3214    fn test_capability_iterators() {
3215        let mut client = DcpClient::with_defaults();
3216        client.register_tool("tool1", "Tool 1", [0; 32], 0);
3217        client.register_tool("tool2", "Tool 2", [0; 32], 0);
3218        client.register_resource(10);
3219        client.register_resource(20);
3220        client.register_prompt(5);
3221
3222        let tool_ids: Vec<_> = client.tool_ids().collect();
3223        assert_eq!(tool_ids, vec![1, 2]);
3224
3225        let resource_ids: Vec<_> = client.resource_ids().collect();
3226        assert_eq!(resource_ids, vec![10, 20]);
3227
3228        let prompt_ids: Vec<_> = client.prompt_ids().collect();
3229        assert_eq!(prompt_ids, vec![5]);
3230    }
3231}