Write #+RESULTS: blocks after code execution

- Execute returns results and writes them to the org file
- Replaces existing #+RESULTS: block if present
- Results formatted with : prefix per org convention
- Process blocks in reverse order to preserve byte offsets
- Blank line added after results for proper separation

crates/ork-server/src/state.rs

@@ -1,6 +1,6 @@
 //! Server state: manages org documents and their parsed state.
 
-use org_ast::{Document, ElementKind, Section, SourceBlock, Timestamp, TimestampValue};
+use org_ast::{Document, ElementKind, Section, SourceBlock, Span, Timestamp, TimestampValue};
 use org_parser::Parser;
 use parking_lot::RwLock;
 use std::collections::HashMap;
@@ -321,7 +321,7 @@ impl OrkState {
             .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "document not found"))?;
 
         // Collect blocks to execute based on scope
-        let blocks: Vec<(String, usize, SourceBlock)> = match args.len() {
+        let blocks: Vec<(String, usize, SourceBlock, Span)> = match args.len() {
             // execute doc — all blocks in document
             1 => self.collect_blocks_doc(&doc_state.doc),
             
@@ -374,23 +374,132 @@ impl OrkState {
             ));
         }
 
-        // Execute each block
+        // Get document content for modification
-        let mut results = Vec::new();
+        let mut content = doc_state.content.clone();
-        for (section_id, idx, block) in &blocks {
+        
-            let result = self.execute_block(block)?;
+        // Execute each block and collect results
-            results.push(format!(
+        // Process in reverse order so byte offsets remain valid as we insert
-                "---\nsection: {}\nindex: {}\nlang: {}\nstatus: {}\nstdout: |\n{}\nstderr: |\n{}",
+        let mut results_output = Vec::new();
-                section_id,
+        let mut blocks_with_results: Vec<_> = blocks.iter().map(|(sid, idx, block, span)| {
-                idx,
+            let result = self.execute_block(block);
-                block.language.as_deref().unwrap_or("unknown"),
+            (sid.clone(), *idx, block.clone(), *span, result)
-                result.status,
+        }).collect();
-                indent_lines(&result.stdout, "  "),
+        
-                indent_lines(&result.stderr, "  "),
+        // Sort by span end position descending so we modify from end to start
-            ));
+        blocks_with_results.sort_by(|a, b| b.3.end.offset.cmp(&a.3.end.offset));
+        
+        for (section_id, idx, block, span, result) in &blocks_with_results {
+            match result {
+                Ok(exec_result) => {
+                    // Insert results into document
+                    content = self.insert_results(&content, *span, &exec_result.stdout);
+                    
+                    results_output.push(format!(
+                        "---\nsection: {}\nindex: {}\nlang: {}\nstatus: {}\nstdout: |\n{}\nstderr: |\n{}",
+                        section_id,
+                        idx,
+                        block.language.as_deref().unwrap_or("unknown"),
+                        exec_result.status,
+                        indent_lines(&exec_result.stdout, "  "),
+                        indent_lines(&exec_result.stderr, "  "),
+                    ));
+                }
+                Err(e) => {
+                    results_output.push(format!(
+                        "---\nsection: {}\nindex: {}\nlang: {}\nerror: {}",
+                        section_id,
+                        idx,
+                        block.language.as_deref().unwrap_or("unknown"),
+                        e,
+                    ));
+                }
+            }
         }
+        
+        // Write updated document back
+        self.write_doc_raw(doc_name, &content)?;
 
-        Ok(results.join("\n"))
+        // Reverse to get original order
+        results_output.reverse();
+        Ok(results_output.join("\n"))
     }
