Learn RustChapter 10: Streaming replies

Lesson 10.4 · 20 min

Stopping a long answer

Answers now stream in word by word, and some of them go on for a long time. Today the only way to stop one is Ctrl-C, and Ctrl-C ends the whole program: the conversation is never saved, because saving happens after the loop. This lesson gives Chatterbox a real stop button. It is the last lesson of the course.

Adds to ChatterboxCtrl-C stops an answer, or saves and quits at the prompt. Version 1.1.

Rust async, await, tokio and Streaming: Live Replies in Your Chat App (Learn Rust, Ep 10)

Whichever comes first

While an answer streams in, Chatterbox is waiting for one thing: the next piece. To stop it, it has to wait for two things at once, the next piece or you pressing Ctrl-C, and do whatever the first one to happen says.

That is what tokio::select! does. Like println!, it is a macro (the ! says so), and its arms look like a match: pattern = future => code. It waits on all the futures at once. When the first one finishes, its value is matched against the pattern and that arm's code runs. Every other future is dropped: thrown away where it stood. A future is a recipe that only does something while it is awaited (lesson 10.1), so a dropped one simply never takes another step.

PredictA race with a slow answer

use std::time::Duration;

use tokio::time::sleep;

// Types one word every 200 ms, like a slow model.

async fn long_answer() -> &'static str {

for word in ["Once", "upon", "a", "time"] {

println!("bot types: {word}");

sleep(Duration::from_millis(200)).await;

}

"the end"

}

#[tokio::main]

async fn main() {

tokio::select! {

end = long_answer() => println!("{end}"),

_ = sleep(Duration::from_millis(500)) => println!("stopped"),

}

// Wait a whole second more: does the long answer carry on?

sleep(Duration::from_millis(1000)).await;

println!("main ends");

}

What does this program print?

bot types: Once bot types: upon bot types: a stopped main ends
bot types: Once bot types: upon bot types: a stopped bot types: time main ends
bot types: Once bot types: upon bot types: a bot types: time the end main ends

Reveal the answer

The program prints: bot types: Once bot types: upon bot types: a stopped main ends The words come at 0, 200 and 400 ms. At 500 ms the short sleep finishes first, so its arm prints stopped and long_answer is dropped in the middle of its own sleep. A dropped future is gone: during the extra second nothing types time, and the end is never reached. _ as the pattern means "I do not need the value", the same as in a match.

Ctrl-C as an event

Pressing Ctrl-C does not type anything. The terminal sends the program a signal, a short message from the operating system (this one is called SIGINT, "interrupt"). A program that does nothing about it is ended on the spot, which is what has happened to Chatterbox so far.

tokio::signal::ctrl_c() turns that signal into a future: it finishes when Ctrl-C is pressed. A future can go in a select!. So in answer(), the while let loop from lesson 10.2 becomes a loop around a select! with two arms:

piece = pieces.recv(): the next piece arrived. The match inside does what the old loop did: print a piece, report an error, or break on None because the channel closed and the answer is complete.
_ = tokio::signal::ctrl_c(): you pressed Ctrl-C. Print [stopped] and break.

Either way the function returns Ok(Some(whole)), so the part of the answer you saw is added to the history, as if the bot had stopped mid-sentence.

Going deeper Does the losing arm lose a piece?

The select! runs once per piece, and each time the arm that loses is dropped. For recv() that loses nothing: a piece that was on its way stays in the channel for the next recv(). tokio's documentation calls a future like that cancel safe. Not every future is: one that has read half of something before it is dropped can lose that half, which is why tokio's documentation says, method by method, whether each one is cancel safe.

Who stops the task?

When you stop an answer, a task is still running somewhere: the pretend model's typing loop, or the real model's download. Nothing in answer() tells it to stop, and nothing needs to. When answer() returns, its local pieces (the receiving end of the channel) is dropped, and a channel with no receiver is closed. The task finds out the next time it sends:

Sending into a channel whose receiver is gone (a real run)

use tokio::sync::mpsc;

#[tokio::main]

