Restore model download cancellation parity

sdk_v2/js/native/src/model.cc

@@ -9,7 +9,10 @@
 #include <foundry_local/foundry_local_c.h>
 #include <foundry_local/foundry_local_cpp.h>
 
+#include <atomic>
+#include <condition_variable>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <utility>
 #include <vector>
@@ -276,34 +279,91 @@ Napi::Value Model::Unload(const Napi::CallbackInfo& info) {
 
 namespace {
 
+class ProgressDispatch {
+ public:
+  explicit ProgressDispatch(float percent) : percent_(percent) {}
+
+  float Percent() const { return percent_; }
+
+  void Wait() {
+    std::unique_lock<std::mutex> lock(mutex_);
+    completed_cv_.wait(lock, [this] { return completed_; });
+  }
+
+  void Complete() {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      completed_ = true;
+    }
+    completed_cv_.notify_one();
+  }
+
+ private:
+  const float percent_;
+  std::mutex mutex_;
+  std::condition_variable completed_cv_;
+  bool completed_ = false;
+};
+
+class ProgressDispatchCompletion {
+ public:
+  explicit ProgressDispatchCompletion(ProgressDispatch& dispatch) : dispatch_(dispatch) {}
+  ~ProgressDispatchCompletion() { dispatch_.Complete(); }
+
+ private:
+  ProgressDispatch& dispatch_;
+};
+
 // AsyncWorker variant that drives IModel::Download with an optional JS
 // progress callback. The callback runs on the libuv worker thread; we bounce
 // each (float percent) to JS via a ThreadSafeFunction acquired before the
 // worker queues and released in OnOK/OnError.
 class DownloadWorker : public Napi::AsyncWorker {
  public:
   DownloadWorker(Napi::Env env, foundry_local::IModel* impl, Napi::ObjectReference owner,
-                 Napi::ThreadSafeFunction tsfn)
+                 Napi::ThreadSafeFunction tsfn, std::shared_ptr<std::atomic<bool>> abort_requested,
+                 Napi::ObjectReference abort_signal, Napi::FunctionReference abort_listener)
       : Napi::AsyncWorker(env),
         deferred_(Napi::Promise::Deferred::New(env)),
         impl_(impl),
         owner_(std::move(owner)),
-        tsfn_(std::move(tsfn)) {}
+        tsfn_(std::move(tsfn)),
+        abort_requested_(std::move(abort_requested)),
+        abort_signal_(std::move(abort_signal)),
+        abort_listener_(std::move(abort_listener)) {}
 
   Napi::Promise Promise() { return deferred_.Promise(); }
 
   void Execute() override {
     try {
-      auto progress_cb = tsfn_ ? std::function<int(float)>([this](float percent) {
-        // BlockingCall keeps backpressure on the worker thread: if JS is
-        // slow to drain the queue we'll wait rather than dropping reports.
-        // Callback return value is unused on the JS side; we always continue.
-        tsfn_.BlockingCall([percent](Napi::Env env, Napi::Function js_cb) {
-          js_cb.Call({Napi::Number::New(env, static_cast<double>(percent))});
-        });
-        return 0;  // 0 = continue per flProgressCallback contract.
+      const bool has_cancellation = abort_requested_ != nullptr;
+      auto progress_cb = (tsfn_ || has_cancellation) ? std::function<int(float)>([this](float percent) {
+        if (IsAbortRequested()) {
+          cancelled_by_signal_ = true;
+          return 1;
+        }
+        if (tsfn_) {
+          ProgressDispatch dispatch(percent);
+          const napi_status status = tsfn_.BlockingCall(
+              &dispatch, [](Napi::Env env, Napi::Function js_cb, ProgressDispatch* pending) {
+                ProgressDispatchCompletion completion(*pending);
+                if (env == nullptr || js_cb.IsEmpty()) {
+                  return;
+                }
+                js_cb.Call({Napi::Number::New(env, static_cast<double>(pending->Percent()))});
+              });
+          if (status != napi_ok) {
+            return 1;
+          }
+          dispatch.Wait();
+        }
+        if (IsAbortRequested()) {
+          cancelled_by_signal_ = true;
+          return 1;
+        }
+        return 0;
       })
-                               : std::function<int(float)>(nullptr);
+                                                       : std::function<int(float)>(nullptr);
       impl_->Download(std::move(progress_cb));
     } catch (const foundry_local::Error& e) {
       err_code_ = static_cast<int>(e.Code());
@@ -321,18 +381,20 @@ class DownloadWorker : public Napi::AsyncWorker {
 
   void OnOK() override {
     Napi::HandleScope scope(Env());
-    ReleaseTsfn();
+    CleanupJsReferences();
     deferred_.Resolve(Env().Undefined());
   }
 
   void OnError(const Napi::Error& /*unused*/) override {
     Napi::Env env = Env();
     Napi::HandleScope scope(env);
-    ReleaseTsfn();
+    CleanupJsReferences();
     if (tagged_) {
       Napi::Error err = Napi::Error::New(env, err_msg_);
       Napi::Object value = err.Value();
-      value.Set("name", Napi::String::New(env, "FoundryLocalError"));
+      const bool is_signal_cancellation =
+          cancelled_by_signal_ && err_code_ == FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED;
+      value.Set("name", Napi::String::New(env, is_signal_cancellation ? "AbortError" : "FoundryLocalError"));
       value.Set("code", Napi::Number::New(env, err_code_));
       deferred_.Reject(value);
     } else {
@@ -341,7 +403,21 @@ class DownloadWorker : public Napi::AsyncWorker {
   }
 
  private:
-  void ReleaseTsfn() {
+  bool IsAbortRequested() const {
+    return abort_requested_ != nullptr && abort_requested_->load(std::memory_order_acquire);
+  }
+
+  void CleanupJsReferences() {
+    if (!abort_signal_.IsEmpty() && !abort_listener_.IsEmpty()) {
+      Napi::Object signal = abort_signal_.Value();
+      Napi::Value remove_value = signal.Get("removeEventListener");
+      if (remove_value.IsFunction()) {
+        remove_value.As<Napi::Function>().Call(
+            signal, {Napi::String::New(Env(), "abort"), abort_listener_.Value()});
+      }
+      abort_listener_.Reset();
+      abort_signal_.Reset();
+    }
     if (tsfn_) {
       tsfn_.Release();
       tsfn_ = Napi::ThreadSafeFunction();
@@ -352,9 +428,13 @@ class DownloadWorker : public Napi::AsyncWorker {
   foundry_local::IModel* impl_;
   Napi::ObjectReference owner_;
   Napi::ThreadSafeFunction tsfn_;
+  std::shared_ptr<std::atomic<bool>> abort_requested_;
+  Napi::ObjectReference abort_signal_;
+  Napi::FunctionReference abort_listener_;
   std::string err_msg_;
   int err_code_ = 0;
   bool tagged_ = false;
+  bool cancelled_by_signal_ = false;
 };
 
 }  // namespace
@@ -378,8 +458,40 @@ Napi::Value Model::Download(const Napi::CallbackInfo& info) {
     return env.Undefined();
   }
 
+  std::shared_ptr<std::atomic<bool>> abort_requested;
+  Napi::ObjectReference abort_signal;
+  Napi::FunctionReference abort_listener;
+  if (info.Length() >= 2 && !info[1].IsUndefined() && !info[1].IsNull()) {
+    if (!info[1].IsObject()) {
+      Napi::TypeError::New(env, "Model.download: signal must be an AbortSignal").ThrowAsJavaScriptException();
+      return env.Undefined();
+    }
+    Napi::Object signal = info[1].As<Napi::Object>();
+    Napi::Value aborted = signal.Get("aborted");
+    Napi::Value add_value = signal.Get("addEventListener");
+    Napi::Value remove_value = signal.Get("removeEventListener");
+    if (!aborted.IsBoolean() || !add_value.IsFunction() || !remove_value.IsFunction()) {
+      Napi::TypeError::New(env, "Model.download: signal must be an AbortSignal").ThrowAsJavaScriptException();
+      return env.Undefined();
+    }
+
+    abort_requested = std::make_shared<std::atomic<bool>>(aborted.As<Napi::Boolean>().Value());
+    Napi::Function listener = Napi::Function::New(
+        env, [abort_requested](const Napi::CallbackInfo&) {
+          abort_requested->store(true, std::memory_order_release);
+        });
+    add_value.As<Napi::Function>().Call(
+        signal, {Napi::String::New(env, "abort"), listener});
+    if (signal.Get("aborted").As<Napi::Boolean>().Value()) {
+      abort_requested->store(true, std::memory_order_release);
+    }
+    abort_signal = Napi::Persistent(signal);
+    abort_listener = Napi::Persistent(listener);
+  }
+
   Napi::ObjectReference owner = Napi::Reference<Napi::Object>::New(manager_.Value(), 1);
-  auto* w = new DownloadWorker(env, impl_, std::move(owner), std::move(tsfn));
+  auto* w = new DownloadWorker(env, impl_, std::move(owner), std::move(tsfn), std::move(abort_requested),
+                               std::move(abort_signal), std::move(abort_listener));
   Napi::Promise p = w->Promise();
   w->Queue();
   return p;

sdk_v2/js/src/detail/native.ts

@@ -92,7 +92,7 @@ export interface NativeModel {
   selectVariant(variant: NativeModel): void;
   load(): Promise<void>;
   unload(): Promise<void>;
-  download(progress?: (percent: number) => void): Promise<void>;
+  download(progress?: (percent: number) => void, signal?: AbortSignal): Promise<void>;
   removeFromCache(): void;
 }
 

sdk_v2/js/src/imodel.ts

@@ -19,7 +19,7 @@ export interface IModel {
   get capabilities(): string | null;
   get supportsToolCalling(): boolean | null;
 
-  download(progressCallback?: (progress: number) => void): Promise<void>;
+  download(progressCallback?: (progress: number) => void, signal?: AbortSignal): Promise<void>;
   get path(): string;
   load(): Promise<void>;
   removeFromCache(): void;

sdk_v2/js/src/model.ts

@@ -16,6 +16,22 @@ const internalCtorKey = Symbol("Model.internal");
 
 const nativeByModel = new WeakMap<Model, NativeModel>();
 
+function isAbortSignal(value: unknown): value is AbortSignal {
+  return (
+    typeof value === "object" &&
+    value !== null &&
+    typeof (value as AbortSignal).aborted === "boolean" &&
+    typeof (value as AbortSignal).addEventListener === "function" &&
+    typeof (value as AbortSignal).removeEventListener === "function"
+  );
+}
+
+function makeAbortError(message: string): Error {
+  const error = new Error(message);
+  error.name = "AbortError";
+  return error;
+}
+
 function toDeviceType(value: NativeModelInfo["deviceType"]): DeviceType {
   switch (value) {
     case "CPU":
@@ -24,7 +40,6 @@ function toDeviceType(value: NativeModelInfo["deviceType"]): DeviceType {
       return DeviceType.GPU;
     case "NPU":
       return DeviceType.NPU;
-    case "Invalid":
     default:
       return DeviceType.Invalid;
   }
@@ -157,8 +172,18 @@ export class Model implements IModel {
     await this.#native.unload();
   }
 
-  async download(progressCallback?: (progress: number) => void): Promise<void> {
-    await this.#native.download(progressCallback);
+  async download(
+    progressCallback?: (progress: number) => void,
+    signal?: AbortSignal,
+  ): Promise<void> {
+    if (signal !== undefined && !isAbortSignal(signal)) {
+      throw new TypeError("Model.download: second argument must be an AbortSignal");
+    }
+    if (signal?.aborted === true) {
+      throw makeAbortError("Model download aborted before start");
+    }
+
+    await this.#native.download(progressCallback, signal);
   }
 
   removeFromCache(): void {

sdk_v2/js/test/model-download.types.ts

@@ -0,0 +1,23 @@
+import type { IModel } from "../src/imodel.js";
+
+declare const model: IModel;
+declare const progress: (percent: number) => void;
+declare const maybeProgress: ((percent: number) => void) | undefined;
+declare const signal: AbortSignal;
+declare const maybeSignal: AbortSignal | undefined;
+
+void model.download();
+void model.download(undefined);
+void model.download(progress);
+void model.download(maybeProgress);
+void model.download(progress, signal);
+void model.download(undefined, signal);
+void model.download(progress, maybeSignal);
+
+// Existing structural implementations remain compatible after adding the optional signal parameter.
+declare const legacyDownload: (progressCallback?: (percent: number) => void) => Promise<void>;
+const compatibleDownload: IModel["download"] = legacyDownload;
+void compatibleDownload;
+
+// @ts-expect-error AbortSignal remains the optional second argument so the original callback-first API is preserved.
+void model.download(signal);

sdk_v2/js/test/model-lifecycle.test.ts

@@ -55,6 +55,23 @@ describe.skipIf(!haveTestModelCache)("Model lifecycle (real model)", () => {
     2 * 60_000,
   );
 
+  it("download() accepts an AbortSignal and preserves completion when a cache hit wins the race", async () => {
+    const m = fixture?.model;
+    if (m === undefined) throw new Error("fixture missing");
+    const controller = new AbortController();
+
+    await expect(m.download(() => controller.abort(), controller.signal)).resolves.toBeUndefined();
+  });
+
+  it("download() rejects a pre-aborted AbortSignal before native submission", async () => {
+    const m = fixture?.model;
+    if (m === undefined) throw new Error("fixture missing");
+    const controller = new AbortController();
+    controller.abort();
+
+    await expect(m.download(undefined, controller.signal)).rejects.toMatchObject({ name: "AbortError" });
+  });
+
   it("calling load() on an already-loaded model is idempotent (or surfaces a clear error)", async () => {
     const m = fixture?.model;
     if (m === undefined) throw new Error("fixture missing");

sdk_v2/js/test/model.test.ts

@@ -4,6 +4,7 @@
 import { afterAll, beforeAll, describe, expect, it } from "vitest";
 
 import type { Catalog } from "../src/catalog.js";
+import { FlErrorCode } from "../src/detail/errors.js";
 import { Model } from "../src/model.js";
 
 import {
@@ -88,6 +89,28 @@ describeIfBuilt("Model (cache-only)", () => {
     expect(typeof model.path).toBe("string");
   });
 
+  it("download() maps AbortSignal cancellation from a native progress checkpoint to AbortError", async () => {
+    const controller = new AbortController();
+
+    await expect(
+      model.download((progress) => {
+        if (progress === 0) controller.abort();
+      }, controller.signal),
+    ).rejects.toMatchObject({
+      name: "AbortError",
+      code: FlErrorCode.OperationCancelled,
+    });
+    expect(model.isCached).toBe(false);
+  });
+
+  it("download() rejects a pre-aborted AbortSignal before native submission", async () => {
+    const controller = new AbortController();
+    controller.abort();
+
+    await expect(model.download(undefined, controller.signal)).rejects.toMatchObject({ name: "AbortError" });
+    expect(model.isCached).toBe(false);
+  });
+
   it("id and alias match info", () => {
     expect(model.id).toBe(model.info.id);
     expect(model.alias).toBe(model.info.alias);

sdk_v2/js/tsconfig.types.json

@@ -5,5 +5,5 @@
     "moduleResolution": "NodeNext",
     "noEmit": true
   },
-  "include": ["src/**/*", "test/tool-definition.types.ts"]
+  "include": ["src/**/*", "test/model-download.types.ts", "test/tool-definition.types.ts"]
 }

sdk_v2/python/README.md

@@ -110,6 +110,20 @@ with ChatSession(model) as session:
 model.unload()
 ```
 
+Pass a `threading.Event` as `cancel_event` to cancel an active download at the next native progress checkpoint:
+
+```python
+from threading import Event
+
+cancel_event = Event()
+
+def on_progress(percent: float) -> None:
+    print(f"\rDownloading: {percent:.1f}%", end="", flush=True)
+    cancel_event.set()
+
+model.download(progress_callback=on_progress, cancel_event=cancel_event)
+```
+
 Runnable end-to-end examples live under [`samples/python/`](https://github.com/microsoft/Foundry-Local/tree/main/samples/python).
 
 ## Usage