+    
+    /// Insert #+RESULTS: block after a source block.
+    fn insert_results(&self, content: &str, block_span: Span, stdout: &str) -> String {
+        let end_offset = block_span.end.offset;
+        
+        // Find the end of the line after the block (in case span doesn't include newline)
+        let insert_pos = content[end_offset..]
+            .find('\n')
+            .map(|i| end_offset + i + 1)
+            .unwrap_or(content.len());
+        
+        // Check if there's already a #+RESULTS: block after this
+        let after_block = &content[insert_pos..];
+        let results_prefix = "#+RESULTS:";
+        
+        // Look for #+RESULTS: at the start of the next non-empty line
+        let (replace_start, replace_end) = if let Some(_results_start) = after_block
+            .lines()
+            .next()
+            .filter(|line| line.trim().starts_with(results_prefix))
+        {
+            // Find the extent of the existing results block
+            let abs_start = insert_pos;
+            let mut abs_end = abs_start;
+            
+            // Results block continues until we hit a non-result line
+            // (lines starting with : or blank lines after #+RESULTS:)
+            let mut in_results = false;
+            for line in after_block.lines() {
+                let trimmed = line.trim();
+                if trimmed.starts_with(results_prefix) {
+                    in_results = true;
+                    abs_end += line.len() + 1; // +1 for newline
+                } else if in_results && (trimmed.starts_with(':') || trimmed.is_empty()) {
+                    abs_end += line.len() + 1;
+                } else if in_results {
+                    // End of results block
+                    break;
+                } else {
+                    // No results block found
+                    break;
+                }
+            }
+            (abs_start, abs_end.min(content.len()))
+        } else {
+            // No existing results, just insert
+            (insert_pos, insert_pos)
+        };
+        
+        // Format the new results
+        let formatted_results = format_results(stdout);
+        
+        // Build new content
+        let mut new_content = String::with_capacity(content.len() + formatted_results.len());
+        new_content.push_str(&content[..replace_start]);
+        new_content.push_str(&formatted_results);
+        new_content.push_str(&content[replace_end..]);
+        
+        new_content
+    }
+}
+
+/// Format stdout as org #+RESULTS: block.
+fn format_results(stdout: &str) -> String {
+    if stdout.is_empty() {
+        return String::new();
+    }
+    
+    let mut result = String::from("#+RESULTS:\n");
+    for line in stdout.lines() {
+        result.push_str(": ");
+        result.push_str(line);
+        result.push('\n');
+    }
+    // Add blank line after results for separation
+    result.push('\n');
+    result
 }
 
 /// Indent each line of text with a prefix.
