Originally published on tamiz.pro.

The prevailing narrative in artificial intelligence has been dominated by cloud-based, API-driven models. While this approach offers scalability, it introduces critical latency, dependency on external services, and significant privacy concerns regarding data exfiltration. For mission-critical applications, financial analysis, or healthcare systems, the inability to guarantee data residency and offline operation is a non-starter.

The solution lies in a "Local-First" architecture, where the AI assistant runs entirely on-premise or on-device. However, building such systems requires more than just downloading an LLM weights file; it demands a robust infrastructure layer capable of managing state, memory safety, and real-time concurrency. This article explores how to construct this infrastructure using two powerhouse languages: Go for its superior concurrency primitives and developer velocity in orchestration, and Rust for its memory safety, zero-cost abstractions, and performance-critical inference execution.

We will dissect the architecture of a secure, local-first AI agent, moving from the conceptual model to the implementation details, focusing on the boundary between the orchestration layer (Go) and the execution layer (Rust).

1. The Architectural Paradigm: Separation of Concerns

Building a local-first AI assistant is not merely a software engineering challenge; it is a systems architecture problem. The core tension lies between flexibility (the ability to swap models, adjust prompts, and handle complex workflows) and performance/security (minimizing latency and preventing memory corruption or data leaks).

To resolve this, we adopt a micro-kernel architecture:

  1. The Orchestrator (Go): Handles the user interface, API gateway, session management, tool calling, and high-level logic. Go’s goroutines allow it to manage thousands of concurrent agent sessions with minimal memory overhead.
  2. The Engine (Rust): Handles the heavy lifting: model loading, tokenization, inference, and memory management. Rust ensures that the critical path—where data is processed and potentially sensitive—is free from race conditions, buffer overflows, and undefined behavior.
  3. The Bridge (FFI/gRPC): A thin, strictly typed boundary between the two languages, ensuring that data crossing the perimeter is validated and serialized efficiently.

This separation allows teams to iterate rapidly on the agent's logic in Go (leveraging its vast ecosystem for HTTP servers, database drivers, and UI frameworks) while maintaining a hardened, performant core in Rust.

2. Why Rust for the Inference Engine?

Local-first AI is computationally intensive. Unlike cloud inference, where you can scale horizontally indefinitely, local inference is bound by the hardware constraints of the host machine (CPU/RAM/GPU). Rust is chosen for the engine layer for three primary reasons:

2.1 Memory Safety and Security Boundaries

AI models often process unstructured text, which can be adversarial. Malformed inputs can lead to buffer overflows in C/C++ libraries. Rust’s ownership model guarantees memory safety at compile time. For a local-first agent, this is a security feature, not just a performance one. If the inference engine crashes due to a memory error, it takes down the entire local service. Rust prevents this.

2.2 Predictable Latency

WebAssembly (Wasm) and real-time systems require deterministic behavior. Rust’s lack of garbage collection pauses ensures that token generation latency remains consistent, which is crucial for chat interfaces that expect real-time streaming responses.

2.3 Interoperability with ML Libraries

The modern Rust ML ecosystem, including Burn, Candle (by Hugging Face), and tch-rs (PyTorch bindings), provides access to state-of-the-art models. These libraries are optimized for SIMD instructions and multi-threaded matrix multiplication, squeezing every bit of performance out of consumer hardware.

3. Why Go for the Orchestrator?

While Rust is superior for raw computation, Go excels at concurrency patterns and ecosystem integration. An AI agent is rarely just a model; it is a system that:

  • Manages multiple conversation sessions.
  • Calls external tools (APIs, databases, file systems).
  • Handles authentication and authorization.
  • Persists state to a database.

Go’s goroutine model is perfectly suited for this. Each user session can be assigned a dedicated goroutine, allowing the system to handle thousands of concurrent users with a small memory footprint. Furthermore, Go’s standard library provides excellent tools for building HTTP/2 servers, handling WebSocket streams for real-time token delivery, and interacting with SQL/NoSQL databases.

4. System Architecture and Data Flow

Let’s visualize the data flow in a typical request:

  1. Client: User sends a message via WebSocket.
  2. Go Orchestrator: Receives the message, validates it, retrieves conversation history from a SQLite/PostgreSQL store, and prepares the context.
  3. Bridge: Go serializes the context (JSON/Protobuf) and sends it to the Rust engine via a local gRPC channel or FFI call.
  4. Rust Engine:
    • Loads the context into the model’s embedding space.
    • Performs inference (token generation).
    • Streams tokens back to Go.
  5. Go Orchestrator: Receives the stream, formats the output, and pushes it to the client via WebSocket.
  6. Persistence: Go saves the new exchange to the database.