async fn main() {

let (tx, pieces) = mpsc::channel(8);

drop(pieces); // what returning from answer() does to it

match tx.send(String::from("Once")).await {

Ok(()) => println!("sent"),

Err(_) => println!("nobody is listening"),

}

$ cargo run
nobody is listening

That Err is the if tx.send(...).await.is_err() check you wrote in lessons 10.2 and 10.3: the task returns, and everything it owned goes with it, including the real model's connection to the server. No stop flag, no message saying "please stop". One owner let go, and ownership did the clean-up.

Ctrl-C at the prompt

There is a catch. The first time a program calls ctrl_c(), tokio replaces what the operating system does on Ctrl-C, for the rest of the program's life. From then on Ctrl-C never ends Chatterbox by itself. If only answer() listened, Ctrl-C at the prompt would do nothing at all.

So the main loop gets a select! too: the next line from the keyboard, or Ctrl-C, which is turned into the line "/quit". Everything after that already knows what /quit means: say goodbye, leave the loop, save. A key that used to throw your conversation away now saves it.

The keyboard gets a thread

Lesson 10.2 read the keyboard with tokio's own stdin, and its documentation warns that it can keep the program from ending until Enter is pressed. Once Ctrl-C can quit from the prompt, that is exactly the moment it would bite: you press Ctrl-C, Chatterbox says goodbye, and then sits there.

The fix is a thread: a second line of work that the operating system runs alongside the main one, each with its own place to wait. std::thread::spawn starts one and gives it a closure to run (lesson 6.4); move in front hands the closure what it uses, here tx, the same as async move did for tasks. The new function keyboard() starts a thread that reads lines the plain, blocking way, as Chatterbox did before chapter 10, and sends each one down a channel, and returns the receiving end. The main loop then waits on keys.recv(), which is a future, so it fits in select!.

Inside the thread, the send is tx.blocking_send(line): a send that just waits, the plain way, until there is room. The thread is not async, so it cannot use .await. Try it and the compiler says so:

Compiler CoachThe compiler, translated

This is the exact output of the real compiler. Toggle the English, then pick a fix and see what actually happens.

use std::io;

use std::thread;

use tokio::sync::mpsc;

fn keyboard() -> mpsc::Receiver<String> {

let (tx, rx) = mpsc::channel(16);

thread::spawn(move || {

for line in io::stdin().lines() {

let line = line.unwrap();

if tx.send(line).await.is_err() {

break;

}

});

}

#[tokio::main]

async fn main() {

let mut keys = keyboard();

while let Some(line) = keys.recv().await {

println!("you typed: {line}");

}

error[E0728]: `await` is only allowed inside `async` functions and blocks
  --> src/main.rs:10:30
   |
 7 |     thread::spawn(move || {
   |                   ------- this is not `async`
...
10 |             if tx.send(line).await.is_err() {
   |                              ^^^^^ only allowed inside `async` functions and blocks

E0728 in the compiler error atlas →

.await only works inside an async fn or an async block, because only async code can pause and let something else run.
The closure given to thread::spawn is ordinary code that runs on its own thread. It is not async, so it has nowhere to pause to.
The thread does not need to pause politely: it has nothing else to do while it waits. It can simply wait.
tokio gives a channel's sender a method for exactly this, blocking_send: it waits the plain way until there is room.

Fixes that were compiled for real

Use blocking_send on the thread — compiles. The thread waits the plain way until the channel has room. This is what Chatterbox v1.1 does.
Trap: drop the move — still E0373. Without <code>move</code> the closure only borrows <code>tx</code>, but the thread may run on after <code>keyboard()</code> has returned and <code>tx</code> is gone. The same rule as for tasks in lesson 10.2.

Going deeper Why a thread of your own, and why blocking_send stays there

tokio's stdin does its reading with an ordinary blocking read on a helper thread of tokio's, and that read cannot be called off. When main returns, shutting the runtime down can hang until that read finishes, which is when you press Enter. A thread started with std::thread::spawn is not the runtime's: when main returns, the program ends, and a thread still waiting for Enter ends with it.

blocking_send is the right send on that thread and the wrong one anywhere else. Its documentation says it panics if it is called from async code. That is lesson 10.1's one rule, enforced: waiting the plain way there would hold up every other task on the same thread.

Try itYour turn: a stop button a test can press

A test cannot press Ctrl-C, so here the stop button is a second channel. Its messages are (), the empty value: the press itself is the news. Write collect, the heart of answer(): keep adding pieces to whole until the channel closes (return the text and false) or a press arrives on stop (return the text so far and true). In main, a pretend model types a word every 100 ms and someone presses stop after 250 ms. In the tests, a sender named _press or _tx is kept alive but never used (a name that starts with _ still owns its value until the end of the test). That matters: if every Sender of stop were dropped, stop.recv() would give None at once, select! would treat that arm as finished, and the stop branch would fire although nobody pressed.

use std::time::Duration;

use tokio::sync::mpsc;

/// Collect an answer's pieces until the channel closes (the answer is

/// complete) or the stop button is pressed. Give back the text so far,

/// and whether it was stopped.

async fn collect(mut pieces: mpsc::Receiver<String>, mut stop: mpsc::Receiver<()>) -> (String, bool) {

let mut whole = String::new();

// wait for whichever comes first, again and again

(whole, false)

}

#[tokio::main]

async fn main() {

let (tx, pieces) = mpsc::channel(8);

let (press, stop) = mpsc::channel(1);

// A slow model: one word every 100 ms.

tokio::spawn(async move {

for word in ["Once ", "upon ", "a ", "time."] {

if tx.send(String::from(word)).await.is_err() {

return; // nobody is listening any more

}

tokio::time::sleep(Duration::from_millis(100)).await;

}

});

// Someone presses the stop button after 250 ms.

tokio::spawn(async move {

tokio::time::sleep(Duration::from_millis(250)).await;

let _ = press.send(()).await;

});

let (text, stopped) = collect(pieces, stop).await;

println!("{text:?} (stopped: {stopped})");

}

Hidden tests (4): a_whole_answer_is_kept an_empty_answer_is_not_a_stop stop_before_anything_arrives stopping_keeps_what_was_shown

This exercise runs real Rust in your browser through the official Rust Playground, so it needs JavaScript.

Chatterbox v1.1

Here is the whole program, with this lesson's changes in green: the two select! loops, keyboard() and its thread, and two new lines in the README. In Cargo.toml tokio's features change: io-std and io-util (tokio's keyboard reading) go, and signal (for ctrl_c) comes in.

A test script cannot press keys, but it can send the same signal Ctrl-C sends. In this run, the build started Chatterbox against a stand-in server that streams slowly, typed the lines under "You type", waited until the answer had begun, and sent SIGINT: that is the "(presses Ctrl-C)" line. The answer stops, the chat carries on, and /history shows the part of the answer you saw, kept.

Chatterbox v1.1 after lesson 10.4: the whole program (green lines are new in this lesson). The run sent SIGINT while the answer was streaming.

src/main.rs51 new lines

// Chatterbox: a small chat program, built one lesson at a time.

mod command;

mod message;

mod replier;

mod store;

use command::{Command, parse};

use message::Message;

use replier::{Echo, Model, Pretend, Replier};

use std::collections::HashMap;

use std::io::{self, Write};

use std::thread;

use tokio::sync::mpsc;

/// What the bot is called. A constant: fixed for the whole program.

const BOT: &str = "Chatterbox";

/// The version as a pair of numbers: (major, minor).

const VERSION: (u32, u32) = (1, 1);

/// The longest message we accept, in bytes.

const MAX_LEN: usize = 200;

/// What /help prints.

const HELP: &str = "commands: /help, /again, /history, /count, /stats, /remember <name> <fact>, /recall <name>, /find <word>, /last <n>, /model <echo|pretend|real>, /say <text>, /shout <text>, /debug <text>, /quit (or /h, /q, /exit)";

/// The first line Chatterbox prints.

fn banner() -> String {

format!("{} v{}.{}", BOT, VERSION.0, VERSION.1)

}

/// The greeting, with a name filled in.

fn greeting(name: &str) -> String {

let mut text = format!("Hi {name}! I'm {BOT}.");

text.push_str(" Ask me anything.");

text

}

/// The last `n` items of any slice: messages, numbers, words, anything.

/// `<T>` says "for any type T"; the body never needs to know which one.

fn last_n<T>(items: &[T], n: usize) -> &[T] {

let start = items.len().saturating_sub(n);

&items[start..]

}

/// Print an answer as its pieces arrive, and give back the whole of it

/// (None if it went wrong: what went wrong has been printed instead).

async fn answer(replier: &dyn Replier, history: &[Message]) -> Result<Option<String>, io::Error> {

let mut pieces = replier.reply(history);

let mut whole = String::new();

print!("bot: ");

loop {

// Wait for whichever comes first: the next piece, or Ctrl-C.

tokio::select! {

piece = pieces.recv() => match piece {

Some(Ok(text)) => {

print!("{text}");

io::stdout().flush()?; // show it now, not at the end of the line

whole.push_str(&text);

}

Some(Err(e)) => {

println!("(that went wrong: {e})");

return Ok(None);

}

None => break, // the channel closed: the answer is complete

_ = tokio::signal::ctrl_c() => {

// Returning drops `pieces`; the task sending them sees that and stops.

print!(" [stopped]");

break;

}

println!();

Ok(Some(whole))

}

/// The keyboard gets a thread of its own. Reading a line blocks until Enter,

/// and a blocked read cannot be called off, so it must not hold anything else

/// up. Each line is sent on the channel; the channel closes at end of input.

fn keyboard() -> mpsc::Receiver<String> {

let (tx, rx) = mpsc::channel(16);

thread::spawn(move || {

for line in io::stdin().lines() {

let line = match line {

Ok(line) => line,

Err(_) => break, // the keyboard could not be read

};

if tx.blocking_send(line).is_err() {

break; // the program has stopped listening

}

});

}

// #[tokio::main] starts tokio's runtime, the part that runs async code, and

// then runs this async main inside it.

#[tokio::main]

async fn main() -> Result<(), io::Error> {

let mut messages_seen: u32 = 0;

let quiet = false;

let mut last: Option<String> = None; // the last thing you said, for /again

let mut history: Vec<Message> = store::load(store::FILE)?; // last time's conversation, if any

let mut facts: HashMap<String, String> = HashMap::new(); // name -> fact, for /recall

// Box<dyn Replier>: "some Replier, decided while the program runs".

// With a key in the environment, start with the real model.

let mut replier: Box<dyn Replier> = match Model::from_env() {

Ok(model) => Box::new(model),

Err(_) => Box::new(Echo { max_len: MAX_LEN }),

};

println!("{} (replies by: {})", banner(), replier.name());

let mut keys = keyboard();

println!("What should I call you?");

let name_line = keys.recv().await.unwrap_or_default();

let name = name_line.trim();

if !quiet {

println!("{}", greeting(name));

}

if !history.is_empty() {

println!(

"(I remember {} messages from last time. /history shows them.)",

history.len()

);

}

loop {

// A line from the keyboard, or Ctrl-C, which at the prompt means /quit.

let line = tokio::select! {

line = keys.recv() => line,

_ = tokio::signal::ctrl_c() => Some(String::from("/quit")),

};

let line = match line {

Some(line) => line,

None => break, // no more input: the keyboard was closed

};

let message = line.trim();

let cmd = match parse(message) {

Ok(cmd) => cmd,

Err(e) => {

println!("bot: {e}");

continue;

}

};

match cmd {

Command::Empty => {} // nothing typed: nothing to answer

Command::Quit => {

println!("bot: bye, {name}!");

break;

}

Command::Help => println!("bot: {HELP}"),

Command::Again => match &last {

Some(text) => {

answer(replier.as_ref(), &[Message::new("you", text)]).await?;

}

None => println!("bot: nothing to repeat yet"),

Command::History => {

if history.is_empty() {

println!("bot: nothing said yet");

} else {

for (i, m) in history.iter().enumerate() {

println!("{}. {}", i + 1, m.to_line());

}

Command::Count => {

let bytes: usize = history.iter().map(|m| m.text.len()).sum();

println!("bot: {} messages, {} bytes of text", history.len(), bytes);

}

Command::Stats => {

// One chain of steps: your messages -> their words -> a count per word.

let mut counts: HashMap<String, u32> = HashMap::new();

for word in history

.iter()

.filter(|m| m.author == "you")

.flat_map(|m| m.text.split_whitespace())

{

*counts.entry(word.to_lowercase()).or_insert(0) += 1;

}

let mut top: Vec<(String, u32)> = counts.into_iter().collect();

// most used first; ties in alphabetical order, so the answer never changes

top.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));

let words: Vec<String> = top

.iter()

.take(3)

.map(|(word, n)| format!("{word} ({n})"))

.collect();

if words.is_empty() {

println!("bot: nothing said yet");

} else {

println!("bot: your top words: {}", words.join(", "));

}

Command::Remember(key, fact) => {

println!("bot: got it: {key} is {fact}");

facts.insert(key, fact);

}

Command::Recall(Some(key)) => match facts.get(&key) {

Some(fact) => println!("bot: {key} is {fact}"),

None => println!("bot: I don't know anything about {key}"),

Command::Recall(None) => {

if facts.is_empty() {

println!("bot: I don't remember anything yet. Try /remember");

} else {

// A HashMap has no order, so sort the names before printing them.

let mut names = Vec::new();

for name in facts.keys() {

names.push(name.as_str());

}

names.sort();

println!("bot: I remember: {}", names.join(", "));

}

Command::Find(word) => {

// The closure borrows `word` from here: a small function made on the spot.

let hits: Vec<&Message> = history

.iter()

.filter(|m| m.author == "you" && m.text.contains(&word))

.collect();

if hits.is_empty() {

println!("bot: you never said {word}");

} else {

for m in hits {

println!("{}", m.to_line());

}

Command::Last(n) => {

for m in last_n(&history, n) {

println!("{}", m.to_line());

}

Command::Model(choice) => {

let next: Result<Box<dyn Replier>, String> = match choice.as_str() {

"echo" => Ok(Box::new(Echo { max_len: MAX_LEN })),

"pretend" => Ok(Box::new(Pretend)),

"real" => match Model::from_env() {

Ok(model) => Ok(Box::new(model)),

Err(e) => Err(e.to_string()),

_ => Err(format!(

"I don't know a model called {choice}. Try echo, pretend or real"

)),

};

match next {

Ok(r) => {

replier = r;

println!("bot: replies now come from {}", replier.name());

}

Err(why) => println!("bot: {why}"),

}

Command::Say(text) => println!("bot: {text}"),

Command::Shout(text) => println!("bot: {}", Message::new("you", &text).shout()),

Command::Debug(text) => println!("bot: {:?}", parse(&text)),

Command::Chat(text) => {

messages_seen += 1;

history.push(Message::new("you", &text));

// the history holds its own copy, so `text` can move into `last`

last = Some(text);

if let Some(reply) = answer(replier.as_ref(), &history).await? {

history.push(Message::new("bot", &reply));

}

println!("messages so far: {messages_seen}");

store::save(store::FILE, &history)?;

let s = if history.len() == 1 { "" } else { "s" };

println!("saved {} message{s} to {}", history.len(), store::FILE);

Ok(())

}

// Tests: `cargo test` builds and runs them; the program itself never sees them.

#[cfg(test)]

mod tests {

use super::*;

#[test]

fn last_n_takes_from_the_end() {

assert_eq!(last_n(&[1, 2, 3, 4], 2), &[3, 4]);

assert_eq!(last_n(&["a", "b"], 5), &["a", "b"]);

assert!(last_n(&[1, 2, 3], 0).is_empty());

}

src/command.rsunchanged

//! Commands: what a line of input means, and what can be wrong with it.

use std::fmt;

/// Everything a line of input can mean. One enum, so a `match` on it must

/// handle every kind of line: the compiler checks that nothing is forgotten.

#[derive(Debug)]

pub enum Command {

Empty,

Help,

Quit,

Again,

History,

Count,

Stats,

Remember(String, String),

Recall(Option<String>),

Find(String),

Last(usize),

Model(String),

Say(String),

Shout(String),

Debug(String),

Chat(String),

}

/// Everything that can be wrong with a line of input. parse() returns one of

/// these instead of a Command, and main prints it.

#[derive(Debug)]

pub enum CommandError {

Unknown(String),

Missing(String, String),

Usage(String),

NotACount(String),

}

/// How a CommandError reads when printed with {}.

impl fmt::Display for CommandError {

fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

match self {

CommandError::Unknown(word) => {

write!(f, "I don't know the command {word}. Try /help.")

}

CommandError::Missing(verb, example) => {

write!(f, "{verb} what? Try /{verb} {example}")

}

CommandError::Usage(form) => write!(f, "try {form}"),

CommandError::NotACount(text) => write!(f, "'{text}' is not a count. Try /last 3"),

}

/// The text after a command word, or the error that says it is missing.

fn need(line: &str, verb: &str, example: &str) -> Result<String, CommandError> {

match arg(line) {

Some(text) => Ok(String::from(text)),

None => Err(CommandError::Missing(

String::from(verb),

String::from(example),

)),

}

/// How many messages /last should show. "two" or "-1" is an error.

fn parse_count(text: &str) -> Result<usize, CommandError> {

match text.parse::<usize>() {

Ok(n) => Ok(n),

Err(_) => Err(CommandError::NotACount(String::from(text))),

}

/// Turn one line of input into a Command, or say what is wrong with it.

/// Each `?` means: if this step failed, return its error right here.

pub fn parse(line: &str) -> Result<Command, CommandError> {

let cmd = match first_word(line) {

"" => Command::Empty,

"/help" | "/h" => Command::Help,

"/quit" | "/q" | "/exit" => Command::Quit,

"/again" => Command::Again,

"/history" => Command::History,

"/count" => Command::Count,

"/stats" => Command::Stats,

"/remember" => {

let usage = CommandError::Usage(String::from("/remember <name> <fact>"));

let rest = match arg(line) {

Some(rest) => rest,

None => return Err(usage),

};

match arg(rest) {

Some(fact) => Command::Remember(String::from(first_word(rest)), String::from(fact)),

None => return Err(usage),

}

"/recall" => match arg(line) {

Some(name) => Command::Recall(Some(String::from(name))),

None => Command::Recall(None),

"/find" => Command::Find(need(line, "find", "hello")?),

"/last" => Command::Last(parse_count(&need(line, "last", "3")?)?),

"/model" => Command::Model(need(line, "model", "pretend")?),

"/debug" => Command::Debug(need(line, "debug", "/say hi")?),

"/say" => Command::Say(need(line, "say", "hello")?),

"/shout" => Command::Shout(need(line, "shout", "hello")?),

word => {

if is_command(word) {

return Err(CommandError::Unknown(String::from(word)));

}

Command::Chat(String::from(line))

}

};

Ok(cmd)

}

/// True when the text is a command: it starts with a slash.

fn is_command(text: &str) -> bool {

text.starts_with('/')

}

/// The first word of the text: a slice (&str) pointing into it, not a copy.

fn first_word(text: &str) -> &str {

match text.find(' ') {

Some(i) => &text[..i],

None => text,

}

/// The text after the command word, if there is any: Some("hello") for

/// "/say hello", None for a bare "/say". No more "" standing in for nothing.

fn arg(line: &str) -> Option<&str> {

match line.find(' ') {

Some(i) => {

let rest = line[i + 1..].trim();

if rest.is_empty() { None } else { Some(rest) }

}

None => None,

}

#[cfg(test)]

mod tests {

use super::*;

#[test]

fn plain_text_is_chat() {

assert!(matches!(parse("hello there"), Ok(Command::Chat(t)) if t == "hello there"));

}

#[test]

fn help_has_a_short_form() {

assert!(matches!(parse("/help"), Ok(Command::Help)));

assert!(matches!(parse("/h"), Ok(Command::Help)));

}

#[test]

fn say_needs_text() {

assert_eq!(

parse("/say").unwrap_err().to_string(),

"say what? Try /say hello"

);

}

#[test]

fn last_needs_a_number() {

assert!(matches!(parse("/last 3"), Ok(Command::Last(3))));

assert_eq!(

parse("/last two").unwrap_err().to_string(),

"'two' is not a count. Try /last 3"

);

}

#[test]

fn unknown_commands_are_errors() {

assert_eq!(

parse("/dance").unwrap_err().to_string(),

"I don't know the command /dance. Try /help."

);

}

#[test]

fn remember_takes_a_name_and_a_fact() {

assert!(

matches!(parse("/remember pizza extra cheese"), Ok(Command::Remember(k, v)) if k == "pizza" && v == "extra cheese")

);

assert!(parse("/remember pizza").is_err());

}

src/message.rsunchanged

//! A Message: one line of the conversation.

use serde::{Deserialize, Serialize};

/// One line of the conversation: who said it, and what they said.

#[derive(Serialize, Deserialize)]

pub struct Message {

pub author: String,

pub text: String,

}

impl Message {

/// Make a Message from two pieces of borrowed text.

pub fn new(author: &str, text: &str) -> Message {

Message {

author: String::from(author),

text: String::from(text),

}

/// The message as one transcript line: "you: hello".

pub fn to_line(&self) -> String {

format!("{}: {}", self.author, self.text)

}

/// The text in capitals, for /shout.

pub fn shout(&self) -> String {

self.text.to_uppercase()

}

src/replier.rsunchanged

//! Who answers: anything that can reply, behind one trait.

use crate::message::Message;

use serde::{Deserialize, Serialize};

use std::env;

use std::fmt;

use std::time::Duration;

use tokio::sync::mpsc;

/// What can go wrong while getting a reply.

#[derive(Debug)]

pub enum ChatError {

/// The request did not get a good answer (no network, bad key, server error).

Http(reqwest::Error),

/// The model answered, but with nothing in it.

Empty,

/// /model real was asked for, but a setting is missing.

NotSet(&'static str),

}

impl fmt::Display for ChatError {

fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

match self {

ChatError::Http(e) => write!(f, "the model call failed: {e}"),

ChatError::Empty => write!(

"there was nothing to reply with (an empty conversation or an empty answer)"

ChatError::NotSet(name) => write!(f, "set {name} first (see lesson 9.2)"),

}

/// Lets `?` turn a reqwest error into a ChatError on its own.

impl From<reqwest::Error> for ChatError {

fn from(e: reqwest::Error) -> Self {

ChatError::Http(e)

}

/// The receiving end of an answer: its pieces arrive in order, as they are

/// ready, and the channel closes when the answer is complete.

pub type Pieces = mpsc::Receiver<Result<String, ChatError>>;

/// A promise with a name: any type that implements Replier can answer a

/// conversation. `reply` returns at once; the answer arrives on the channel.

pub trait Replier {

/// A short name, for "/model" and the banner.

fn name(&self) -> &str;

/// Start answering the conversation so far (the last message is the newest).

fn reply(&self, history: &[Message]) -> Pieces;

}

/// A channel that already holds a whole answer, or an error, and is closed.

fn all_at_once(answer: Result<String, ChatError>) -> Pieces {

let (tx, rx) = mpsc::channel(1);

// A new channel with room for one cannot be full, so this send always works.

tx.try_send(answer).expect("a new channel has room for one");

}

/// The echo Chatterbox has had since chapter 3. It answers all at once.

pub struct Echo {

pub max_len: usize,

}

impl Replier for Echo {

fn name(&self) -> &str {

"echo"

}

fn reply(&self, history: &[Message]) -> Pieces {

let message = match history.last() {

Some(m) => &m.text,

None => return all_at_once(Err(ChatError::Empty)),

};

let text = if message.len() > self.max_len {

format!("that is too long (the limit is {} bytes)", self.max_len)

} else if message.ends_with('?') {

String::from("Good question! I can only echo. Try /model real to ask a real model.")

} else {

format!("you said: {message}")

};

all_at_once(Ok(text))

}

/// A stand-in for a real model. It types its answer a word at a time, the way

/// a real model does, so you can watch an answer arrive in pieces.

pub struct Pretend;

impl Replier for Pretend {

fn name(&self) -> &str {

"pretend"

}

fn reply(&self, history: &[Message]) -> Pieces {

let last = match history.last() {

Some(m) => &m.text,

None => return all_at_once(Err(ChatError::Empty)),

};

let text = format!(

"(pretend model) I read all {} message{}. You last said: {last}",

history.len(),

if history.len() == 1 { "" } else { "s" },

);

let (tx, rx) = mpsc::channel(8);

// The task runs on after this function has returned, so it cannot

// borrow anything from here: `move` hands `text` and `tx` over to it.

tokio::spawn(async move {

for (i, word) in text.split(' ').enumerate() {

let piece = if i == 0 {

word.to_string()

} else {

format!(" {word}")

};

if tx.send(Ok(piece)).await.is_err() {

return; // nobody is listening any more

}

tokio::time::sleep(Duration::from_millis(40)).await;

}

});

}

/// One message as a chat API expects it: a role and some content.

#[derive(Serialize)]

struct ApiMessage {

role: String,

content: String,

}

/// The body we send: which model, the whole conversation, and a request to

/// send the answer as it is written instead of all at the end.

#[derive(Serialize)]

struct Request {

model: String,

messages: Vec<ApiMessage>,

stream: bool,

}

/// One piece of a streamed answer. serde ignores every other field.

#[derive(Deserialize)]

struct Chunk {

choices: Vec<ChunkChoice>,

}

#[derive(Deserialize)]

struct ChunkChoice {

delta: Delta,

}

/// The new text in this piece. Some pieces carry none (the first one often

/// only says who is speaking), so it is an Option.

#[derive(Deserialize)]

struct Delta {

content: Option<String>,

}

/// What one line of the stream means.

#[derive(Debug, PartialEq)]

enum Event {

/// More of the answer.

Piece(String),

/// The answer is complete.

Done,

/// Anything else: a blank line between events, a comment, a piece with no text.

Other,

}

/// Read one line of a server-sent event stream as the chat API sends it:

/// `data: {json}` for each piece, then `data: [DONE]`.

fn parse_event(line: &str) -> Event {

let data = match line.strip_prefix("data:") {

Some(rest) => rest.trim(),

None => return Event::Other,

};

if data == "[DONE]" {

return Event::Done;

}

let chunk: Chunk = match serde_json::from_str(data) {

Ok(chunk) => chunk,

Err(_) => return Event::Other,

};

match chunk.choices.into_iter().next() {

Some(choice) => match choice.delta.content {

Some(text) if !text.is_empty() => Event::Piece(text),

_ => Event::Other,

None => Event::Other,

}

/// A real model behind an OpenAI-style chat endpoint, configured from

/// environment variables so the key never sits in the code.

pub struct Model {

url: String,

key: String,

model: String,

}

impl Model {

/// Read the three settings, or say which one is missing.

pub fn from_env() -> Result<Model, ChatError> {

let key =

env::var("CHATTERBOX_API_KEY").map_err(|_| ChatError::NotSet("CHATTERBOX_API_KEY"))?;

let url = env::var("CHATTERBOX_URL").map_err(|_| ChatError::NotSet("CHATTERBOX_URL"))?;

let model =

env::var("CHATTERBOX_MODEL").map_err(|_| ChatError::NotSet("CHATTERBOX_MODEL"))?;

Ok(Model { url, key, model })

}

/// Send the conversation and pass each piece of the answer to `tx` as it

/// arrives over the network.

async fn stream(

url: &str,

key: &str,

body: &Request,

tx: &mpsc::Sender<Result<String, ChatError>>,

) -> Result<(), ChatError> {

let mut response = reqwest::Client::new()

.post(url)

.bearer_auth(key)

.json(body)

.send()

.await?

.error_for_status()?;

// The network delivers bytes in chunks of any size: a chunk can end in the

// middle of a line, or in the middle of a letter that takes several bytes.

// So keep the bytes until a whole line has arrived, and only then read it.

let mut pending: Vec<u8> = Vec::new();

let mut sent_any = false;

while let Some(bytes) = response.chunk().await? {

pending.extend_from_slice(&bytes);

while let Some(end) = pending.iter().position(|&b| b == b'\n') {

let line: Vec<u8> = pending.drain(..=end).collect();

match parse_event(String::from_utf8_lossy(&line).trim()) {

Event::Piece(text) => {

sent_any = true;

if tx.send(Ok(text)).await.is_err() {

return Ok(()); // nobody is listening any more

}

Event::Done => return finished(sent_any),

Event::Other => {}

}

finished(sent_any)

}

/// An answer that ended without a single piece of text is an empty answer.

fn finished(sent_any: bool) -> Result<(), ChatError> {

if sent_any {

Ok(())

} else {

Err(ChatError::Empty)

}

impl Replier for Model {

fn name(&self) -> &str {

"real"

}

fn reply(&self, history: &[Message]) -> Pieces {

// Our names ("you", "bot") become the API's ("user", "assistant").

let mut messages = Vec::new();

for m in history {

let role = if m.author == "you" {

"user"

} else {

"assistant"

};

messages.push(ApiMessage {

role: String::from(role),

content: m.text.clone(),

});

}

let body = Request {

model: self.model.clone(),

messages,

stream: true,

};

// Owned copies: the task may outlive `self`, so it cannot borrow from it.

let url = self.url.clone();

let key = self.key.clone();

let (tx, rx) = mpsc::channel(16);

tokio::spawn(async move {

if let Err(e) = stream(&url, &key, &body, &tx).await {

// If nobody is listening any more, there is no one to tell.

let _ = tx.send(Err(e)).await;

}

});

}

#[cfg(test)]

mod tests {

use super::*;

fn said(text: &str) -> Vec<Message> {

vec![Message::new("you", text)]

}

/// Wait for every piece of an answer and join them.

async fn whole(mut pieces: Pieces) -> Result<String, ChatError> {

let mut text = String::new();

while let Some(piece) = pieces.recv().await {

text.push_str(&piece?);

}

Ok(text)

}

#[tokio::test]

async fn echo_repeats_you() {

let answer = whole(Echo { max_len: 200 }.reply(&said("hi"))).await;

assert_eq!(answer.unwrap(), "you said: hi");

}

#[tokio::test]

async fn echo_refuses_long_messages() {

let long = "x".repeat(10);

let answer = whole(Echo { max_len: 5 }.reply(&said(&long))).await;

assert!(answer.unwrap().starts_with("that is too long"));

}

#[tokio::test]

async fn pretend_reads_the_whole_history() {

let history = vec![

Message::new("you", "a"),

Message::new("bot", "b"),

Message::new("you", "c"),

];

assert_eq!(

whole(Pretend.reply(&history)).await.unwrap(),

"(pretend model) I read all 3 messages. You last said: c"

);

}

#[tokio::test]

async fn pretend_answers_a_word_at_a_time() {

let mut pieces = Pretend.reply(&said("one two"));

let mut count = 0;

while let Some(piece) = pieces.recv().await {

piece.unwrap();

count += 1;

}

assert!(count > 1, "the answer came in {count} piece(s)");

}

#[test]

fn a_data_line_is_a_piece() {

let line = r#"data: {"choices":[{"delta":{"content":"Hel"}}]}"#;

assert_eq!(parse_event(line), Event::Piece(String::from("Hel")));

}

#[test]

fn done_ends_the_answer() {

assert_eq!(parse_event("data: [DONE]"), Event::Done);

}

#[test]

fn lines_without_text_are_skipped() {

assert_eq!(parse_event(""), Event::Other);

assert_eq!(parse_event(": keep-alive"), Event::Other);

let role_only = r#"data: {"choices":[{"delta":{"role":"assistant"}}]}"#;

assert_eq!(parse_event(role_only), Event::Other);

}

#[tokio::test]

async fn nothing_to_reply_to_is_an_error() {

let answer = whole(Echo { max_len: 200 }.reply(&[])).await;

assert!(matches!(answer, Err(ChatError::Empty)));

}

src/store.rsunchanged

//! Saving the conversation to a file and reading it back, as JSON.

use crate::message::Message;

use std::fs;

use std::io;

/// Where the history is kept, next to wherever you run Chatterbox.

pub const FILE: &str = "chatterbox.json";

/// The saved history, or an empty one the very first time (no file yet).

pub fn load(path: &str) -> Result<Vec<Message>, io::Error> {

match fs::read_to_string(path) {

// `?` turns a JSON error into an io::Error for us (serde_json says how).

Ok(text) => Ok(serde_json::from_str(&text)?),

Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(Vec::new()),

Err(e) => Err(e),

}

/// Write the whole history to the file as pretty-printed JSON.

pub fn save(path: &str, history: &[Message]) -> Result<(), io::Error> {

let text = serde_json::to_string_pretty(history)?;

fs::write(path, text)

}

Cargo.toml1 new line

[package]

name = "chatterbox"

version = "1.1.0"

edition = "2024"

[dependencies]

reqwest = { version = "0.13", features = ["json"] }

serde = { version = "1", features = ["derive"] }

serde_json = "1"

tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] }

You type

Sam
tell me a long story
(presses Ctrl-C)
/history
/quit

Chatterbox prints (a real run)

Chatterbox v1.1 (replies by: real)
What should I call you?
Hi Sam! I'm Chatterbox. Ask me anything.
bot: mock model [stopped]
1. you: tell me a long story
2. bot: mock model
bot: bye, Sam!
messages so far: 1
saved 2 messages to chatterbox.json

Download this version (zip)Unzip, then cargo run inside the folder.

And Ctrl-C at the prompt, with the keyboard still open: goodbye, and the conversation saved.

SIGINT sent at the prompt (a real run)

You type

Sam
hello
(presses Ctrl-C)

Chatterbox prints (a real run)

Chatterbox v1.1 (replies by: echo)
What should I call you?
Hi Sam! I'm Chatterbox. Ask me anything.
bot: you said: hello
bot: bye, Sam!
messages so far: 1
saved 2 messages to chatterbox.json

Ctrl-C in Python

In Python, Ctrl-C raises KeyboardInterrupt wherever the program happens to be. tokio turns it into a future instead, and you choose the places where it is waited for.

What you built

That is the course. You started with "what is a program", and you finish with one you wrote: Chatterbox remembers facts, saves the conversation, answers through a trait so the echo, the pretend model and a real one are interchangeable, streams the real model's answer as it is written, and stops it when you ask. Every version of Chatterbox in the course was built, tested and run with scripted conversations, and every compiler message you saw is the compiler's own.

Some of what you met along the way was hard: one owner per value, borrowing, lifetimes, futures that do nothing until awaited. Those are also the ideas that made this lesson short. Stopping a download took no special code, because when the receiver went away, the owner of everything else noticed.

Where to go next

A program that does a real thing is a better place to learn from than any exercise. Roughly in order of size:

A /clear command. Empty the history: a new Command variant, one arm in parse, one in main, a line in HELP and the README, and a test that parse("/clear") gives it.
A system message. Chat-completions services accept a first message with the role system that tells the model how to behave. Add it in front of the history inside Model::reply, so it is sent every time but never saved.
More tests. store.rs has none: save a history to a file in a temporary folder and load it back. /stats has none either, because its counting lives inside main; move it into a function that takes &[Message], and it can be tested like last_n.
A timeout. A third arm in answer()'s select! with tokio::time::sleep, so a model that goes quiet for 30 seconds is stopped too.
The tokio tutorial. It builds a small server and goes deeper into tasks, channels and select!.
The Rust Book. It is free, and you have met most of it. Its async chapter covers what chapter 10 did from another side.

Thank you for building it.

Self-check

Say it out loud, then reveal

Self-checkRetrieval beats re-reading

What does <code>tokio::select!</code> do with the futures that do not finish first?Reveal

It drops them. A dropped future takes no more steps; for recv() nothing is lost, the piece stays in the channel.

When you stop an answer, how does the task sending the pieces find out?Reveal

Returning from answer() drops the receiver, which closes the channel. The task's next send returns an Err, and it returns.

Why does the prompt have to handle Ctrl-C too?Reveal

Once ctrl_c() has been called, Ctrl-C no longer ends the program by itself. Without its own arm, Ctrl-C at the prompt would do nothing; with it, it means /quit.

Why does the keyboard thread use <code>blocking_send</code> and not <code>send(...).await</code>?Reveal

The thread is not async, so it cannot use .await (error E0728). blocking_send waits the plain way, which is fine on a thread that has nothing else to do.

That is lesson 10.4. Tick it off so the course remembers where you are.

Next: 10.5 Chapter review: Streaming replies →

Go deeper

Questions people ask

How do I handle Ctrl-C in a tokio program?

Await tokio::signal::ctrl_c(), usually as one arm of tokio::select!, next to the work it should interrupt. It needs the signal feature of tokio.

Does stopping an answer leave the download running?

Only until the next piece arrives. Returning drops the channel receiver; the task sending pieces gets an error on its next send, which happens when the next piece comes in, and returns, and the connection it owned is dropped then. A server that stalls and sends nothing keeps the connection open until then. A timeout would fix that; this course does not add one.

Why not read the keyboard with tokio::io::stdin?

Its documentation warns that it uses a blocking read that cannot be cancelled, which can keep the program from exiting until Enter is pressed. A thread of your own ends with the program.

What is tokio::select! for?

Waiting on several futures at once. The first one to finish runs its arm of code, and the others are dropped, which for a stopped answer is exactly what you want.

Sources tokio::select! on docs.rs · tokio::signal::ctrl_c on docs.rs · tokio::io::stdin on docs.rs · tokio::sync::mpsc::Sender::blocking_send on docs.rs · The Rust Programming Language: Async and Await · The tokio tutorial