feat: support event driven for sequences

Cargo.toml

@@ -2,7 +2,7 @@
 name = "puma"
 version = "0.0.3"
 edition = "2021"
-description = "A lightweight, high-performance inference engine for local AI."
+description = "A lightweight, high-performance model engine for local AI."
 license = "Apache-2.0"
 repository = "https://github.com/InftyAI/PUMA"
 homepage = "https://github.com/InftyAI/PUMA"
@@ -31,6 +31,7 @@ sysinfo = "0.32"
 rusqlite = { version = "0.32", features = ["bundled"] }
 rusqlite_migration = "1.3"
 regex = "1.11"
+thiserror = "2.0"
 
 # Web server
 axum = "0.7"
@@ -39,6 +40,10 @@ tower-http = { version = "0.5", features = ["cors", "trace"] }
 uuid = { version = "1.0", features = ["v4", "serde"] }
 futures = "0.3"
 tokio-stream = "0.1"
+rustyline = "14.0"
+rustyline-derive = "0.10"
+tokenizers = "0.20"
+async-stream = "0.3"
 
 [dev-dependencies]
 tempfile = "3.12"

Makefile

@@ -1,4 +1,5 @@
 build:
+	rm -rf ./puma
 	cargo build && cp target/debug/puma ./puma
 
 test:

README.md

@@ -6,7 +6,7 @@
   <img alt="PUMA Logo" src="https://raw.githubusercontent.com/InftyAI/PUMA/main/site/images/logo-light.svg" width="240">
 </picture>
 