4.1 The Bridge: Minimizing Overhead

The cost of inter-process communication (IPC) or Foreign Function Interface (FFI) calls can be significant if data is copied excessively. To mitigate this, we use zero-copy strategies where possible.

For gRPC, we define a schema that minimizes payload size. For FFI (Go calling Rust), we use unsafe blocks carefully to pass pointers to pre-allocated buffers, avoiding the serialization/deserialization overhead of JSON.

5. Implementation: The Rust Inference Core

We will use Candle (by Hugging Face) for its simplicity and performance in local inference. Candle is designed to be embeddable and works well on CPU and GPU.

5.1 Project Structure 

rust-engine/
├── Cargo.toml
├── src/
│   ├── main.rs          # Entry point, gRPC server
│   ├── inference.rs     # Model loading and execution logic
│   ├── tokenizer.rs     # Tokenization utilities
│   └── lib.rs           # Exports for FFI/gRPC
└── proto/
    └── agent.proto      # gRPC definitions

Enter fullscreen mode Exit fullscreen mode

5.2 Defining the gRPC Service

First, we define the contract between Go and Rust. This ensures type safety and clear expectations.

// proto/agent.proto
syntax = "proto3";

package agent;

service InferenceService {
  // Stream tokens to the orchestrator
  rpc GenerateStream (Request) returns (stream Response);

  // Health check and model status
  rpc GetModelStatus (Empty) returns (Status);
}

message Request {
  string prompt = 1;
  int32 max_tokens = 2;
  float temperature = 3;
  repeated string history = 4; // Previous conversation turns
}

message Response {
  string token = 1;
  bool is_end = 2;
}

message Status {
  string model_name = 1;
  bool is_ready = 2;
}

message Empty {}

Enter fullscreen mode Exit fullscreen mode

5.3 Implementing the Rust Server

We use tonic for gRPC and candle for inference. This setup allows the Rust binary to act as a standalone service that Go clients connect to.

// src/inference.rs
use candle::{Device, Tensor, DType};
use candle_transformers::generation::LogitsProcessor;
use candle_transformers::models::llama as llama_model;
use std::sync::Arc;

pub struct LlamaEngine {
    model: Arc<llama_model::Llama>,
    device: Device,
    tokenizer: tokenizers::Tokenizer,
}

impl LlamaEngine {
    pub fn new(model_path: &str, tokenizer_path: &str, device: Device) -> Result<Self, Box<dyn std::error::Error>> {
        // Load tokenizer
        let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)?;

        // Load model weights
        let model = llama_model::Llama::from_preferred(model_path, &device)?;

        Ok(LlamaEngine {
            model: Arc::new(model),
            device,
            tokenizer,
        })
    }

    pub fn generate(&self, prompt: &str, max_tokens: usize, temperature: f64) -> Result<String, Box<dyn std::error::Error>> {
        // Tokenize input
        let tokens = self.tokenizer.encode(prompt, true)?;
        let tokens = tokens.get_ids();

        // Prepare input tensor
        let input = Tensor::new(tokens, &self.device)?;

        // Initialize logits processor
        let mut logits_processor = LogitsProcessor::new(temperature, Some(42), None);

        let mut output_tokens = Vec::new();
        let mut current_tokens = input.clone();

        // Inference loop
        for _ in 0..max_tokens {
            let logits = self.model.forward(&current_tokens, 0)?;
            let logits = logits.squeeze(0)?;

            let next_token = logits_processor.sample(&logits)?;
            output_tokens.push(next_token);

            // Break if end of sequence
            if next_token == self.tokenizer.tokenizer.eos_token_id() {
                break;
            }

            // Prepare next input
            current_tokens = Tensor::new(&[next_token], &self.device)?;
        }

        // Decode output tokens
        let decoded = self.tokenizer.decode(&output_tokens, true)?;
        Ok(decoded)
    }
}

Enter fullscreen mode Exit fullscreen mode

5.4 Integrating with gRPC 

// src/main.rs
use tonic::{transport::Server, Request, Response, Status};
use tokio_stream::{Stream, StreamExt};

