diff --git a/Cargo.lock b/Cargo.lock index 1e7ddee..0038dcd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1042,7 +1042,7 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] name = "openai-api-dispatch" -version = "0.1.1" +version = "0.2.0" dependencies = [ "anyhow", "async-nats", diff --git a/Cargo.toml b/Cargo.toml index 6da7018..b390068 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "openai-api-dispatch" -version = "0.1.1" +version = "0.2.0" authors = ["Victor Lopez "] description = "OpenAI-compatible chat requests through memory or NATS queues." edition = "2024" @@ -38,9 +38,20 @@ uuid = { version = "1.25.0", default-features = false, features = ["serde"] } tokio = { version = "1.53.1", features = ["full"] } [features] -default = ["memory-queue", "nats-queue", "std"] +default = ["jetstream-queue", "memory-queue", "nats-queue", "std"] memory-queue = ["spin"] nats-queue = ["async-nats", "std", "tokio-stream"] +jetstream-queue = [ + "async-nats", + "std", + "tokio-stream", + "async-nats/jetstream", + "async-nats/kv", + "async-nats/server_2_10", + "tokio/macros", + "tokio/rt", + "tokio/time", +] std = [ "anyhow/default", "async-openai", diff --git a/README.md b/README.md index 236cfa3..8cfce76 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,11 @@ [![Documentation](https://docs.rs/openai-api-dispatch/badge.svg)](https://docs.rs/openai-api-dispatch/) [![License](https://img.shields.io/crates/l/openai-api-dispatch.svg)](#license) -OpenAI-compatible chat requests through memory or NATS queues. +OpenAI-compatible chat requests through memory, Core NATS, or durable JetStream queues. Producers submit typed tasks, workers call the configured API, and responses return through the queue. -## NATS quick start +## Core NATS quick start Run a NATS server, then configure the model and OpenAI-compatible endpoint: @@ -41,7 +41,6 @@ async fn main() -> anyhow::Result<()> { .build_chat()?; let response = task.send_and_wait(&producer, Some(30)).await?; - anyhow::ensure!(response.success, "{}", response.contents); println!("{}", response.contents); Ok(()) @@ -52,9 +51,39 @@ async fn main() -> anyhow::Result<()> { Workers subscribe to `openai-api-queue/` by default. Override the prefix with `OPENAI_API_NATS_PREFIX`. In production, run workers and producers as separate processes. +## Durable JetStream queues + +The `jetstream-queue` feature is enabled by default. Use +`queue::jetstream::JetStreamProducer` to submit tasks, and combine +`queue::jetstream::JetStreamWorker` with an executor through `worker::Worker::new`. +The `queue::jetstream` module documents all connection, provisioning, and worker +environment settings. + +JetStream stores tasks in a work-queue stream and keeps status and responses in +a KV bucket. Workers save a response before acknowledging a successful task or +terminating redelivery of a failed task. Defaults are three total deliveries per +message, a 30-second acknowledgment wait, and a one-hour maximum age for task +messages and KV entries. KV updates start a new entry age. Use fresh task IDs for +new work: an existing KV record suppresses publication of that ID. + +Unacknowledged tasks can be redelivered while an active worker is available and +delivery and retention limits permit it. The worker does not renew acknowledgment +deadlines or impose an execution deadline. API calls can repeat if execution +outlasts the acknowledgment wait or a worker crashes before acknowledgment. +Reaching the server's delivery limit does not itself produce an error response, +so callers should bound their response wait. + +Producer and worker constructors open existing resources or create missing ones. +Existing resource configurations are not updated. Set +`OPENAI_API_NATS_STREAM_FORBID_CREATE` and `OPENAI_API_NATS_STORE_FORBID_CREATE` to +require an existing stream and KV bucket; a worker can still create its durable +consumer. Both flags are enabled by any value, including `0` or `false`. + ## Configuration and behavior Environment settings are read when constructing producers, workers, and executors. +The NATS connection and routing notes below describe Core NATS; JetStream uses +the resource and delivery settings documented in `queue::jetstream`. | Variable | Default | | --- | --- | @@ -62,14 +91,14 @@ Environment settings are read when constructing producers, workers, and executor | `OPENAI_API_NATS_WORKERS_GROUP` | `task_workers` | | `OPENAI_API_NATS_PREFIX` | `openai-api-queue/` | | `OPENAI_API_URL` | `http://127.0.0.1:8000/v1` | -| `OPENAI_API_DEFAULT_MODEL` | Unset; required by NATS workers | +| `OPENAI_API_DEFAULT_MODEL` | Unset; required by Core NATS workers | -- A task's explicit model overrides the default and selects its NATS subject. -- Constructors do not wait for server confirmation of subscriptions, so startup can race with publishing. +- A task's explicit model overrides the executor default and selects its Core NATS subject. JetStream uses one configured subject for all task models. +- Core NATS constructors do not wait for server confirmation of subscriptions, so startup can race with publishing. - Each worker handles one task at a time. - Queue/API errors stop its loop without an error reply; supervise spawned workers. -- `send_and_wait` limits only the reply wait, not submission, and expiry does not cancel execution. -- Check `response.success` even when the call returns `Ok`. +- `send_and_wait` returns a `ValidatedResponse` after checking the response's success flag; an unsuccessful response returns `Err`. Its optional timeout limits only the reply wait, not submission, and expiry does not cancel execution. +- Direct calls to `Executor::execute` or `QueueProducer::receive_response` return raw responses; check their `success` flag even when the call returns `Ok`. - Only non-streaming chat is implemented. - The current prompt is sent as a user message; system entries in input history are ignored (use `with_system`). - Schemas request strict JSON output without local validation. @@ -78,17 +107,22 @@ Environment settings are read when constructing producers, workers, and executor ## Features -Default features are `std`, `memory-queue`, and `nats-queue`; NATS implies `std`. The API executor requires `std`. +Default features are `std`, `memory-queue`, `nats-queue`, and `jetstream-queue`. +Both NATS backends imply `std` and Tokio support. The API executor requires `std`. #### no_std -For memory queues, use the feature `memory-queue`. An allocator and pointer/32-bit atomics are required. +Disable default features and enable `memory-queue` to use the polling memory +backend without `std`. An allocator and pointer/32-bit atomics are required. -This backend polls once, has no timeout support, and evicts old items when bounded; the `std` backend uses bounded Tokio channels and backpressure. +This backend polls once, has no timeout support, and evicts old items when bounded; the `std` backend uses bounded Tokio channels and backpressure. ## Development Run `cargo test` for local tests. `just check` also requires `cargo-hack` and the `thumbv7em-none-eabi` target. Live integration tests are ignored by default: run `just check-nats MODEL [URL]` or `just check-openai MODEL [URL]` against configured servers. +Run `just check-jetstream [MODEL] [URL]` against a JetStream-enabled NATS server. +Each invocation uses a separate stream and KV bucket with memory storage, and +each task uses a fresh ID so retained results do not suppress test submissions. ## License diff --git a/justfile b/justfile index 0acb042..45ee0d5 100644 --- a/justfile +++ b/justfile @@ -7,18 +7,27 @@ check-format: cargo fmt --all -- --check check-clippy: - cargo clippy --workspace --all-targets -- -D warnings + cargo clippy --workspace --all-targets --all-features -- -D warnings check-test: cargo hack test --feature-powerset check-doc: - RUSTDOCFLAGS="-D warnings" cargo doc --no-deps + RUSTDOCFLAGS="-D warnings" cargo doc --all-features --no-deps check-nats model="test" url="nats://127.0.0.1:4222": OPENAI_API_DEFAULT_MODEL="{{ model }}" OPENAI_API_NATS_URL="{{ url }}" \ cargo test --features nats-queue nats -- --ignored +check-jetstream model="test" url="nats://127.0.0.1:4222": + test_namespace="openai_api_dispatch_tests_{{ uuid() }}"; \ + OPENAI_API_DEFAULT_MODEL="{{ model }}" \ + OPENAI_API_NATS_URL="{{ url }}" \ + OPENAI_API_NATS_STREAM="$test_namespace" \ + OPENAI_API_NATS_BUCKET="$test_namespace" \ + OPENAI_API_NATS_MEMORY_STORAGE=1 \ + cargo test --features jetstream-queue jetstream -- --ignored + check-openai model="test" url="http://127.0.0.1:8000/v1": OPENAI_API_DEFAULT_MODEL="{{ model }}" OPENAI_API_URL="{{ url }}" \ cargo test openai -- --ignored @@ -31,5 +40,5 @@ check-features: RUSTFLAGS="-D warnings" cargo hack check \ --target thumbv7em-none-eabi \ --feature-powerset \ - --exclude-features std,default,nats-queue \ + --exclude-features std,default,nats-queue,jetstream-queue \ --no-dev-deps diff --git a/src/executor/openai/mod.rs b/src/executor/openai/mod.rs index 248c35f..f2c93fc 100644 --- a/src/executor/openai/mod.rs +++ b/src/executor/openai/mod.rs @@ -30,14 +30,14 @@ mod tests; /// an unsuccessful [`Response`]. Reported usage is total tokens, or zero if absent. /// /// Sends the optional system instruction, prior user/assistant turns, and the -/// current prompt as a user message, in that order. +/// current prompt as a user message, in that order. Successful responses contain +/// those user/assistant turns with the new assistant reply appended; the current +/// prompt appears once and system turns are omitted from returned history. /// /// # Current limitations /// /// - System entries in input history are discarded; use [`TaskChat::system`] /// for a persistent system instruction. -/// - Returned history currently duplicates the latest user prompt: it is kept -/// from the request messages and appended again before the assistant reply. /// - Schemas request strict JSON output named `response`; returned text is not /// locally parsed or validated against the schema. /// - `max_tokens` is cast from `u64` to `u32` without range checking. Payload diff --git a/src/lib.rs b/src/lib.rs index 40efea0..090af7d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,5 @@ #![cfg_attr(not(feature = "std"), no_std)] -#![warn(missing_docs)] +#![deny(missing_docs)] #![doc = include_str!("../README.md")] extern crate alloc; diff --git a/src/queue/jetstream/mod.rs b/src/queue/jetstream/mod.rs new file mode 100644 index 0000000..88235fb --- /dev/null +++ b/src/queue/jetstream/mod.rs @@ -0,0 +1,80 @@ +//! Durable task delivery through NATS JetStream, with task status and responses +//! retained in a key-value (KV) bucket. +//! +//! [crate::queue::jetstream::JetStreamProducer] publishes JSON tasks to a work-queue stream and watches +//! their KV records for responses. [crate::queue::jetstream::JetStreamWorker] shares a durable pull +//! consumer with other workers in its group, waits for tasks, and records results +//! before acknowledging deliveries. This module requires the `jetstream-queue` +//! feature, a Tokio runtime, and a server with JetStream enabled. +//! +//! # Environment configuration +//! +//! [crate::queue::jetstream::JetStreamProducer::from_env_or_default] and +//! [crate::queue::jetstream::JetStreamWorker::from_env_or_default] read the following variables during +//! construction. Existing instances do not track subsequent environment changes. +//! Defaults apply when a variable is unset or contains non-Unicode data. Invalid +//! numeric text returns an error when the corresponding setting is parsed; +//! resource names and numeric limits are also subject to NATS validation. +//! +//! ## Connection and resource names +//! +//! | Variable | Default | Behavior | +//! | --- | --- | --- | +//! | `OPENAI_API_NATS_URL` | `nats://localhost:4222` | NATS connection URL used by both producers and workers. | +//! | `OPENAI_API_NATS_STREAM` | `OPENAI_API_DISPATCH` | Task stream name, also used as the prefix of the task subject. | +//! | `OPENAI_API_NATS_SUBJECT` | `dispatch.task` | Subject suffix. The full subject is `.`, with a literal dot inserted between the values. | +//! | `OPENAI_API_NATS_BUCKET` | `openai_api_dispatch_task` | KV bucket containing task records keyed by the decimal task ID. | +//! +//! The default task subject is `OPENAI_API_DISPATCH.dispatch.task`. Producers +//! publish to this subject and workers filter on it; task models do not select +//! subjects in this backend. Producers and workers must agree on the stream, +//! subject, and bucket. The bucket name is independent of the stream name, so use +//! distinct names for both when isolating queues. An existing task record causes +//! the producer to reuse that record instead of publishing the task again. +//! +//! ## Stream and KV provisioning +//! +//! Both constructors reuse existing resources and create missing ones by default. +//! Creation settings below do not update existing streams or buckets. +//! +//! | Variable | Default | Behavior | +//! | --- | --- | --- | +//! | `OPENAI_API_NATS_MEMORY_STORAGE` | Unset (file storage) | Presence selects memory storage for newly created task streams and KV buckets. | +//! | `OPENAI_API_NATS_TTL_SECS` | `3600` | Maximum age in seconds (`u64`) of task stream messages and KV entries. Each KV update has its own age; this is not a response-wait timeout. | +//! | `OPENAI_API_NATS_DUPLICATE_WINDOW_SECS` | `600` | Stream deduplication window in seconds (`u64`) for message IDs, which the producer sets to task IDs. KV record deduplication is separate. | +//! | `OPENAI_API_NATS_REPLICAS` | `1` | Replica count (`usize`) for the task stream. KV bucket replication uses the client's default configuration. | +//! | `OPENAI_API_NATS_MAX_BYTES` | `0` | Total KV bucket size limit in bytes (`i64`), passed to NATS through the KV configuration. Applies only when creating a bucket. | +//! | `OPENAI_API_NATS_STREAM_FORBID_CREATE` | Unset | Presence requires the task stream to exist; fetching it must succeed. Stream creation settings are skipped. | +//! | `OPENAI_API_NATS_STORE_FORBID_CREATE` | Unset | Presence requires the KV bucket to be accessible; a failed lookup returns an error instead of attempting creation. | +//! +//! The three presence flags (`MEMORY_STORAGE`, `STREAM_FORBID_CREATE`, and +//! `STORE_FORBID_CREATE`, each prefixed with `OPENAI_API_NATS_`) accept any Unicode +//! value, including an empty string, `0`, or `false`, as enabled. Unset a flag to +//! disable it. The creation restrictions apply to the stream and bucket only; +//! workers still create their durable consumer if it is missing. +//! +//! ## Worker settings +//! +//! These variables are read by [crate::queue::jetstream::JetStreamWorker::from_env_or_default]. Existing +//! durable consumers are reused without updating their server configuration. +//! +//! | Variable | Default | Behavior | +//! | --- | --- | --- | +//! | `OPENAI_API_NATS_CONSUMER_GROUP` | `openai_api_dispatch_workers` | Durable pull consumer name. Workers using the same stream and consumer share task deliveries. | +//! | `OPENAI_API_NATS_ACK_WAIT` | `30` | Time in seconds (`u64`) the consumer waits for an acknowledgment before a delivery becomes eligible for redelivery. | +//! | `OPENAI_API_NATS_MAX_DELIVER` | `3` | Consumer delivery limit per message (`i64`), including the initial delivery. Positive values also enable the worker's local delivery-count check. | +//! | `OPENAI_API_NATS_MAX_ACK_PENDING` | `100` | Maximum number of outstanding, unacknowledged messages (`i64`) across the entire consumer, shared by all its workers. | +//! | `OPENAI_API_DEFAULT_MODEL` | Unset | Optional fallback model when the worker constructs a delivery-limit error response for a task without an explicit model. | +//! +//! The default model is not required to construct a JetStream worker and does not +//! affect its subscription. Executing tasks may impose separate model requirements. +//! +mod producer; +mod types; +mod worker; + +#[cfg(test)] +mod tests; + +pub use producer::JetStreamProducer; +pub use worker::JetStreamWorker; diff --git a/src/queue/jetstream/producer.rs b/src/queue/jetstream/producer.rs new file mode 100644 index 0000000..949cca1 --- /dev/null +++ b/src/queue/jetstream/producer.rs @@ -0,0 +1,183 @@ +use async_nats::{ + header::{NATS_EXPECTED_STREAM, NATS_MESSAGE_ID}, + jetstream::{self, kv::Store}, +}; +use tokio_stream::StreamExt as _; + +use crate::{ + queue::{ + QueueProducer, + jetstream::types::{StoreStreamSubject, TaskRecord, TaskStatus}, + }, + task::{Response, Task}, + utils, +}; + +#[derive(Debug, Clone)] +/// Publishes durable tasks and retrieves their responses from a JetStream KV bucket. +/// +/// Tasks are published to the configured subject, independently of their model. +/// An existing, valid KV record for a task ID is reused without publishing again. +/// Receiving a response waits for a completed or failed record and returns the +/// stored [`Response`], whose success flag may be false. Use [`Task::send_and_wait`] +/// to check that flag and optionally limit the response wait. +/// +/// KV writes and task publication are separate operations. A publish error can +/// leave a queued record without a corresponding stream message; reusing that +/// task ID returns the record without retrying publication. +pub struct JetStreamProducer { + js: jetstream::Context, + store: Store, + stream: String, + subject: String, +} + +impl JetStreamProducer { + /// Connects to NATS and opens or creates the configured task stream and KV bucket. + /// + /// Reads the connection and provisioning variables documented in the + /// [module configuration](crate::queue::jetstream). Existing resources are + /// reused without updating their configuration. A creation-forbidden flag + /// requires the corresponding resource to be accessible already. + /// + /// Returns errors for connection failures, invalid configuration, or failed + /// resource lookup or creation. No default model is required. + pub async fn from_env_or_default() -> anyhow::Result { + let client = utils::nats_client_from_env_or_default().await?; + let js = jetstream::new(client); + + TaskRecord::get_or_maybe_create_stream(&js).await?; + + let StoreStreamSubject { + store, + stream, + subject, + } = TaskRecord::get_or_maybe_create_store(&js).await?; + + Ok(Self { + js, + store, + stream, + subject, + }) + } +} + +impl QueueProducer for JetStreamProducer { + type Message = TaskRecord; + + async fn send_task(&self, task: Task) -> anyhow::Result { + tracing::debug!("nats-jetstream queue; sending task `{}`", task.id,); + + let record = TaskRecord::from(&task); + + match self.store.entry(&record.id).await? { + Some(entry) => match TaskRecord::try_from_json_bytes(entry.value) { + Ok(r) => { + tracing::debug!("nats-jetstream queue; already queued `{}`", task.id,); + return Ok(r); + } + Err(e) => { + tracing::warn!( + "nats-jetstream queue; overwriting invalid task `{}` on store({}): {e}", + entry.revision, + task.id, + ); + self.store + .update(&record.id, record.to_json_bytes().into(), entry.revision) + .await?; + } + }, + None => { + self.store + .put(&record.id, record.to_json_bytes().into()) + .await?; + } + } + + let mut headers = async_nats::HeaderMap::new(); + + headers.insert(NATS_MESSAGE_ID, task.id.to_string()); + headers.insert(NATS_EXPECTED_STREAM, self.stream.as_str()); + + let payload = task.to_json_bytes(); + let ack = self + .js + .publish_with_headers(self.subject.clone(), headers, payload.into()) + .await? + .await?; + + if ack.duplicate { + tracing::warn!( + "nats-jetstream queue; duplicated task `{}` post KV check", + task.id, + ); + anyhow::bail!("task `{}` discarded; duplicated", task.id); + } + + tracing::debug!("nats-jetstream queue; task `{}` sent", task.id,); + + Ok(record) + } + + async fn receive_response(&self, record: Self::Message) -> anyhow::Result> { + let mut watcher = self.store.watch_with_history(&record.id).await?; + + while let Some(maybe_entry) = watcher.next().await { + let entry = maybe_entry?; + + if entry.value.is_empty() { + tracing::trace!( + "nats-jetstream queue; skipping empty tombstone `{}`", + record.id, + ); + continue; + } + + match TaskRecord::try_from_json_bytes(&entry.value)?.status { + TaskStatus::Queued => { + tracing::trace!("nats-jetstream queue; waiting task queued `{}`", record.id,); + } + TaskStatus::Running { worker } => { + tracing::trace!( + "nats-jetstream queue; waiting task queued `{}` with worker `{}`", + record.id, + worker + ); + } + TaskStatus::Completed { result } => { + anyhow::ensure!( + result.success, + "nats-jetstream queue; task `{}` completed but with failed `{}`", + record.id, + result.contents + ); + + tracing::debug!( + "nats-jetstream queue; received response on task `{}`", + record.id, + ); + + return Ok(Some(result)); + } + TaskStatus::Failed { error } => { + anyhow::ensure!( + !error.success, + "nats-jetstream queue; task `{}` completed but with failed `{}`", + record.id, + error.contents + ); + + tracing::debug!( + "nats-jetstream queue; received error response on task `{}`", + record.id, + ); + + return Ok(Some(error)); + } + } + } + + anyhow::bail!("nats-jetstream queue; no response for task `{}`", record.id,); + } +} diff --git a/src/queue/jetstream/tests.rs b/src/queue/jetstream/tests.rs new file mode 100644 index 0000000..6d31ba3 --- /dev/null +++ b/src/queue/jetstream/tests.rs @@ -0,0 +1,161 @@ +use std::{env, time::Duration}; + +use tokio::{sync::Mutex, time::timeout}; + +use super::{JetStreamProducer, JetStreamWorker}; +use crate::{ + queue::{QueueProducer as _, QueueWorker as _}, + task::{Response, Task, TaskBuilder}, + utils, +}; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +// The default workers share a durable consumer, so tests within one invocation +// must not compete for tasks. The just recipe isolates stream and bucket names. +static NATS_TEST: Mutex<()> = Mutex::const_new(()); + +fn worker_model() -> String { + utils::get_default_model() + .as_ref() + .clone() + .or_else(|| env::var("OPENAI_API_DEFAULT_MODEL").ok()) + .expect("set OPENAI_API_DEFAULT_MODEL when running NATS tests") +} + +fn task(model: &str) -> Task { + TaskBuilder::new() + .with_prompt("user prompt") + .with_model(model) + .build_chat() + .unwrap() +} + +async fn default_queue() -> (JetStreamProducer, JetStreamWorker) { + let worker = JetStreamWorker::from_env_or_default().await.unwrap(); + let producer = JetStreamProducer::from_env_or_default().await.unwrap(); + + (producer, worker) +} + +#[tokio::test] +#[ignore = "requires a NATS server and a configured worker model"] +async fn default_worker_receives_tasks_from_default_producer() { + let _guard = NATS_TEST.lock().await; + let (producer, worker) = default_queue().await; + let expected = task(&worker_model()); + + producer.send_task(expected.clone()).await.unwrap(); + let received = timeout(TEST_TIMEOUT, worker.receive_task()) + .await + .expect("timed out waiting for the NATS task") + .unwrap() + .expect("the NATS worker subscription ended"); + + assert_eq!(received.task, expected); + + worker + .send_response( + received.message, + Response::success(expected, 0, worker_model(), "response"), + ) + .await + .unwrap(); +} + +#[tokio::test] +#[ignore = "requires a NATS server and a configured worker model"] +async fn default_producer_receives_responses_from_default_worker() { + let _guard = NATS_TEST.lock().await; + let (producer, worker) = default_queue().await; + let expected_task = task(&worker_model()); + let producer_message = producer.send_task(expected_task.clone()).await.unwrap(); + let worker_message = timeout(TEST_TIMEOUT, worker.receive_task()) + .await + .expect("timed out waiting for the NATS task") + .unwrap() + .expect("the NATS worker subscription ended"); + let expected_response = Response::success(expected_task, 42, worker_model(), "response"); + + worker + .send_response(worker_message.message, expected_response.clone()) + .await + .unwrap(); + + let received = timeout(TEST_TIMEOUT, producer.receive_response(producer_message)) + .await + .expect("timed out waiting for the NATS response") + .unwrap(); + + assert_eq!(received, Some(expected_response)); +} + +#[tokio::test] +#[ignore = "requires a NATS server and a configured worker model"] +async fn default_worker_waits_for_tasks_when_queue_is_empty() { + let _guard = NATS_TEST.lock().await; + let (producer, worker) = default_queue().await; + + // Check both the initial idle queue and an idle period after completing work. + for _ in 0..2 { + let receive = worker.receive_task(); + tokio::pin!(receive); + assert!( + timeout(Duration::from_millis(100), &mut receive) + .await + .is_err(), + "the worker stopped waiting while the queue was empty" + ); + + let expected = task(&worker_model()); + producer.send_task(expected.clone()).await.unwrap(); + let received = timeout(TEST_TIMEOUT, receive) + .await + .expect("timed out waiting for the JetStream task") + .unwrap() + .expect("the JetStream worker subscription ended"); + + assert_eq!(received.task, expected); + worker + .send_response( + received.message, + Response::success(expected, 0, worker_model(), "response"), + ) + .await + .unwrap(); + } +} + +#[tokio::test] +#[ignore = "requires a NATS server and a configured worker model"] +async fn default_worker_preserves_pending_tasks_across_calls_and_clones() { + let _guard = NATS_TEST.lock().await; + let (producer, worker) = default_queue().await; + let cloned_worker = worker.clone(); + let tasks = [ + task(&worker_model()), + task(&worker_model()), + task(&worker_model()), + ]; + + for expected in &tasks { + producer.send_task(expected.clone()).await.unwrap(); + } + + for (receiver, expected) in [&worker, &cloned_worker, &worker].into_iter().zip(tasks) { + let received = timeout(TEST_TIMEOUT, receiver.receive_task()) + .await + .expect("timed out waiting for a pending JetStream task") + .unwrap() + .expect("the JetStream worker subscription ended"); + + assert_eq!(received.task, expected); + receiver + .send_response( + received.message, + Response::success(expected, 0, worker_model(), "response"), + ) + .await + .unwrap(); + } +} diff --git a/src/queue/jetstream/types.rs b/src/queue/jetstream/types.rs new file mode 100644 index 0000000..40b3957 --- /dev/null +++ b/src/queue/jetstream/types.rs @@ -0,0 +1,232 @@ +use core::time::Duration; +use std::env; + +use async_nats::jetstream::{ + Context, + kv::{self, Store}, + stream::{Config, DiscardPolicy, RetentionPolicy, StorageType, Stream}, +}; +use serde::{Deserialize, Serialize}; + +use crate::task::{Response, Task}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +/// Task lifecycle state stored in the JetStream KV bucket. +pub enum TaskStatus { + /// Recorded as queued, awaiting a worker. + Queued, + /// Received by a worker for execution. + Running { + /// ID of the worker that most recently received the task. + worker: u128, + }, + /// Recorded as completed with a retained response. + Completed { + /// Response saved when the task completed successfully. + result: Response, + }, + /// Recorded as failed with a retained error response. + Failed { + /// Unsuccessful response saved for the task. + error: Response, + }, +} + +#[derive(Debug, Clone)] +/// KV store and task stream addressing resolved during queue construction. +pub struct StoreStreamSubject { + /// KV bucket containing task status and retained responses. + pub store: Store, + /// Name of the stream that stores task messages. + pub stream: String, + /// Full subject used to publish and receive task messages. + pub subject: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +/// Task ID and status retained in JetStream KV and used as a producer response handle. +/// +/// Converting from a task creates a queued record. Status methods update only +/// this in-memory value; they do not write to KV, acknowledge deliveries, or +/// validate that a response matches the record's ID or status. +pub struct TaskRecord { + /// Decimal task ID, used as the KV key. + pub id: String, + /// Latest recorded lifecycle state and any terminal response. + pub status: TaskStatus, +} + +impl From<&Task> for TaskRecord { + fn from(task: &Task) -> Self { + Self { + id: task.id.to_string(), + status: TaskStatus::Queued, + } + } +} + +impl TaskRecord { + /// Serializes the task record into JSON bytes. + pub fn to_json_bytes(&self) -> Vec { + serde_json::to_vec(self).expect("infallible serialization") + } + + /// Deserializes the task record from JSON bytes. + pub fn try_from_json_bytes>(bytes: B) -> anyhow::Result { + let bytes = bytes.as_ref(); + + Ok(serde_json::from_slice(bytes)?) + } + + /// Marks this record completed and retains a copy of the supplied response. + pub fn completed(&mut self, response: Response) { + self.status = TaskStatus::Completed { + result: response.clone(), + }; + } + + /// Marks this record failed and retains a copy of the supplied response. + pub fn failed(&mut self, response: Response) { + self.status = TaskStatus::Failed { + error: response.clone(), + }; + } + + /// Marks this record running under the supplied worker ID. + pub fn running(&mut self, worker: u128) { + self.status = TaskStatus::Running { worker }; + } + + pub(crate) fn get_max_age() -> anyhow::Result { + let max_age = env::var("OPENAI_API_NATS_TTL_SECS") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_TTL_SECS: {e}"))? + .unwrap_or(3600); + + Ok(Duration::from_secs(max_age)) + } + + fn get_stream() -> String { + env::var("OPENAI_API_NATS_STREAM").unwrap_or_else(|_| "OPENAI_API_DISPATCH".into()) + } + + fn get_storage_type() -> StorageType { + if env::var("OPENAI_API_NATS_MEMORY_STORAGE").ok().is_some() { + StorageType::Memory + } else { + StorageType::File + } + } + + pub(crate) fn get_subject() -> String { + let stream = Self::get_stream(); + let subject = + env::var("OPENAI_API_NATS_SUBJECT").unwrap_or_else(|_| "dispatch.task".into()); + + format!("{stream}.{subject}") + } + + pub(crate) async fn get_or_maybe_create_stream(js: &Context) -> anyhow::Result { + let stream = Self::get_stream(); + + tracing::info!("nats-jetstream queue; using stream `{stream}`",); + + if env::var("OPENAI_API_NATS_STREAM_FORBID_CREATE") + .ok() + .is_some() + { + tracing::info!( + "nats-jetstream queue; stream creation forbidden by OPENAI_API_NATS_STREAM_FORBID_CREATE, fetching...", + ); + return Ok(js.get_stream(&stream).await?); + } + + let subjects = vec![Self::get_subject()]; + + let max_age = Self::get_max_age()?; + let storage = Self::get_storage_type(); + + let duplicate_window = env::var("OPENAI_API_NATS_DUPLICATE_WINDOW_SECS") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_DUPLICATE_WINDOW_SECS: {e}"))? + .unwrap_or(600); + + let num_replicas = env::var("OPENAI_API_NATS_REPLICAS") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_REPLICAS: {e}"))? + .unwrap_or(1); + + let config = Config { + name: stream, + subjects, + retention: RetentionPolicy::WorkQueue, + discard: DiscardPolicy::New, + duplicate_window: Duration::from_secs(duplicate_window), + max_age, + storage, + num_replicas, + ..Default::default() + }; + + Ok(js.get_or_create_stream(config).await?) + } + + pub(crate) async fn get_or_maybe_create_store( + js: &Context, + ) -> anyhow::Result { + let stream = Self::get_stream(); + let subject = Self::get_subject(); + let bucket = env::var("OPENAI_API_NATS_BUCKET") + .unwrap_or_else(|_| "openai_api_dispatch_task".into()); + + tracing::info!( + "nats-jetstream queue; using stream `{stream}`, subject `{subject}`, bucket `{bucket}`" + ); + + let store = match js.get_key_value(&bucket).await { + Ok(s) => s, + Err(e) => { + tracing::info!("nats-jetstream queue; bucket `{bucket}` not present, creating...",); + + anyhow::ensure!( + env::var("OPENAI_API_NATS_STORE_FORBID_CREATE") + .ok() + .is_none(), + "nats-jetstream queue; store unavailable and OPENAI_API_NATS_STORE_FORBID_CREATE set: {e}" + ); + + let max_age = TaskRecord::get_max_age()?; + let storage = Self::get_storage_type(); + let max_bytes = env::var("OPENAI_API_NATS_MAX_BYTES") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_MAX_BYTES: {e}"))? + .unwrap_or(0); + + js.create_key_value(kv::Config { + bucket, + max_age, + max_bytes, + description: "openai-api-dispatch tasks bucket".to_string(), + storage, + history: 1, + ..Default::default() + }) + .await? + } + }; + + Ok(StoreStreamSubject { + store, + stream, + subject, + }) + } +} diff --git a/src/queue/jetstream/worker.rs b/src/queue/jetstream/worker.rs new file mode 100644 index 0000000..d371bfc --- /dev/null +++ b/src/queue/jetstream/worker.rs @@ -0,0 +1,206 @@ +use core::time::Duration; +use std::env; + +use alloc::sync::Arc; +use async_nats::jetstream::{ + self, AckKind, Message, + consumer::{ + AckPolicy, DeliverPolicy, ReplayPolicy, + pull::{Config, Stream}, + }, + kv::Store, +}; +use tokio::sync::Mutex; +use tokio_stream::StreamExt as _; + +use crate::{ + queue::{QueueWorker, WrappedTask, jetstream::types::TaskRecord}, + task::{Response, Task}, + utils, +}; + +#[derive(Clone)] +/// Receives durable tasks and saves their responses before acknowledging delivery. +/// +/// Clones share one pull stream and worker ID. Receiving waits when the queue is +/// idle and marks each returned task as running in the KV bucket. A successful +/// response is stored and acknowledged; an unsuccessful response is stored and +/// terminates redelivery of that message. +/// +/// Acknowledgment deadlines are not renewed during execution. Unacknowledged +/// tasks can be redelivered within the consumer's delivery and retention limits, +/// so API calls may repeat. Reaching the server's delivery limit does not itself +/// create a failed KV record or a response for the producer. +pub struct JetStreamWorker { + id: u128, + store: Store, + messages: Arc>, + max_deliver: i64, + default_model: Arc>, +} + +impl JetStreamWorker { + /// Connects to NATS and opens or creates the queue resources and durable consumer. + /// + /// Reads the variables documented in the + /// [module configuration](crate::queue::jetstream), including the consumer + /// group, acknowledgment wait, and delivery limits. Existing resources are + /// reused without updating their server configuration. Stream and bucket + /// creation restrictions do not prevent creation of a missing consumer. + /// + /// Returns errors for connection failures, invalid configuration, or failed + /// resource operations. The default model is optional and only supplies a + /// fallback when this queue constructs a delivery-limit error response. + pub async fn from_env_or_default() -> anyhow::Result { + let id = utils::id(); + let default_model = utils::get_default_model(); + let client = utils::nats_client_from_env_or_default().await?; + let js = jetstream::new(client); + let store = TaskRecord::get_or_maybe_create_store(&js).await?.store; + let stream = TaskRecord::get_or_maybe_create_stream(&js).await?; + + let durable_name = env::var("OPENAI_API_NATS_CONSUMER_GROUP") + .unwrap_or_else(|_| "openai_api_dispatch_workers".into()); + let consumers_description = "OpenAI API dispatch worker consumer pool".to_string(); + + let ack_wait = env::var("OPENAI_API_NATS_ACK_WAIT") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_ACK_WAIT: {e}"))? + .unwrap_or(30); + + let max_deliver = env::var("OPENAI_API_NATS_MAX_DELIVER") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_MAX_DELIVER: {e}"))? + .unwrap_or(3); + + let max_ack_pending = env::var("OPENAI_API_NATS_MAX_ACK_PENDING") + .ok() + .map::, _>(|ttl| ttl.parse()) + .transpose() + .map_err(|e| anyhow::anyhow!("invalid OPENAI_API_NATS_MAX_ACK_PENDING: {e}"))? + .unwrap_or(100); + + // TODO add suport per worker model capability + let subject = TaskRecord::get_subject(); + + let config = Config { + durable_name: Some(durable_name.clone()), + description: Some(consumers_description), + filter_subject: subject, + max_deliver, + max_ack_pending, + ack_wait: Duration::from_secs(ack_wait), + ack_policy: AckPolicy::Explicit, + deliver_policy: DeliverPolicy::All, + replay_policy: ReplayPolicy::Instant, + ..Default::default() + }; + + let consumer = stream.get_or_create_consumer(&durable_name, config).await?; + let messages = consumer + .stream() + .max_messages_per_batch(1) + .messages() + .await?; + + Ok(Self { + id, + store, + messages: Arc::new(Mutex::new(messages)), + max_deliver, + default_model, + }) + } +} + +impl QueueWorker for JetStreamWorker { + type Message = Message; + + async fn receive_task(&self) -> anyhow::Result>> { + let mut messages = self.messages.lock().await; + + while let Some(message) = messages.next().await { + let message = message.map_err(|e| anyhow::anyhow!(e))?; + + let task = Task::try_from_json_bytes(&message.payload)?; + let mut record = TaskRecord::from(&task); + + tracing::debug!("nats-jetstream queue; task `{}` received", task.id,); + + if 0 < self.max_deliver { + let info = message.info().map_err(|e| anyhow::anyhow!(e))?; + + if self.max_deliver < info.delivered { + tracing::debug!( + "nats-jetstream queue; task `{}` rejected with delivery attempts `{}`", + task.id, + info.delivered + ); + + let tokens = 0; + let model = task.model(&self.default_model)?; + let error = format!( + "maximum deliveries `{}` for task `{}` reached", + info.delivered, task.id + ); + + let response = Response::error(task, tokens, model, error); + + self.send_response(message, response).await?; + + continue; + } + } + + tracing::debug!("nats-jetstream queue; task `{}` accepted", task.id,); + + record.running(self.id); + + self.store + .put(&record.id, record.to_json_bytes().into()) + .await?; + + return Ok(Some(WrappedTask { message, task })); + } + + Ok(None) + } + + async fn send_response( + &self, + message: Self::Message, + response: Response, + ) -> anyhow::Result<()> { + let id = response.task.id; + + tracing::debug!("nats-jetstream queue; response `{id}` received"); + + let mut record = TaskRecord::from(&response.task); + + // TODO implement message retry Nack(Duration) + let ack = if response.success { + record.completed(response); + AckKind::Ack + } else { + record.failed(response); + AckKind::Term + }; + + self.store + .put(&record.id, record.to_json_bytes().into()) + .await?; + + message + .ack_with(ack) + .await + .map_err(|e| anyhow::anyhow!(e))?; + + tracing::debug!("nats-jetstream queue; response `{id}` accepted"); + + Ok(()) + } +} diff --git a/src/queue/mod.rs b/src/queue/mod.rs index 3671a27..fcd1175 100644 --- a/src/queue/mod.rs +++ b/src/queue/mod.rs @@ -10,6 +10,10 @@ pub mod memory; /// JSON request/reply over Core NATS; requires `std` and a Tokio runtime. pub mod nats; +#[cfg(feature = "jetstream-queue")] +/// Durable tasks, retained results, and optional JetStream provisioning. +pub mod jetstream; + /// Submits tasks and retrieves their responses. /// /// Waiting and cancellation semantics depend on the backend. Returned futures @@ -23,8 +27,11 @@ pub trait QueueProducer { /// Consumes a handle to retrieve a response. /// + /// A returned [`Response`] may have `success == false`; use + /// [`Task::send_and_wait`] to turn that outcome into an error. /// `None` may mean no response is ready (`no_std` memory) or the response - /// stream ended (NATS); consult the backend before treating it as terminal. + /// stream ended (Core NATS). JetStream returns an error if its KV watch ends + /// without a result; consult the backend before treating `None` as terminal. fn receive_response( &self, message: Self::Message, @@ -50,7 +57,7 @@ pub trait QueueWorker { /// Receives a task, or `None` when none is available or the stream ends. /// - /// NATS and `std` memory wait for work; `no_std` memory polls once. + /// Core NATS, JetStream, and `std` memory wait for work; `no_std` memory polls once. fn receive_task( &self, ) -> impl Future>>>; diff --git a/src/queue/nats/mod.rs b/src/queue/nats/mod.rs index bb0a0f4..205ce2f 100644 --- a/src/queue/nats/mod.rs +++ b/src/queue/nats/mod.rs @@ -55,17 +55,7 @@ impl NatsProducer { let default_model = utils::get_default_model(); let prefix = env::var("OPENAI_API_NATS_PREFIX").unwrap_or_else(|_| "openai-api-queue/".into()); - let url = - env::var("OPENAI_API_NATS_URL").unwrap_or_else(|_| "nats://localhost:4222".into()); - let client = async_nats::ConnectOptions::new() - .request_timeout(None) - .connect(url) - .await?; - - tracing::info!( - "nats client connected to `{:?}`", - client.server_info().connect_urls - ); + let client = utils::nats_client_from_env_or_default().await?; Ok(Self { client, diff --git a/src/task.rs b/src/task.rs index d678b79..62c316a 100644 --- a/src/task.rs +++ b/src/task.rs @@ -10,9 +10,9 @@ use crate::{queue::QueueProducer, utils}; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] /// Builds a chat task; only the prompt is required at build time. /// -/// Prefer [`Self::new`] for a generated ID: derived `Default` uses ID zero. +/// Both [`Self::new`] and [`Default::default`] generate an ID through [`utils::id`]. pub struct TaskBuilder { - /// Correlation ID; must be unique among outstanding memory-queue tasks. + /// Correlation ID; must be unique for new work, including retained JetStream tasks. pub id: u128, /// Explicit model, overriding the producer/executor default when present. pub model: Option, @@ -51,7 +51,7 @@ impl TaskBuilder { } } - /// Overrides the default model; NATS uses this model to select a subject. + /// Overrides the default model; Core NATS also uses it to select a subject. pub fn with_model(mut self, model: M) -> Self { self.model.replace(model.to_string()); self @@ -144,7 +144,7 @@ impl TaskBuilder { /// Use [`TaskBuilder`] for chat tasks. Derived `Default` has ID zero and null /// contents, so it is not an executable chat request. pub struct Task { - /// Correlation ID; memory queues require uniqueness for outstanding tasks. + /// Correlation ID; must be unique for new work, including retained JetStream tasks. pub id: u128, /// Operation encoded in [`Self::contents`]. pub task_type: TaskType, @@ -159,21 +159,23 @@ pub struct Task { } impl Task { - /// Submits this task and retrieves its response, which may be unsuccessful. + /// Submits this task and returns its response after checking the success flag. /// - /// With `std`, `Some(seconds)` limits only the response wait after submission - /// and requires a Tokio runtime with time enabled. Expiry does not cancel - /// worker execution. `None` adds no timeout; backend polling semantics still - /// apply. Without `std`, a supplied timeout errors before sending. + /// Returns [`ValidatedResponse`] when the received response's `success` flag + /// is true. An unsuccessful response, missing response, or queue error returns + /// `Err`. This check does not validate output text against a JSON schema. /// - /// Submission/retrieval failures, timeout expiry, and `None` responses return - /// errors. Check the response's `success` field separately from the `Result`. + /// With `std`, `Some(seconds)` limits only the response wait after submission + /// and requires a Tokio runtime with time enabled. Expiry returns an error + /// without cancelling worker execution. `None` adds no timeout; backend + /// polling semantics still apply. Without `std`, supplying a timeout returns + /// an error before submitting the task. pub async fn send_and_wait( self, queue: &Q, timeout_secs: Option, - ) -> anyhow::Result { - match timeout_secs { + ) -> anyhow::Result { + let response = match timeout_secs { #[cfg(feature = "std")] Some(t) => { let timeout = core::time::Duration::from_secs(t); @@ -193,7 +195,72 @@ impl Task { response.ok_or_else(|| anyhow::anyhow!("task response not available")) } - } + }?; + + ValidatedResponse::try_from(response) + } + + /// Serializes the task into JSON bytes. + pub fn to_json_bytes(&self) -> Vec { + serde_json::to_vec(self).expect("infallible serialization") + } + + /// Deserializes the task from JSON bytes. + pub fn try_from_json_bytes>(bytes: B) -> anyhow::Result { + let bytes = bytes.as_ref(); + + Ok(serde_json::from_slice(bytes)?) + } + + /// Replaces the correlation ID; use a fresh ID for new work. + /// + /// Memory queues correlate outstanding responses by this ID. JetStream also + /// uses it for deduplication while task records or message IDs are retained. + pub fn with_id(mut self, id: u128) -> Self { + self.id = id; + self + } + + /// Sets the output limit; the API executor casts it to `u32` unchecked. + pub fn with_max_tokens(mut self, max_tokens: u64) -> Self { + self.max_tokens.replace(max_tokens); + self + } + + /// Attaches opaque caller metadata to be returned with the task. + /// + /// # Panics + /// + /// Panics if the payload cannot be serialized as JSON. + pub fn with_payload(mut self, payload: P) -> Self { + // failure serialization is considered unrecoverable from the API documentation + let payload = serde_json::to_value(&payload).expect("failed to serialize payload"); + self.payload.replace(payload); + self + } + + /// Overrides the model used for Core NATS routing and API execution. + pub fn with_model(mut self, model: M) -> Self { + self.model.replace(model.to_string()); + self + } + + /// Decodes contents as [`TaskChat`], returning an error for incompatible JSON. + /// + /// Does not inspect [`Self::task_type`] or validate schema/model settings. + pub fn try_to_chat(&self) -> anyhow::Result { + Ok(serde_json::from_value(self.contents.clone())?) + } + + #[cfg(feature = "std")] + pub(crate) fn model( + &self, + default_model: &std::sync::Arc>, + ) -> anyhow::Result { + self.model + .clone() + .or_else(|| default_model.as_ref().clone()) + .ok_or_else(|| anyhow::anyhow!("no model provided")) } } @@ -212,6 +279,47 @@ pub struct Response { pub contents: String, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +/// Response returned after its success flag has been checked. +/// +/// [`Task::send_and_wait`] and [`TryFrom`] reject responses whose +/// `success` flag is false. Validation does not inspect the output +/// text, parse JSON, or check a schema. Direct construction and deserialization +/// do not perform the success-flag check. +pub struct ValidatedResponse { + /// Original task; the API executor updates its history on successful output. + pub task: Task, + /// Token usage; the API executor reports total tokens, or zero if absent. + pub tokens: u64, + /// Resolved request model; not necessarily the model name returned by the API. + pub model: String, + /// Output text; structured output remains JSON text without schema validation. + pub contents: String, +} + +impl TryFrom for ValidatedResponse { + type Error = anyhow::Error; + + fn try_from(response: Response) -> anyhow::Result { + let Response { + task, + success, + tokens, + model, + contents, + } = response; + + anyhow::ensure!(success); + + Ok(Self { + task, + tokens, + model, + contents, + }) + } +} + #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)] /// Supported operations; only chat is currently implemented. pub enum TaskType { @@ -274,56 +382,6 @@ pub enum Interaction { User(String), } -impl Task { - /// Replaces the correlation ID; avoid duplicates among outstanding tasks. - pub fn with_id(mut self, id: u128) -> Self { - self.id = id; - self - } - - /// Sets the output limit; the API executor casts it to `u32` unchecked. - pub fn with_max_tokens(mut self, max_tokens: u64) -> Self { - self.max_tokens.replace(max_tokens); - self - } - - /// Attaches opaque caller metadata to be returned with the task. - /// - /// # Panics - /// - /// Panics if the payload cannot be serialized as JSON. - pub fn with_payload(mut self, payload: P) -> Self { - // failure serialization is considered unrecoverable from the API documentation - let payload = serde_json::to_value(&payload).expect("failed to serialize payload"); - self.payload.replace(payload); - self - } - - /// Overrides the model used for NATS routing and API execution. - pub fn with_model(mut self, model: M) -> Self { - self.model.replace(model.to_string()); - self - } - - /// Decodes contents as [`TaskChat`], returning an error for incompatible JSON. - /// - /// Does not inspect [`Self::task_type`] or validate schema/model settings. - pub fn try_to_chat(&self) -> anyhow::Result { - Ok(serde_json::from_value(self.contents.clone())?) - } - - #[cfg(feature = "std")] - pub(crate) fn model( - &self, - default_model: &std::sync::Arc>, - ) -> anyhow::Result { - self.model - .clone() - .or_else(|| default_model.as_ref().clone()) - .ok_or_else(|| anyhow::anyhow!("no model provided")) - } -} - impl Response { /// Creates a successful response without changing the task or its history. pub fn success( @@ -364,7 +422,8 @@ impl Response { } /// Decodes returned history; missing, null, or malformed history is an error. - /// The API executor currently includes the latest user prompt twice. + /// The API executor returns prior user/assistant turns, the current user + /// prompt once, and the new assistant reply; system turns are omitted. pub fn chat_history(&self) -> anyhow::Result> { let history = self .task diff --git a/src/utils.rs b/src/utils.rs index cbeb401..f1c9876 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -42,6 +42,23 @@ pub fn get_default_model() -> std::sync::Arc> { std::sync::Arc::new(default_model) } +#[cfg(any(feature = "nats-queue", feature = "jetstream-queue"))] +pub(crate) async fn nats_client_from_env_or_default() -> anyhow::Result { + let url = + std::env::var("OPENAI_API_NATS_URL").unwrap_or_else(|_| "nats://localhost:4222".into()); + let client = async_nats::ConnectOptions::new() + .request_timeout(None) + .connect(url) + .await?; + + tracing::info!( + "nats client connected to `{:?}`", + client.server_info().connect_urls + ); + + Ok(client) +} + #[test] fn id_increments() { let a = (0..1_000).map(|_| id()); diff --git a/src/worker/mod.rs b/src/worker/mod.rs index c921072..df2fc3f 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -9,25 +9,41 @@ use crate::{ #[cfg(test)] mod tests; -#[cfg(all(feature = "nats-queue", feature = "std"))] -use crate::{executor::openai::ExecutorAsyncOpenai, queue::nats::NatsWorker}; - -#[cfg(all(feature = "nats-queue", feature = "std"))] +#[cfg(any(feature = "nats-queue", feature = "jetstream-queue"))] +use crate::executor::openai::ExecutorAsyncOpenai; +#[cfg(feature = "jetstream-queue")] +use crate::queue::jetstream::JetStreamWorker; +#[cfg(feature = "nats-queue")] +use crate::queue::nats::NatsWorker; + +#[cfg(feature = "nats-queue")] /// A Core NATS consumer backed by the OpenAI-compatible chat executor. pub type NatsOpenaiWorker = Worker; +#[cfg(feature = "jetstream-queue")] +/// A durable JetStream worker backed by the OpenAI-compatible executor. +pub type JetStreamOpenaiWorker = Worker; + #[derive(Debug, Clone)] /// Connects a queue consumer to an executor, processing one task at a time. /// -/// Only `NatsOpenaiWorker::from_env_or_default` currently has a public -/// constructor (with `nats-queue` enabled). Other queue/executor combinations -/// require a caller-managed loop using their respective traits. +/// Use [`Self::new`] for any queue/executor combination, then [`Self::run`] to +/// receive tasks, execute them, and send responses through the queue. pub struct Worker { queue: Q, executor: E, } -#[cfg(all(feature = "nats-queue", feature = "std"))] +//#[cfg(feature = "jetstream-queue")] +//impl JetStreamOpenaiWorker { +// /// Binds to provisioned JetStream resources and creates the API executor. +// pub async fn from_env_or_default() -> anyhow::Result { +// let queue = JetStreamWorker::from_env_or_default().await?; +// Ok(Self::new(queue, ExecutorAsyncOpenai::from_env_or_default())) +// } +//} + +#[cfg(feature = "nats-queue")] impl NatsOpenaiWorker { /// Builds the NATS consumer and API executor from their environment settings. /// @@ -43,11 +59,18 @@ impl NatsOpenaiWorker { } impl Worker { + /// Connects a queue to an executor without starting the processing loop. + pub fn new(queue: Q, executor: E) -> Self { + Self { queue, executor } + } + /// Processes tasks sequentially until the queue returns `None` or an error. /// - /// Queue and executor errors stop the loop immediately; execution errors - /// are not converted into replies, and no application-level retry occurs. - /// Responses with `success == false` are sent normally and do not stop it. + /// Queue and executor errors stop the loop. An executor error leaves the + /// task without a response; JetStream can redeliver an unacknowledged task + /// to an active worker within its delivery and retention limits. This loop + /// does not restart itself or retry execution errors locally. + /// Responses with `success == false` are sent back and do not stop the loop. /// A polling queue returning `None` ends the loop even if more work may arrive. /// There is no shutdown signal; callers must arrange cancellation themselves. pub async fn run(self) -> anyhow::Result<()> { diff --git a/src/worker/tests.rs b/src/worker/tests.rs index 0c86f8f..4185d36 100644 --- a/src/worker/tests.rs +++ b/src/worker/tests.rs @@ -4,11 +4,11 @@ use core::time::Duration; use tokio::task::JoinHandle; -use super::Worker; use crate::{ executor::{DummyExecutor, Executor as _}, queue::{QueueProducer as _, QueueWorker as _, memory::MemoryQueue}, - task::{Interaction, Response, Task, TaskBuilder}, + task::{Interaction, Response, Task, TaskBuilder, ValidatedResponse}, + worker::Worker, }; const RESPONSE_TIMEOUT_SECS: u64 = 5; @@ -60,7 +60,7 @@ async fn receive_response(producer: &MemoryQueue, message: u128) -> Response { .expect("the response channel should remain open") } -async fn send_to_worker_with_executor(task: Task, executor: DummyExecutor) -> Response { +async fn send_to_worker_with_executor(task: Task, executor: DummyExecutor) -> ValidatedResponse { let (producer, worker_handle) = start_worker(executor); let response = task .send_and_wait(&producer, Some(RESPONSE_TIMEOUT_SECS)) @@ -73,14 +73,16 @@ async fn send_to_worker_with_executor(task: Task, executor: DummyExecutor) -> Re response.expect("the worker should return a response before the timeout") } -async fn send_to_worker(task: Task) -> Response { +async fn send_to_worker(task: Task) -> ValidatedResponse { send_to_worker_with_executor(task, DummyExecutor::default()).await } #[tokio::test] async fn worker_processes_a_memory_queue_task_with_the_default_model() { let expected_task = task(0, None); - let expected_response = Response::success(expected_task.clone(), 100, "dummy", "0"); + let expected_response = Response::success(expected_task.clone(), 100, "dummy", "0") + .try_into() + .unwrap(); assert_eq!(send_to_worker(expected_task).await, expected_response); } @@ -88,35 +90,13 @@ async fn worker_processes_a_memory_queue_task_with_the_default_model() { #[tokio::test] async fn worker_processes_a_memory_queue_task_with_its_requested_model() { let expected_task = task(2, Some("requested-model")); - let expected_response = Response::success(expected_task.clone(), 100, "requested-model", "2"); + let expected_response = Response::success(expected_task.clone(), 100, "requested-model", "2") + .try_into() + .unwrap(); assert_eq!(send_to_worker(expected_task).await, expected_response); } -#[tokio::test] -async fn worker_sends_unsuccessful_responses_and_continues_processing() { - let mut executor = DummyExecutor::default(); - executor.set_fail(); - let (producer, worker_handle) = start_worker(executor); - let first_task = task(3, None); - let second_task = task(4, Some("requested-model")); - - let first_message = producer.send_task(first_task.clone()).await.unwrap(); - let second_message = producer.send_task(second_task.clone()).await.unwrap(); - let first_response = receive_response(&producer, first_message).await; - let second_response = receive_response(&producer, second_message).await; - stop_worker(worker_handle).await; - - assert_eq!( - first_response, - Response::error(first_task, 100, "dummy", "3") - ); - assert_eq!( - second_response, - Response::error(second_task, 100, "requested-model", "4") - ); -} - #[tokio::test] async fn worker_preserves_boundary_and_optional_task_values() { let expected_task = TaskBuilder::new() @@ -135,7 +115,9 @@ async fn worker_preserves_boundary_and_optional_task_values() { .unwrap() .with_id(u128::MAX); let expected_response = - Response::success(expected_task.clone(), 100, "", u128::MAX.to_string()); + Response::success(expected_task.clone(), 100, "", u128::MAX.to_string()) + .try_into() + .unwrap(); assert_eq!(send_to_worker(expected_task).await, expected_response); } @@ -189,10 +171,14 @@ async fn worker_handles_concurrent_producer_clones() { assert_eq!( first_response.unwrap(), Response::success(expected_first_task, 100, "dummy", "20") + .try_into() + .unwrap() ); assert_eq!( second_response.unwrap(), Response::success(expected_second_task, 100, "second-model", "21") + .try_into() + .unwrap() ); }