-**A lightweight, high-performance inference engine for local AI**
+**A lightweight, high-performance model engine for local AI**
 
 [![Stability: Active](https://img.shields.io/badge/stability-active-brightgreen.svg)](https://github.com/InftyAI/PUMA)
 [![Latest Release](https://img.shields.io/github/v/release/InftyAI/PUMA)](https://github.com/InftyAI/PUMA/releases)

docs/fsm_architecture.md

@@ -0,0 +1,240 @@
+# FSM Architecture for PUMA
+
+Finite State Machine pattern for sequence lifecycle management. The FSM states
+and events are defined in `src/fsm/`; the transition logic lives in the
+`SequenceManager` (`src/sequence_manager/mod.rs`), which the `Scheduler` drives.
+
+## Core Concepts
+
+### State Machine per Sequence
+
+Each sequence is an independent state machine. Multiple sequences can be in different states simultaneously.
+
+**Implementation:** `src/fsm/states.rs`
+
+**States:**
+- `Empty` - Pre-birth placeholder; the `from` state of `Create`, never persisted
+- `Waiting` - In queue, no resources allocated
+- `Scheduling` - Being scheduled (transitional)
+- `Prefilling` - Processing prompt
+- `Decoding` - Generating tokens
+- `Preempted` - Temporarily suspended (blocks freed)
+- `Finished` - Completed successfully
+- `Aborted` - Error or cancelled
+
+### States Own Resources
+
+Each state owns the resources it needs:
+- `WaitingState` has no blocks (not allocated yet)
+- `DecodingState` owns blocks via `Arc<Vec<BlockId>>`
+- `PreemptedState` has no blocks (freed)
+
+**Benefit:** Type system enforces resource invariants - can't access blocks that don't exist.
+
+### Events as State Transformations
+
+Events transform states: `Event + OldState → NewState`
+
+**Implementation:** `src/fsm/events.rs`
+
+**Events:**
+- `Create` - Empty → Waiting (birth a new sequence)
+- `Schedule` - Waiting → Prefilling (allocate blocks)
+- `AppendTokens` - Prefilling → Decoding or Decoding → Decoding/Finished
+- `Preempt` - Decoding → Preempted (free blocks on OOM)
+- `Resume` - Preempted → Waiting (re-queue)
+- `Complete` - * → Finished (free blocks)
+- `Abort` - * → Aborted (cleanup)
+
+Invalid transitions return `Error::InvalidTransition`.
+
+Forking (Decoding → (Decoding, Decoding) for beam search) is **not** an event:
+it is a dedicated `SequenceManager::fork` method, since it produces a second
+sequence rather than transforming one state into another.
+
+### FSM Integration with SequenceManager
+
+The FSM logic lives in the **`SequenceManager`** (the *mechanism* layer), which
+owns `BlockManager`. The `Scheduler` (the *policy* layer) drives it with events
+and never touches `BlockManager` or transition logic directly.
+
+**Public API:** `SequenceManager::advance(seq_id, event)` - Event-based abstraction
+- Single entry point for all state transitions
+- Easy to add logging, metrics, debugging
+- Backs up state and restores it on error (no block leaks)
+
+**Internal Implementation:** Private `transition()` / `transition_*()` methods
+- Type-safe helpers that enforce correct state types
+- Direct access to `self.block_manager` - no parameter passing
+- Called by `advance()` after looking up the current state
+
+```rust
+impl SequenceManager {
+    pub fn advance(&mut self, seq_id: SequenceId, event: Event) -> Result<&SequenceState> {
+        // Remove current state, run transition(), reinsert on success
+        // or restore the backup on error.
+    }
+
+    fn transition(&mut self, state: SequenceState, event: Event) -> Result<SequenceState> {
+        let seq_id = state.seq_id();
+        match (state, event) {
+            (SequenceState::Waiting(s), Event::Schedule{..}) =>
+                self.transition_schedule(s),
+            // ... dispatches to internal methods
+        }
+    }
+
+    fn transition_schedule(&mut self, state: WaitingState) -> Result<SequenceState> {
+        // Direct access to self.block_manager, self.tokens_per_block
+    }
+}
+```
+
+## State Diagram
+
+```
+                    AddRequest
+                        ↓
+    ┌─────────────→ Waiting
+    │                   ↓ ScheduleEvent
+    │              Prefilling
+    │                   ↓ AppendTokens
+    │              Decoding
+    │                   ├─→ AppendTokens (continue)
+    │                   │   ├─ Decoding (more tokens)
+    │                   │   └─ Finished (max_tokens)
+    │                   │
+    │                   ├─→ fork() method (beam search, not an event)
+    │                   │   ├─ Parent: Decoding
+    │                   │   └─ Child: Decoding
+    │                   │
+    │                   ├─→ PreemptEvent (OOM)
+    │                   │   ↓ Preempted
+    │                   │   ↓ ResumeEvent
+    │                   └───┘
+    │
+    └─ CompleteEvent → Finished
+```
+
+## Scheduler / SequenceManager Integration
+
+The `SequenceManager` owns the FSM and executes all state transitions; the
+`Scheduler` owns the queues/batches and decides *which* events to fire.
+
+**Implementation:** `src/sequence_manager/mod.rs` (mechanism),
+`src/scheduler/core.rs` (policy)
+
+**Key features:**
+- Synchronous API (LLMEngine handles async coordination)
+- `SequenceManager` maintains `HashMap<SequenceId, SequenceState>`
+- `SequenceManager` owns `BlockManager` directly - no parameter passing needed
+- Clones state before transitions (prevents loss on failure)
+- Restores backup on error (no block leaks)
+- Event-based `advance()` method for maintainability
+
+**Tests:** `cargo test --lib`
+
+## Arc Optimization
+
+States use `Arc<Vec<BlockId>>` for efficient cloning:
+
+- **State cloning** (for backup/fork): **O(1)** - just increment ref count
+- **Modifying blocks**: Copy-on-write via `Arc::make_mut()` - only copies if shared
+- **Fork sequences**: Share blocks via Arc, diverge only when modified
+
+**Benefit:** Hot path performance - state cloning costs ~50 bytes instead of O(n) vector copy.
+
+## Error Handling
+
+All state transitions are wrapped with backup/restore:
+
+1. Remove state from HashMap
+2. Clone state as backup
+3. Apply event transformation
+4. On success: insert new state
+5. On failure: restore backup, log warning
+
+**Prevents:**
+- Sequence loss on transition failure
+- Block leaks (blocks remain owned by restored state)
+- Duplicate sequences (check before insert)
+
+## Type Safety
+
+**Runtime checks:**
+- Invalid transitions caught by pattern matching
+- Return `Error::InvalidTransition` with clear message
+
+**Resource ownership:**
+- States own their blocks via Arc
+- Type system prevents accessing freed blocks
+- Drop trait ensures cleanup
+
+## Benefits
+
+1. **Type Safety** - States enforce resource invariants, invalid transitions caught explicitly
+2. **Clear Ownership** - Resources belong to states, automatic cleanup on drop
+3. **Testability** - Events are pure functions, test state transitions in isolation
+4. **Explicitness** - Every transition is explicit, state diagram maps to code
+5. **Performance** - Arc enables O(1) cloning, copy-on-write for efficiency
+6. **Reliability** - Backup/restore pattern prevents data loss on errors
+
+## Two-Level Event Architecture
+
+PUMA uses a two-level event system (inspired by TokenSpeed):
+
+**Level 1: FSM Events (Internal)** - `src/fsm/events.rs`
+- Created by the Scheduler internally
+- Include scheduler context (tokens_per_block, etc.)
+- Applied via `SequenceManager::advance(seq_id, event)`
+- Examples: `Event::Schedule`, `Event::AppendTokens`
+
+**Level 2: Scheduler Events (External)** - `src/scheduler/events.rs`
+- Created by external components (clients, GPU workers)
+- Simple data, no resource pointers
+- LLMEngine translates these to FSM events
+- Examples: `SchedulerEvent::AddRequest`, `SchedulerEvent::CancelRequest`
+
+**Why two levels?** External components cannot create FSM events because:
+- They don't own `BlockManager`
+- They don't know scheduler policies (tokens_per_block)
+- FSM events require scheduler context
+
+## Key Files
+
+- `src/fsm/states.rs` - State definitions and helper methods
+- `src/fsm/events.rs` - FSM event enum (internal transitions)
+- `src/sequence_manager/mod.rs` - Owns memory + FSM transition logic (`advance()`)
+- `src/scheduler/core.rs` - Scheduling policy that drives the SequenceManager
+- `src/scheduler/events.rs` - External scheduler events
+- `src/backend/llm_engine.rs` - Event coordinator
+
+## Usage Example
+
+```rust
+// External component sends a high-level event
+scheduler_tx.send(SchedulerEvent::AddRequest {
+    seq_id: SequenceId(1),
+    token_ids: vec![/* tokenized prompt */],
+    max_tokens: 100,
+    response_tx, // channel the result is delivered on
+})?;
+
+// The Scheduler translates it to an FSM event and drives the SequenceManager,
+// which owns BlockManager and performs the transition.
+impl Scheduler {
+    fn schedule_prefill(&mut self) {
+        while let Some(seq_id) = self.waiting_queue.pop_front() {
+            // Build the FSM event with scheduler context
+            let event = Event::Schedule {
+                tokens_per_block: self.tokens_per_block,
+            };
+
+            match self.sequences.advance(seq_id, event) {
+                Ok(_) => self.prefill_batch.push(seq_id),
+                Err(e) => { /* apply OOM / queueing policy */ }
+            }
+        }
+    }
+}
+```

src/api/chat.rs

@@ -7,7 +7,6 @@ use axum::{
     Json,
 };
 use futures::stream::StreamExt;
-use std::sync::Arc;
 use tokio_stream::wrappers::ReceiverStream;
 use uuid::Uuid;
 
@@ -16,11 +15,11 @@ use crate::api::types::{
     ChatChoice, ChatChoiceDelta, ChatCompletionChunk, ChatCompletionRequest,
     ChatCompletionResponse, ChatMessage, ChatMessageDelta, ErrorResponse, Usage,
 };
-use crate::backend::InferenceEngine;
+use crate::backend::EngineHandle;
 
 /// Main handler for chat completions
-pub async fn chat_completions<E: InferenceEngine + 'static>(
-    State(state): State<AppState<E>>,
+pub async fn chat_completions(
+    State(state): State<AppState>,
     Json(req): Json<ChatCompletionRequest>,
 ) -> Response {
     let engine = state.engine;
@@ -83,8 +82,8 @@ pub async fn chat_completions<E: InferenceEngine + 'static>(
 }
 
 /// Non-streaming chat completion
-async fn chat_completions_non_stream<E: InferenceEngine>(
-    engine: Arc<E>,
+async fn chat_completions_non_stream(
+    engine: EngineHandle,
     req: ChatCompletionRequest,
 ) -> Result<ChatCompletionResponse, Box<dyn std::error::Error>> {
     let id = format!("chatcmpl-{}", Uuid::new_v4());
@@ -125,8 +124,8 @@ async fn chat_completions_non_stream<E: InferenceEngine>(
 }
 
 /// Streaming chat completion
-async fn chat_completions_stream<E: InferenceEngine + 'static>(
-    engine: Arc<E>,
+async fn chat_completions_stream(
+    engine: EngineHandle,
     req: ChatCompletionRequest,
 ) -> Sse<impl futures::Stream<Item = Result<Event, std::convert::Infallible>>> {
     let id = format!("chatcmpl-{}", Uuid::new_v4());
@@ -236,19 +235,11 @@ async fn chat_completions_stream<E: InferenceEngine + 'static>(
     Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default())
 }
 
-/// Format chat messages into a prompt
+/// Format chat messages into a prompt via the shared prompt formatter.
 fn format_chat_messages(messages: &[ChatMessage]) -> String {
-    messages
-        .iter()
-        .map(|m| {
-            if m.role == "system" {
-                format!("System: {}", m.content)
-            } else if m.role == "user" {
-                format!("User: {}", m.content)
-            } else {
-                format!("Assistant: {}", m.content)
-            }
-        })
-        .collect::<Vec<_>>()
-        .join("\n")
+    crate::utils::prompt::format_conversation(
+        messages
+            .iter()
+            .map(|m| (m.role.as_str(), m.content.as_str())),
+    )
 }

src/api/completions.rs

@@ -5,11 +5,10 @@ use crate::api::routes::AppState;
 use crate::api::types::{
     CompletionChoice, CompletionRequest, CompletionResponse, ErrorResponse, Usage,
 };
-use crate::backend::InferenceEngine;
 
 /// Handler for legacy text completions
-pub async fn completions<E: InferenceEngine + 'static>(
-    State(state): State<AppState<E>>,
+pub async fn completions(
+    State(state): State<AppState>,
     Json(req): Json<CompletionRequest>,
 ) -> impl IntoResponse {
     let engine = state.engine;

src/api/models.rs

@@ -6,12 +6,9 @@ use axum::{
 
 use crate::api::routes::AppState;
 use crate::api::types::{ErrorResponse, Model, ModelList};
-use crate::backend::InferenceEngine;
 
 /// List all available models
-pub async fn list_models<E: InferenceEngine + 'static>(
-    State(state): State<AppState<E>>,
-) -> impl IntoResponse {
+pub async fn list_models(State(state): State<AppState>) -> impl IntoResponse {
     let registry = state.registry;
     match registry.load_models(None) {
         Ok(models) => {
@@ -43,8 +40,8 @@ pub async fn list_models<E: InferenceEngine + 'static>(
 }
 
 /// Get a specific model by ID
-pub async fn get_model<E: InferenceEngine + 'static>(
-    State(state): State<AppState<E>>,
+pub async fn get_model(
+    State(state): State<AppState>,
     Path(model_id): Path<String>,
 ) -> impl IntoResponse {
     let registry = state.registry;