mod inference;
mod agent {
    include!(concat!(env!("OUT_DIR"), "/agent.rs"));
}

struct InferenceServiceImpl {
    engine: inference::LlamaEngine,
}

#[tonic::async_trait]
impl agent::inference_service_server::InferenceService for InferenceServiceImpl {
    async fn generate_stream(
        &self,
        request: Request<agent::Request>,
    ) -> Result<Response<Self::GenerateStreamStream>, Status> {
        let req = request.into_inner();
        let prompt = req.prompt;
        let max_tokens = req.max_tokens as usize;
        let temperature = req.temperature;

        // Generate tokens synchronously for simplicity, 
        // but in production, you'd stream intermediate tokens
        let result = match self.engine.generate(&prompt, max_tokens, temperature) {
            Ok(text) => text,
            Err(e) => return Err(Status::internal(e.to_string())),
        };

        // For streaming, we'd split the result and yield chunks
        // Here we return a single response for demonstration
        Ok(Response::new(agent::Response {
            token: result,
            is_end: true,
        }))
    }

    async fn get_model_status(
        &self,
        _request: Request<agent::Empty>,
    ) -> Result<Response<agent::Status>, Status> {
        Ok(Response::new(agent::Status {
            model_name: "Llama-2-7b".to_string(),
            is_ready: true,
        }))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let engine = inference::LlamaEngine::new(
        "./models/llama-2-7b",
        "./models/tokenizer.json",
        Device::Cpu,
    )?;

    let svc = agent::inference_service_server::InferenceServiceServer::new(
        InferenceServiceImpl { engine }
    );

    Server::builder()
        .add_service(svc)
        .serve("127.0.0.1:50051".parse()?)
        .await?;

    Ok(())
}

Enter fullscreen mode Exit fullscreen mode

6. The Go Orchestrator: Concurrency and Session Management

The Go application acts as the brain, managing the state of the conversation and coordinating with the Rust engine.

6.1 Session Store

We need a way to manage multiple users. A simple in-memory map is sufficient for demonstration, but in production, you’d use Redis or PostgreSQL.

// go-orchestrator/internal/session/store.go
package session

import (
    "sync"
)

type Session struct {
    ID      string
    History []string
}

type Store struct {
    mu    sync.RWMutex
    sessions map[string]*Session
}

func NewStore() *Store {
    return &Store{
        sessions: make(map[string]*Session),
    }
}

func (s *Store) GetOrCreate(id string) *Session {
    s.mu.RLock()
    sess, ok := s.sessions[id]
    s.mu.RUnlock()

    if !ok {
        s.mu.Lock()
        // Double-check locking
        if sess, ok = s.sessions[id]; !ok {
            sess = &Session{ID: id, History: make([]string, 0)}
            s.sessions[id] = sess
        }
        s.mu.Unlock()
    }
    return sess
}

Enter fullscreen mode Exit fullscreen mode

6.2 gRPC Client for Rust Engine

The Go app needs to connect to the Rust gRPC server. We use the generated protobuf client.

// go-orchestrator/internal/ai/engine.go
package ai

import (
    "context"
    "log"
    "github.com/yourname/agent/proto/agent"
    "google.golang.org/grpc"
)

type RustEngineClient struct {
    client agent.InferenceServiceClient
}

func NewRustEngineClient(addr string) (*RustEngineClient, error) {
    conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    if err != nil {
        return nil, err
    }
    return &RustEngineClient{
        client: agent.NewInferenceServiceClient(conn),
    }, nil
}

func (c *RustEngineClient) Generate(ctx context.Context, prompt string, history []string, maxTokens int, temperature float32) (string, error) {
    resp, err := c.client.GenerateStream(ctx, &agent.Request{
        Prompt:      prompt,
        MaxTokens:   int32(maxTokens),
        Temperature: temperature,
        History:     history,
    })
    if err != nil {
        return "", err
    }
    return resp.Token, nil
}

Enter fullscreen mode Exit fullscreen mode

6.3 WebSocket Handler with Concurrency

Go’s strength is handling many concurrent connections. We use a WebSocket server to stream responses back to the client.

// go-orchestrator/internal/handler/ai_handler.go
package handler

import (
    "context"
    "log"
    "net/http"
    "time"

    "github.com/gorilla/websocket"
    "github.com/yourname/agent/go-orchestrator/internal/ai"
    "github.com/yourname/agent/go-orchestrator/internal/session"
)

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

func AIHandler(engine *ai.RustEngineClient, store *session.Store) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        conn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
            log.Println("Upgrade error:", err)
            return
        }
        defer conn.Close()

        // Generate a unique session ID for this connection
        sessionID := "user-" + fmt.Sprintf("%d", time.Now().UnixNano())
        sess := store.GetOrCreate(sessionID)

        ctx := context.Background()

        for {
            _, msg, err := conn.ReadMessage()
            if err != nil {
                break
            }

            // Process message
            response, err := engine.Generate(ctx, string(msg), sess.History, 500, 0.7)
            if err != nil {
                conn.WriteMessage(websocket.TextMessage, []byte("Error: "+err.Error()))
                continue
            }

            // Send response
            conn.WriteMessage(websocket.TextMessage, []byte(response))

            // Update history
            sess.History = append(sess.History, string(msg), response)
            // Limit history size to manage memory
            if len(sess.History) > 20 {
                sess.History = sess.History[len(sess.History)-20:]
            }
        }
    }
}

