Source code
Revision control
Copy as Markdown
Other Tools
/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
#ifndef mozilla_dom_LlamaRunner_h
#define mozilla_dom_LlamaRunner_h
// Primary WebIDL interface for llama.cpp-based streaming chat model
// integration. Exposes LLM-backed ReadableStream generation and chat prompt
// formatting to JavaScript via LlamaRunner.
#include "MediaInfo.h"
#include "mozilla/dom/LlamaRunnerBinding.h"
#include "mozilla/dom/MessagePort.h"
#include "mozilla/dom/ReadableStream.h"
#include "nsIFileStreams.h"
#include "nsISupports.h"
#include "nsWrapperCache.h"
#include "mozilla/llama/LlamaBackend.h"
#include "prsystem.h"
#include "mozilla/Casting.h"
#include "mozilla/SPSCQueue.h"
#include "mozilla/Array.h"
#include <cstdint>
#include "mozilla/dom/Promise.h"
#include "mozilla/dom/ReadableStream.h"
#include "mozilla/dom/ReadableStreamDefaultController.h"
#include "mozilla/dom/UnderlyingSourceCallbackHelpers.h"
#include "nsThreadUtils.h"
#include "mozilla/TaskQueue.h"
#include "mozilla/ScopeExit.h"
#include "mozilla/RefPtr.h"
#include "nsIThread.h"
#include "nsError.h"
#include "mozilla/dom/ContentChild.h"
#include "mozilla/Atomics.h"
#include "mozilla/llama/LlamaBackend.h"
#include "mozilla/Logging.h"
#include "nsFmtString.h"
#include "nsThread.h"
#include "nsThreadManager.h"
#include "mozilla/ThreadEventQueue.h"
#include "mozilla/EventQueue.h"
#include "nsDeque.h"
#include "mozilla/dom/Blob.h"
namespace mozilla::dom {
class LlamaStreamSource;
}
namespace mozilla::llama {
// When the Maybe has no value, it indicates the model has finished generating.
// Otherwise, it contains a LlamaChatResponse with partial or final content.
using LlamaGenerateTaskPromise =
MozPromise<mozilla::Maybe<LlamaChatResponse>, nsCString,
/* IsExclusive = */ true>;
/**
* LlamaGenerateTask runs the orchestration of model inference on a background
* thread, but delegates all compute-intensive operations (e.g., token
* generation, prompt processing, decoding) to the LlamaBackend's internal
* threadpool.
*
* The task itself manages:
* - Calling LlamaBackend::Generate with user options and callbacks.
* - Buffering responses and applying streaming heuristics (e.g., flush size,
* phase boundaries).
* - Monitoring for cancellation requests.
*
* This task is launched during the first JS stream pull and executes
* independently. It posts all intermediate/final results back to the
* LlamaStreamSource on the JS stream thread.
*
* Note: This class does not do the heavy lifting; it schedules and collects
* results from backend-owned threads. This enables concurrency across multiple
* LlamaRunner instances with minimal blocking.
*/
class LlamaGenerateTask final : public mozilla::CancelableRunnable {
public:
using LlamaChatResponse = mozilla::dom::LlamaChatResponse;
enum class TaskState : uint32_t {
// Task has not started yet.
Idle,
// Task is actively running.
Running,
// Task completed successfully.
CompletedSuccess,
// Task failed due to an error.
CompletedFailure,
// Task was externally cancelled (if not already completed).
Cancelled
};
LlamaGenerateTask(RefPtr<LlamaBackend> aBackend,
const LlamaChatOptions& aOptions);
NS_IMETHOD Run() override;
nsresult Cancel() override;
// Returns the next message if available, or a promise that will resolve once
// a message is ready. Rejects immediately if the task has failed.
RefPtr<LlamaGenerateTaskPromise> GetMessage();
private:
// Attempts to push a message only if a consumer is actively waiting.
// Returns true if the message was consumed immediately or queued to resolve a
// pending promise.
bool MaybePushMessage(mozilla::Maybe<LlamaChatResponse> aMessage);
// Unconditionally pushes a message. First tries MaybePushMessage(); if that
// fails, enqueues the message into the internal queue. Returns true if the
// message was accepted.
bool PushMessage(mozilla::Maybe<LlamaChatResponse> aMessage);
// Thread-safe task state
mozilla::Atomic<TaskState> mState{TaskState::Idle};
// Error message set if the task fails. We store it separately instead of
// using the message-passing queue, so that errors can still be surfaced even
// if the queue mechanism fails.
nsCString mErrorMessage;
// Shared reference to backend (not exclusive owner)
RefPtr<LlamaBackend> mBackend;
const LlamaChatOptions mChatOptions;
// Index of the currently active promise holder (0 or 1).
// The producer resolves the promise at this index, then toggles it.
// The consumer uses the *other* index to safely create a new promise.
int mCurrentPromiseHolderIdx{0};
// Double-buffered promise holders (index toggles between 0 and 1) to avoid a
// rare race where a resolved promise might be consumed by the other thread
// *before* we’ve officially marked it as resolved. While this race is highly
// unlikely (would require the consumer thread to request a new promise
// between the resolve() call and its internal state update), this design
// fully eliminates the possibility. The producer resolves the current index,
// then switches to the next; the consumer always creates new promises from
// the non-current index.
Array<MozPromiseHolder<LlamaGenerateTaskPromise>, 2> mPromiseHolders;
// Thread-safe flag indicating whether a consumer is waiting for data.
Atomic<bool> mHasPendingConsumer{false};
// Thread-safe buffer for messages to be sent back to the consumer.
SPSCQueue<mozilla::Maybe<LlamaChatResponse>> mMessagesQueue;
~LlamaGenerateTask();
};
} // namespace mozilla::llama
namespace mozilla::dom {
using LlamaBackend = ::mozilla::llama::LlamaBackend;
/**
* LlamaStreamSource is a bridge between LlamaGenerateTask and JS
* ReadableStream.
*
* Implements UnderlyingSourceAlgorithmsWrapper so it can be used with
* ReadableStream::CreateNative. This object owns the lifecycle of the
* generation task (LlamaGenerateTask) and buffers intermediate results
* for consumption by JS pull callbacks.
*
* It holds a shared strong reference to the backend (LlamaBackend),
* which may also be retained by other components (e.g., LlamaRunner).
* All compute-heavy work is performed by the backend’s internal threadpool.
*
* Generation results are delivered via LlamaGenerateTask::Generate(), which
* returns a promise. Once resolved, the result is forwarded to the JS consumer.
*
* The stream starts when PullCallbackImpl is first called from JS,
* launching a background generation task and associating it with a thread.
*/
class LlamaStreamSource final : public UnderlyingSourceAlgorithmsWrapper {
public:
MOZ_DECLARE_REFCOUNTED_TYPENAME(LlamaStreamSource)
NS_DECL_ISUPPORTS_INHERITED
NS_DECL_CYCLE_COLLECTION_CLASS_INHERITED(LlamaStreamSource,
UnderlyingSourceAlgorithmsWrapper)
LlamaStreamSource(RefPtr<LlamaBackend> aBackend,
const LlamaChatOptions& aOptions);
MOZ_CAN_RUN_SCRIPT
already_AddRefed<Promise> CancelCallbackImpl(
JSContext* aCx, const Optional<JS::Handle<JS::Value>>& aReason,
ErrorResult& aRv) override;
MOZ_CAN_RUN_SCRIPT already_AddRefed<Promise> PullCallbackImpl(
JSContext* aCx, ReadableStreamControllerBase& aController,
ErrorResult& aRv) override;
// Links the JS-side stream controller to this source
void SetControllerStream(RefPtr<ReadableStream> aStream);
private:
~LlamaStreamSource();
RefPtr<LlamaBackend> mBackend;
const LlamaChatOptions mChatOptions;
// Background generation task
RefPtr<mozilla::llama::LlamaGenerateTask> mTask;
// Background worker thread
nsCOMPtr<nsIThread> mGenerateThread;
// Holds the event queue where PullCallbackImpl is
// called from
nsCOMPtr<nsISerialEventTarget> mOriginalEventTarget;
// Associated JS stream object
RefPtr<ReadableStream> mControllerStream;
};
class MetadataCallback;
/**
* LlamaRunner is the primary WebIDL-exposed controller for llama.cpp chat.
*
* It provides JavaScript with an API to format prompts, launch inference, and
* receive output as a `ReadableStream`. It delegates inference to a
* thread-safe LlamaBackend and manages stream logic via LlamaStreamSource.
*
* This class is designed for use in JS.
*/
class LlamaRunner final : public nsISupports, public nsWrapperCache {
public:
MOZ_DECLARE_REFCOUNTED_TYPENAME(LlamaRunner)
NS_DECL_CYCLE_COLLECTING_ISUPPORTS
NS_DECL_CYCLE_COLLECTION_WRAPPERCACHE_CLASS(LlamaRunner)
explicit LlamaRunner(const GlobalObject& aGlobal);
nsISupports* GetParentObject() const { return mGlobal; }
static already_AddRefed<LlamaRunner> Constructor(const GlobalObject& aGlobal,
ErrorResult& aRv);
already_AddRefed<Promise> Initialize(const LlamaModelOptions& aOptions,
Blob& aModelBlob, ErrorResult& aRv);
virtual JSObject* WrapObject(JSContext* aCx,
JS::Handle<JSObject*> aGivenProto) override;
/**
* Creates a readable stream that incrementally yields language model
* responses.
*
* This function initiates a new generation session using the provided options
* and returns a JavaScript `ReadableStream`. The stream will asynchronously
* emit `LlamaChatResponse` chunks as the model produces output.
*
* @param aOptions Configuration for the generation session.
* @param aRv Used to report any errors during stream setup or execution.
*
* @returns A `ReadableStream` that yields `LlamaChatResponse` objects from
* the language model, suitable for consumption in JavaScript via async
* iteration or stream readers.
*
* @note This function is designed for use in JavaScript via WebIDL. It
* supports streaming output for real-time use cases such as chat UIs or
* progressive rendering.
*
* @example JavaScript usage:
* const stream = CreateGenerationStream(chatOptions);
* const reader = stream.getReader();
*
* while (true) {
* const { value, done } = await reader.read();
* if (done) break;
* console.log(value); // `value` is a LlamaChatResponse
* }
*/
already_AddRefed<ReadableStream> CreateGenerationStream(
const LlamaChatOptions& aOptions, ErrorResult& aRv);
/**
* Formats a sequence of chat messages into a prompt string for LLM inference.
*
* This function takes a structured list of chat messages (user, assistant, or
* system roles) and formats them into a prompt string suitable for processing
* by a llama.cpp-based language model. The function is asynchronous and
* returns a JavaScript Promise that resolves to the formatted string.
*
* @param aOptions An object containing the input messages and formatting
* options.
*
* @param aRv Used to report any errors that occur during formatting.
*
* @returns A Promise that resolves to the final prompt string formatted for
* LLM input.
*
* @throws DOMException via `aRv` on failure (e.g., invalid message roles or
* empty input).
*
* @example JavaScript usage:
* FormatChat({
* messages: [
* { role: "user", content: "What's the weather like?" },
* { role: "assistant", content: "It's sunny and 25°C." }
* ],
* addAssistant: true
* }).then(prompt => {
* // Pass prompt to LLM for inference
* });
*/
already_AddRefed<Promise> FormatChat(const LlamaFormatChatOptions& aOptions,
ErrorResult& aRv);
static bool InInferenceProcess(JSContext*, JSObject*);
void OnMetadataReceived();
private:
~LlamaRunner() = default;
RefPtr<LlamaBackend> mBackend;
RefPtr<LlamaStreamSource> mStreamSource;
protected:
LlamaModelOptions mModelOptions;
nsCOMPtr<nsIGlobalObject> mGlobal;
RefPtr<Promise> mInitPromise;
nsCOMPtr<nsIInputStream> mStream;
RefPtr<MetadataCallback> mMetadataCallback;
};
} // namespace mozilla::dom
#endif