@@ -410,7 +519,8 @@ fn indent_lines(text: &str, prefix: &str) -> String {
 
 impl OrkState {
     /// Collect all source blocks from a document.
-    fn collect_blocks_doc(&self, doc: &Document) -> Vec<(String, usize, SourceBlock)> {
+    /// Returns (section_id, element_index, SourceBlock, element_span).
+    fn collect_blocks_doc(&self, doc: &Document) -> Vec<(String, usize, SourceBlock, Span)> {
         let mut blocks = Vec::new();
         for section in &doc.sections {
             blocks.extend(self.collect_blocks_subtree(section));
@@ -419,14 +529,14 @@ impl OrkState {
     }
 
     /// Collect all source blocks from a section and its children (DFS).
-    fn collect_blocks_subtree(&self, section: &Section) -> Vec<(String, usize, SourceBlock)> {
+    fn collect_blocks_subtree(&self, section: &Section) -> Vec<(String, usize, SourceBlock, Span)> {
         let mut blocks = Vec::new();
         let section_id = self.section_id(section);
         
         // Blocks in this section
         for (idx, elem) in section.content.iter().enumerate() {
             if let ElementKind::SourceBlock(sb) = &elem.kind {
-                blocks.push((section_id.clone(), idx, sb.clone()));
+                blocks.push((section_id.clone(), idx, sb.clone(), elem.span));
             }
         }
         
@@ -439,7 +549,7 @@ impl OrkState {
     }
 
     /// Find a named block anywhere in the document.
-    fn find_named_block_doc(&self, doc: &Document, name: &str) -> Result<Vec<(String, usize, SourceBlock)>> {
+    fn find_named_block_doc(&self, doc: &Document, name: &str) -> Result<Vec<(String, usize, SourceBlock, Span)>> {
         for section in &doc.sections {
             if let Some(result) = self.find_named_block_recursive(section, name) {
                 return Ok(vec![result]);
@@ -452,7 +562,7 @@ impl OrkState {
     }
 
     /// Recursively search for a named block.
-    fn find_named_block_recursive(&self, section: &Section, name: &str) -> Option<(String, usize, SourceBlock)> {
+    fn find_named_block_recursive(&self, section: &Section, name: &str) -> Option<(String, usize, SourceBlock, Span)> {
         let section_id = self.section_id(section);
         
         // Check elements for SourceBlock with #+NAME:
@@ -469,14 +579,14 @@ impl OrkState {
                         if aff.key.eq_ignore_ascii_case("NAME") {
                             if let org_ast::AffiliatedValue::Text(ref val) = aff.value {
                                 if val == name {
-                                    return Some((section_id, idx, sb.clone()));
+                                    return Some((section_id, idx, sb.clone(), elem.span));
                                 }
                             }
                         }
                     }
                     // Fall back to preceding Keyword element
                     if pending_name == Some(name) {
-                        return Some((section_id, idx, sb.clone()));
+                        return Some((section_id, idx, sb.clone(), elem.span));
                     }
                     pending_name = None;
                 }
@@ -497,7 +607,7 @@ impl OrkState {
     }
 
     /// Find a named block within a section.
-    fn find_named_block_section(&self, section: &Section, section_id: &str, name: &str) -> Result<Vec<(String, usize, SourceBlock)>> {
+    fn find_named_block_section(&self, section: &Section, section_id: &str, name: &str) -> Result<Vec<(String, usize, SourceBlock, Span)>> {
         let mut pending_name: Option<&str> = None;
         for (idx, elem) in section.content.iter().enumerate() {
             match &elem.kind {
@@ -510,14 +620,14 @@ impl OrkState {
                         if aff.key.eq_ignore_ascii_case("NAME") {
                             if let org_ast::AffiliatedValue::Text(ref val) = aff.value {
                                 if val == name {
-                                    return Ok(vec![(section_id.to_string(), idx, sb.clone())]);
+                                    return Ok(vec![(section_id.to_string(), idx, sb.clone(), elem.span)]);
                                 }
                             }
                         }
                     }
                     // Fall back to preceding Keyword element
                     if pending_name == Some(name) {
-                        return Ok(vec![(section_id.to_string(), idx, sb.clone())]);
+                        return Ok(vec![(section_id.to_string(), idx, sb.clone(), elem.span)]);
                     }
                     pending_name = None;
                 }
@@ -533,12 +643,12 @@ impl OrkState {
     }
 
     /// Get a block by index within a section.
-    fn get_block_by_index(&self, section: &Section, section_id: &str, idx: usize) -> Result<Vec<(String, usize, SourceBlock)>> {
+    fn get_block_by_index(&self, section: &Section, section_id: &str, idx: usize) -> Result<Vec<(String, usize, SourceBlock, Span)>> {
         let blocks: Vec<_> = section.content.iter()
             .enumerate()
             .filter_map(|(i, elem)| {
                 if let ElementKind::SourceBlock(sb) = &elem.kind {
-                    Some((i, sb))
+                    Some((i, sb, elem.span))
                 } else {
                     None
                 }
@@ -552,8 +662,8 @@ impl OrkState {
             ));
         }
         
-        let (elem_idx, sb) = blocks[idx];
+        let (elem_idx, sb, span) = blocks[idx];
-        Ok(vec![(section_id.to_string(), elem_idx, sb.clone())])
+        Ok(vec![(section_id.to_string(), elem_idx, sb.clone(), span)])
     }
 
     /// Execute a source block and return the result.