Enter fullscreen mode Exit fullscreen mode

7. Security Considerations in Local-First AI

Building a local-first AI assistant introduces unique security challenges. Unlike cloud APIs, where you can implement rate limiting and WAFs, your local application is the first line of defense.

7.1 Input Sanitization and Prompt Injection

Even locally, users can attempt prompt injection. The Go orchestrator should sanitize inputs before sending them to the Rust engine. Implement a pre-processing layer that filters out malicious patterns or restricts the context window to prevent buffer overflow attempts in the Rust layer.

7.2 Memory Isolation

If the Rust engine is compiled as a shared library (.so or .dylib) and linked into the Go binary, a vulnerability in Rust could compromise the entire process. To mitigate this, consider running the Rust engine as a separate process (as shown in the gRPC example) and communicating via IPC. This way, a crash in the Rust engine does not take down the Go orchestrator, and you can implement stricter OS-level sandboxing for the Rust process.

7.3 Data Persistence

Local-first implies data stays on the device. Ensure that session histories and any sensitive data processed by the AI are encrypted at rest. Use Go’s crypto package to encrypt the session store before writing to disk.

8. Performance Optimization: Batching and Quantization

Local inference is constrained by hardware. To make the assistant responsive, you must optimize the model.

8.1 Quantization

Large Language Models are often stored in FP16 or BF16. Quantizing to INT8 or even INT4 can reduce memory usage by 2-4x with minimal loss in quality. The Candle library supports quantized models out of the box. Load quantized models to fit larger context windows on consumer hardware.

8.2 Batching

If you have multiple users, the Rust engine can batch requests. However, for a local-first single-user assistant, batching is less relevant. For multi-user scenarios, the Rust engine can maintain a batch queue, processing multiple prompts in parallel using multi-threaded matrix multiplication.

9. Conclusion

Building a secure, local-first AI assistant is a complex but rewarding engineering challenge. By leveraging Go for its concurrency and ecosystem, and Rust for its performance and memory safety, you can create a system that is both flexible and robust.

This architecture provides:

  • Privacy: Data never leaves the device.
  • Performance: Low-latency inference via optimized Rust code.
  • Security: Memory-safe inference and isolated processes.
  • Scalability: Go’s ability to handle many concurrent sessions.

As AI moves from a cloud-centric model to a distributed, edge-centric one, this hybrid approach will become the standard for high-assurance applications. Start by prototyping the Rust inference engine, then build the Go orchestrator around it, and gradually add security and optimization layers.

Frequently Asked Questions

Q: Can I use Python for the orchestration layer?
A: Yes, Python is often used for AI development. However, for a high-concurrency local service, Go or Rust is superior. Python’s GIL limits true parallelism, making it less ideal for handling many simultaneous WebSocket connections compared to Go’s goroutines.

Q: How do I handle model updates in a local-first app?
A: Implement a versioning system in your gRPC protocol. When a new model is released, the Go orchestrator can trigger a download of the new weights (quantized) and reload the Rust engine process without restarting the entire application.

Q: Is this approach viable for mobile devices?
A: Yes. The same architecture can be adapted for iOS/Android using Go (via gomobile) or Rust (via Swift/Kotlin bindings). The key is to keep the inference engine (Rust) native and the logic (Go) as a thin layer or also native.