Compare commits
15 commits
main
...
pr/daniel-
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d994800bf3 |
||
|
|
a598409f88 | ||
|
|
edea3ef5de | ||
|
|
545a62174d | ||
|
|
36f896a528 | ||
|
|
506cdc6c8a | ||
|
|
e43304a1b9 | ||
|
|
b9f10d484f | ||
|
|
2735cec36c | ||
|
|
751fa0bfaf | ||
|
|
26b341bde6 | ||
|
|
91664ccd62 | ||
|
|
715727b64b | ||
|
|
96774c4974 | ||
|
|
314627cd6a |
12 changed files with 2980 additions and 221 deletions
|
|
@ -86,6 +86,13 @@ Replace `claude` with any supported agent:
|
|||
| OpenCode | `unsloth start opencode` |
|
||||
| Pi Coding Agent | `unsloth start pi` |
|
||||
|
||||
Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local
|
||||
subagent:
|
||||
|
||||
```bash
|
||||
unsloth start claude --as-subagent --model unsloth/model-GGUF:quant
|
||||
```
|
||||
|
||||
## 📥 Install
|
||||
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
|
||||
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"}
|
|||
include-package-data = true
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
unsloth_cli = ["codex_fallback_prompt.md"]
|
||||
unsloth_cli = ["codex_fallback_prompt.md", "pi_subagent.ts"]
|
||||
studio = [
|
||||
"*.sh",
|
||||
"*.ps1",
|
||||
|
|
|
|||
|
|
@ -3406,9 +3406,9 @@ async def _acquire_swap_gate() -> None:
|
|||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
# Counts in-flight auto-switch requests per (target, variant). The busy guard
|
||||
# subtracts same-target waiters so concurrent requests for one model load once
|
||||
# instead of each 409-ing the other.
|
||||
# Counts auto-switch requests waiting to load each (target, variant). These
|
||||
# requests are queued for a swap but are not generating, so the drain wait below
|
||||
# excludes them from the active inference count.
|
||||
_auto_switch_waiters: dict[tuple[str, str], int] = {}
|
||||
_auto_switch_waiters_guard = threading.Lock()
|
||||
|
||||
|
|
@ -3426,35 +3426,31 @@ def _note_switch_waiter(key: tuple[str, str], delta: int) -> None:
|
|||
_auto_switch_waiters.pop(key, None)
|
||||
|
||||
|
||||
def _same_target_waiters(key: tuple[str, str]) -> int:
|
||||
def _switch_waiter_count() -> int:
|
||||
with _auto_switch_waiters_guard:
|
||||
return _auto_switch_waiters.get(key, 0)
|
||||
return sum(max(0, count) for count in _auto_switch_waiters.values())
|
||||
|
||||
|
||||
# A second waiter map keyed by the raw requested model, registered before the
|
||||
# (slow) resolve. The middleware counts a concurrent same-model request as
|
||||
# in-flight before it resolves and joins _auto_switch_waiters, so without this
|
||||
# the first request would see it as an unrelated request and 409.
|
||||
_auto_switch_request_waiters: dict[str, int] = {}
|
||||
_auto_switch_request_waiters_guard = threading.Lock()
|
||||
async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None:
|
||||
"""Wait until a model replacement cannot interrupt active inference.
|
||||
|
||||
|
||||
def _request_waiter_key(requested_model: str) -> str:
|
||||
return requested_model.strip().lower()
|
||||
|
||||
|
||||
def _note_request_waiter(key: str, delta: int) -> None:
|
||||
with _auto_switch_request_waiters_guard:
|
||||
n = _auto_switch_request_waiters.get(key, 0) + delta
|
||||
if n > 0:
|
||||
_auto_switch_request_waiters[key] = n
|
||||
else:
|
||||
_auto_switch_request_waiters.pop(key, None)
|
||||
|
||||
|
||||
def _same_request_waiters(key: str) -> int:
|
||||
with _auto_switch_request_waiters_guard:
|
||||
return _auto_switch_request_waiters.get(key, 0)
|
||||
The caller holds ``inference_lifecycle_gate``, which prevents new inference
|
||||
from starting while existing requests drain. Auto-switch requests that have
|
||||
resolved their targets are scheduler waiters, not active generations, so
|
||||
exclude them to avoid a queue deadlock.
|
||||
"""
|
||||
from core.inference.llama_keepwarm import other_inference_request_count
|
||||
while True:
|
||||
queued_switches = _switch_waiter_count()
|
||||
if current_request_counted and queued_switches > 0:
|
||||
queued_switches -= 1
|
||||
active_others = other_inference_request_count(
|
||||
current_request_counted = current_request_counted,
|
||||
include_pending = False,
|
||||
)
|
||||
if active_others <= queued_switches:
|
||||
return
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]:
|
||||
|
|
@ -3582,7 +3578,6 @@ async def _maybe_auto_switch_model(
|
|||
from core.inference.local_model_resolver import resolve_local_gguf
|
||||
from core.inference.llama_keepwarm import (
|
||||
get_last_unloaded_model,
|
||||
other_inference_request_count,
|
||||
inference_lifecycle_gate,
|
||||
)
|
||||
|
||||
|
|
@ -3603,12 +3598,7 @@ async def _maybe_auto_switch_model(
|
|||
if not auto_switch_on and get_auto_unload_idle_seconds() <= 0:
|
||||
return
|
||||
|
||||
# Register by the raw requested model before resolving (which can be slow):
|
||||
# the middleware already counts a concurrent same-model request as in-flight,
|
||||
# so the busy guard must know it shares this target even while it resolves.
|
||||
request_key = _request_waiter_key(requested_model)
|
||||
_note_request_waiter(request_key, 1)
|
||||
try:
|
||||
async def _resolve_and_switch() -> None:
|
||||
# Off the loop: a cold-cache rebuild walks several model dirs + HF caches.
|
||||
# With auto-switch off (or an omitted-model reload-only request), skip the
|
||||
# resolve so only the reload-stash path runs and no name is ever matched.
|
||||
|
|
@ -3718,31 +3708,6 @@ async def _maybe_auto_switch_model(
|
|||
if _already_serving():
|
||||
_record_serving_alias()
|
||||
return
|
||||
# Single slot: refuse a cross-model swap while another inference
|
||||
# request is active rather than killing its response. Requests
|
||||
# heading to this same target (by resolved id or raw name) are
|
||||
# excluded, so concurrent requests for one model load once. A
|
||||
# pending request is still in the middleware, not generating, so
|
||||
# it is not counted here.
|
||||
same_others = max(
|
||||
_same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0
|
||||
)
|
||||
others = other_inference_request_count(
|
||||
current_request_counted = True, include_pending = False
|
||||
)
|
||||
# Not gated on the GGUF being loaded: _load_model_impl also
|
||||
# tears down an active Unsloth backend before loading a GGUF,
|
||||
# so refuse whenever any other inference request is in flight.
|
||||
if others > same_others:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = openai_error_body(
|
||||
"Cannot switch models while another inference request is in progress.",
|
||||
status = 409,
|
||||
code = "model_switch_busy",
|
||||
param = "model",
|
||||
),
|
||||
)
|
||||
# Apply this model's saved launch flags so the swap honors the config.
|
||||
override = get_model_override(override_id)
|
||||
load_kwargs = {"model_path": target_id, "gguf_variant": variant}
|
||||
|
|
@ -3757,6 +3722,7 @@ async def _maybe_auto_switch_model(
|
|||
LoadRequest(**load_kwargs),
|
||||
fastapi_request,
|
||||
current_subject,
|
||||
current_request_counted = True,
|
||||
)
|
||||
# Advertise the repo id (not the concrete load path) as the loaded
|
||||
# model's public id and override key for /v1/models and idle stash.
|
||||
|
|
@ -3765,8 +3731,8 @@ async def _maybe_auto_switch_model(
|
|||
_auto_switch_process_lock.release()
|
||||
finally:
|
||||
_note_switch_waiter(key, -1)
|
||||
finally:
|
||||
_note_request_waiter(request_key, -1)
|
||||
|
||||
await _resolve_and_switch()
|
||||
|
||||
|
||||
async def _auto_switch_from_request_body(request: Request, current_subject: str):
|
||||
|
|
@ -4186,6 +4152,15 @@ def _maybe_unsupported_message(msg: str) -> str:
|
|||
return msg
|
||||
|
||||
|
||||
def _raise_if_sidecar_swap_in_progress() -> None:
|
||||
from utils.transformers_version import sidecar_swap_in_progress
|
||||
if sidecar_swap_in_progress():
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "A transformers installation is in progress. Retry when it completes.",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/load", response_model = LoadResponse)
|
||||
async def load_model(
|
||||
request: LoadRequest,
|
||||
|
|
@ -4206,24 +4181,23 @@ async def load_model(
|
|||
# install can reserve while this request queues on the gate, so the pre-gate
|
||||
# check alone is only a fast path.
|
||||
from core.inference.llama_keepwarm import inference_lifecycle_gate
|
||||
from utils.transformers_version import sidecar_swap_in_progress
|
||||
|
||||
_swap_409 = HTTPException(
|
||||
status_code = 409,
|
||||
detail = "A transformers installation is in progress. Retry when it completes.",
|
||||
)
|
||||
if sidecar_swap_in_progress():
|
||||
raise _swap_409
|
||||
_raise_if_sidecar_swap_in_progress()
|
||||
# Hold the lifecycle gate across the load so idle auto-unload can't unload the
|
||||
# model mid-load. Auto-switch calls _load_model_impl directly since it already
|
||||
# holds this gate.
|
||||
async with inference_lifecycle_gate():
|
||||
if sidecar_swap_in_progress():
|
||||
raise _swap_409
|
||||
_raise_if_sidecar_swap_in_progress()
|
||||
return await _load_model_impl(request, fastapi_request, current_subject)
|
||||
|
||||
|
||||
async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str):
|
||||
async def _load_model_impl(
|
||||
request: LoadRequest,
|
||||
fastapi_request: Request,
|
||||
current_subject: str,
|
||||
*,
|
||||
current_request_counted: bool = False,
|
||||
):
|
||||
from core.inference.llama_cpp import LlamaServerNotFoundError
|
||||
|
||||
# A new load starts here; arm the progress throttle so this load's first
|
||||
|
|
@ -4557,6 +4531,14 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
),
|
||||
)
|
||||
|
||||
# Keep the resident model alive until every active generation has
|
||||
# finished. The lifecycle gate held by the caller blocks new starts.
|
||||
await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
|
||||
# The installer reserves its sidecar swap before waiting on this gate.
|
||||
# It can do so while active inference drains, after the route-level
|
||||
# checks above, so honor that reservation before replacing either backend.
|
||||
_raise_if_sidecar_swap_in_progress()
|
||||
|
||||
# Unload any active Unsloth model only after every hub conflict check.
|
||||
if unsloth_backend.active_model_name:
|
||||
logger.info(
|
||||
|
|
@ -4767,6 +4749,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
|
|||
|
||||
# Unload any active GGUF model first
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
|
||||
_raise_if_sidecar_swap_in_progress()
|
||||
if llama_backend.is_loaded:
|
||||
logger.info("Unloading GGUF model before loading Unsloth model")
|
||||
llama_backend.unload_model()
|
||||
|
|
@ -7096,7 +7080,7 @@ async def openai_chat_completions(
|
|||
if payload.provider_id or payload.provider_type:
|
||||
# External provider: this request won't touch the local GGUF, so drop it
|
||||
# from the keep-warm count or its in-flight stream would falsely block a
|
||||
# concurrent local auto-switch with model_switch_busy.
|
||||
# concurrent local model switch from proceeding.
|
||||
from core.inference.llama_keepwarm import untrack_current_request
|
||||
|
||||
untrack_current_request(request.scope)
|
||||
|
|
|
|||
|
|
@ -68,7 +68,13 @@ class _LoadRecorder:
|
|||
request,
|
||||
fastapi_request,
|
||||
current_subject = None,
|
||||
*,
|
||||
current_request_counted = False,
|
||||
):
|
||||
# Mirror the production load boundary before recording any replacement.
|
||||
await inference_route._wait_for_model_switch_idle(
|
||||
current_request_counted = current_request_counted
|
||||
)
|
||||
self.calls.append(request)
|
||||
if self.fail:
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -94,7 +100,6 @@ def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder):
|
|||
# gate that auto-switch already owns, so it calls the impl directly).
|
||||
monkeypatch.setattr(inference_route, "_load_model_impl", recorder)
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
|
||||
|
||||
|
||||
def _run_hook(model = "some/model"):
|
||||
|
|
@ -1205,10 +1210,9 @@ def test_middleware_ignores_non_post(monkeypatch):
|
|||
# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ──
|
||||
|
||||
|
||||
def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
|
||||
# A cross-model swap must 409 (not kill) while another inference request is in
|
||||
# flight; the requesting call itself is excluded from the count.
|
||||
from fastapi import HTTPException
|
||||
def test_auto_switch_waits_for_another_inference_to_finish(monkeypatch):
|
||||
# A cross-model swap queues while another request is generating, then loads
|
||||
# after that request drains. The requesting call itself is excluded.
|
||||
from core.inference import llama_keepwarm as kw
|
||||
|
||||
backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M")
|
||||
|
|
@ -1222,10 +1226,18 @@ def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
|
|||
)
|
||||
monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one
|
||||
monkeypatch.setattr(kw, "_pending", 0)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert exc.value.status_code == 409
|
||||
assert rec.calls == []
|
||||
|
||||
async def _drive():
|
||||
task = asyncio.create_task(
|
||||
inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
assert rec.calls == []
|
||||
kw._note_end() # the other generation finishes; this request remains counted
|
||||
await asyncio.wait_for(task, timeout = 1)
|
||||
|
||||
asyncio.run(_drive())
|
||||
assert len(rec.calls) == 1
|
||||
|
||||
|
||||
def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch):
|
||||
|
|
@ -1411,13 +1423,12 @@ def test_concurrent_same_target_requests_load_once(monkeypatch):
|
|||
monkeypatch.setattr(kw, "_pending", 0)
|
||||
inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert len(rec.calls) == 1 # loads once, no 409
|
||||
assert len(rec.calls) == 1
|
||||
|
||||
|
||||
def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch):
|
||||
# A concurrent request heading to a different target still blocks the swap: the
|
||||
# same-target exclusion must not swallow a genuinely conflicting request.
|
||||
from fastapi import HTTPException
|
||||
def test_queued_different_target_does_not_deadlock_current_swap(monkeypatch):
|
||||
# A concurrent request already queued for another target is not generating,
|
||||
# so it must not prevent the current serialized swap from proceeding.
|
||||
from core.inference import llama_keepwarm as kw
|
||||
|
||||
backend = _FakeBackend("org/A-GGUF")
|
||||
|
|
@ -1432,10 +1443,8 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat
|
|||
monkeypatch.setattr(kw, "_inflight", 2)
|
||||
monkeypatch.setattr(kw, "_pending", 0)
|
||||
inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert exc.value.status_code == 409
|
||||
assert rec.calls == []
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert len(rec.calls) == 1
|
||||
|
||||
|
||||
def test_v1_models_advertises_repo_id_not_load_path(monkeypatch):
|
||||
|
|
@ -1481,6 +1490,25 @@ def test_load_route_holds_lifecycle_gate(monkeypatch):
|
|||
assert "_load_model_impl" in src
|
||||
|
||||
|
||||
def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded():
|
||||
# Both replacement directions drain active inference, then recheck whether a
|
||||
# sidecar install reserved the lifecycle gate during that wait. Exact-model
|
||||
# reuse exits earlier, so an already-loaded model never waits on unrelated inference.
|
||||
import inspect
|
||||
|
||||
src = inspect.getsource(inference_route._load_model_impl)
|
||||
gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:"))
|
||||
gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait)
|
||||
unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait)
|
||||
standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1)
|
||||
standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
|
||||
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
|
||||
already_loaded = src.index('status = "already_loaded"')
|
||||
|
||||
assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth
|
||||
assert standard_wait < standard_sidecar_check < unload_gguf
|
||||
|
||||
|
||||
def _anthropic_payload(max_tokens = None):
|
||||
from models.inference import AnthropicMessagesRequest, AnthropicMessage
|
||||
return AnthropicMessagesRequest(
|
||||
|
|
@ -1519,9 +1547,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch):
|
|||
# ── review round 6: concurrency ordering, external untrack, unload gate, ids ──
|
||||
|
||||
|
||||
def test_pending_same_target_request_does_not_force_409(monkeypatch):
|
||||
def test_pending_same_target_request_does_not_block_swap(monkeypatch):
|
||||
# A second same-target request blocked in the middleware (pending, not yet
|
||||
# generating) must not make the first request 409: pending is excluded.
|
||||
# generating) must not block the first request: pending is excluded.
|
||||
from core.inference import llama_keepwarm as kw
|
||||
|
||||
backend = _FakeBackend("org/A-GGUF")
|
||||
|
|
@ -1536,13 +1564,13 @@ def test_pending_same_target_request_does_not_force_409(monkeypatch):
|
|||
monkeypatch.setattr(kw, "_inflight", 1) # just the caller
|
||||
monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert len(rec.calls) == 1 # loads once, no 409
|
||||
assert len(rec.calls) == 1
|
||||
|
||||
|
||||
def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch):
|
||||
def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch):
|
||||
# The real middleware counts a concurrent same-model request as in-flight
|
||||
# before it resolves and registers a target waiter. The raw-request waiter,
|
||||
# registered before resolve, must still exclude it so the first request loads.
|
||||
# before it resolves and registers a target waiter. Treat it as active until
|
||||
# its target is known, then recognize it as another queued switch request.
|
||||
from core.inference import llama_keepwarm as kw
|
||||
|
||||
backend = _FakeBackend("org/A-GGUF")
|
||||
|
|
@ -1556,10 +1584,20 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat
|
|||
)
|
||||
monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin
|
||||
monkeypatch.setattr(kw, "_pending", 0)
|
||||
# The twin has only registered its raw requested model (not yet a target waiter).
|
||||
inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1)
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert len(rec.calls) == 1 # loads once, no 409
|
||||
# The twin is still resolving, so it is counted in-flight but has not joined
|
||||
# the concrete target queue yet.
|
||||
|
||||
async def _drive():
|
||||
task = asyncio.create_task(
|
||||
inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
assert rec.calls == []
|
||||
inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
|
||||
await asyncio.wait_for(task, timeout = 1)
|
||||
|
||||
asyncio.run(_drive())
|
||||
assert len(rec.calls) == 1
|
||||
|
||||
|
||||
def test_external_untrack_decrements_inflight_and_is_idempotent():
|
||||
|
|
@ -1595,11 +1633,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch):
|
|||
assert not backend.is_loaded # torn down despite the active request
|
||||
|
||||
|
||||
def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
|
||||
def test_auto_switch_waits_when_unsloth_stream_active(monkeypatch):
|
||||
# The GGUF slot is empty but an Unsloth model is streaming (counted in-flight).
|
||||
# _load_model_impl would unload it, so auto-switch must 409, not only when a
|
||||
# GGUF is loaded.
|
||||
from fastapi import HTTPException
|
||||
# The replacement waits for it just as it does for a GGUF generation.
|
||||
from core.inference import llama_keepwarm as kw
|
||||
|
||||
backend = _FakeBackend(None) # no GGUF loaded
|
||||
|
|
@ -1613,10 +1649,18 @@ def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
|
|||
)
|
||||
monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request
|
||||
monkeypatch.setattr(kw, "_pending", 0)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_run_hook("org/B-GGUF:Q8_0")
|
||||
assert exc.value.status_code == 409
|
||||
assert rec.calls == [] # the active Unsloth model is not torn down
|
||||
|
||||
async def _drive():
|
||||
task = asyncio.create_task(
|
||||
inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
assert rec.calls == []
|
||||
kw._note_end()
|
||||
await asyncio.wait_for(task, timeout = 1)
|
||||
|
||||
asyncio.run(_drive())
|
||||
assert len(rec.calls) == 1
|
||||
|
||||
|
||||
def test_public_model_id_prefers_advertised_over_path():
|
||||
|
|
@ -3097,6 +3141,8 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
|
|||
request,
|
||||
fastapi_request,
|
||||
current_subject = None,
|
||||
*,
|
||||
current_request_counted = False,
|
||||
):
|
||||
with slock:
|
||||
state["cur"] += 1
|
||||
|
|
@ -3114,7 +3160,6 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
|
|||
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
|
||||
monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load)
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
|
||||
monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
|
||||
|
||||
barrier = threading.Barrier(2)
|
||||
|
||||
|
|
|
|||
366
unsloth_cli/claude_subagent_mcp.py
Normal file
366
unsloth_cli/claude_subagent_mcp.py
Normal file
|
|
@ -0,0 +1,366 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Small stdio MCP bridge from cloud Claude Code to a local Claude Code child."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
|
||||
from unsloth_cli.commands.start import (
|
||||
_CLAUDE_ENV_UNSET,
|
||||
_SUBAGENT_DESCRIPTION,
|
||||
_SUBAGENT_INSTRUCTIONS,
|
||||
_claude_flags,
|
||||
_claude_local_env,
|
||||
_wsl_shim_env,
|
||||
)
|
||||
|
||||
_MAX_RESULT_CHARACTERS = 100_000
|
||||
_CANCEL_POLL_SECONDS = 0.1
|
||||
_CANCEL_GRACE_SECONDS = 2.0
|
||||
|
||||
|
||||
def _required_env(name: str) -> str:
|
||||
value = os.environ.get(name, "").strip()
|
||||
if not value:
|
||||
raise RuntimeError(f"Missing {name}.")
|
||||
return value
|
||||
|
||||
|
||||
def _bounded(text: str) -> str:
|
||||
if len(text) <= _MAX_RESULT_CHARACTERS:
|
||||
return text
|
||||
return text[:_MAX_RESULT_CHARACTERS] + "\n\n[Local agent output truncated]"
|
||||
|
||||
|
||||
def _result_text(stdout: str) -> str:
|
||||
lines = [line for line in stdout.splitlines() if line.strip()]
|
||||
candidates = [stdout.strip(), *reversed(lines)]
|
||||
for candidate in candidates:
|
||||
try:
|
||||
payload = json.loads(candidate)
|
||||
except ValueError:
|
||||
continue
|
||||
if not isinstance(payload, dict):
|
||||
continue
|
||||
result = payload.get("result")
|
||||
if payload.get("is_error"):
|
||||
raise RuntimeError(str(result or "The local Claude agent failed."))
|
||||
if isinstance(result, str) and result.strip():
|
||||
return _bounded(result.strip())
|
||||
raise RuntimeError("The local Claude agent returned no readable result.")
|
||||
|
||||
|
||||
def _stop_child(process: subprocess.Popen) -> None:
|
||||
"""Stop the Claude child and any tool processes it started."""
|
||||
if process.poll() is not None:
|
||||
if os.name != "nt":
|
||||
# Leader exited, but its tool processes may still be running.
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
except OSError:
|
||||
return
|
||||
time.sleep(_CANCEL_GRACE_SECONDS)
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
return
|
||||
if os.name == "nt":
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
["taskkill", "/PID", str(process.pid), "/T", "/F"],
|
||||
capture_output = True,
|
||||
timeout = 15,
|
||||
check = False,
|
||||
)
|
||||
except Exception:
|
||||
completed = None
|
||||
# A failed taskkill must not leave the child running through the grace wait.
|
||||
if (completed is None or completed.returncode != 0) and process.poll() is None:
|
||||
process.terminate()
|
||||
else:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
except OSError:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout = _CANCEL_GRACE_SECONDS)
|
||||
except subprocess.TimeoutExpired:
|
||||
if os.name == "nt":
|
||||
process.kill()
|
||||
else:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
process.kill()
|
||||
process.wait()
|
||||
else:
|
||||
if os.name != "nt":
|
||||
# Leader is gone; kill any surviving group members.
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> str:
|
||||
base = _required_env("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL")
|
||||
key = _required_env("UNSLOTH_CLAUDE_SUBAGENT_API_KEY")
|
||||
model = _required_env("UNSLOTH_CLAUDE_SUBAGENT_MODEL")
|
||||
window = int(os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW", "0") or 0)
|
||||
entry = {"id": model, "context_length": window}
|
||||
local_env = _claude_local_env(base, key, entry)
|
||||
child_env = dict(os.environ)
|
||||
|
||||
executable = shutil.which("claude")
|
||||
if executable is None:
|
||||
raise RuntimeError("`claude` is not installed or is not on PATH.")
|
||||
cancel_event = cancel_event or threading.Event()
|
||||
if cancel_event.is_set():
|
||||
raise RuntimeError("The local Claude agent was cancelled.")
|
||||
command = [
|
||||
"claude",
|
||||
"--model",
|
||||
model,
|
||||
*_claude_flags(model),
|
||||
"--permission-mode",
|
||||
(
|
||||
"bypassPermissions"
|
||||
if os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS") == "1"
|
||||
else "acceptEdits"
|
||||
),
|
||||
"--print",
|
||||
"--output-format",
|
||||
"json",
|
||||
"--no-session-persistence",
|
||||
"--append-system-prompt",
|
||||
_SUBAGENT_INSTRUCTIONS,
|
||||
f"Task: {task}",
|
||||
]
|
||||
bridged, wsl_names = _wsl_shim_env(command, local_env, _CLAUDE_ENV_UNSET)
|
||||
if wsl_names:
|
||||
from unsloth_cli.commands.start import _merge_wslenv
|
||||
|
||||
bridged = {**bridged, "PWD": os.getcwd()}
|
||||
child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_names)
|
||||
for name in _CLAUDE_ENV_UNSET:
|
||||
child_env[name] = ""
|
||||
else:
|
||||
for name in _CLAUDE_ENV_UNSET:
|
||||
child_env.pop(name, None)
|
||||
child_env.update(bridged)
|
||||
popen_kwargs: dict[str, Any] = {
|
||||
"cwd": os.environ.get("CLAUDE_PROJECT_DIR") or os.getcwd(),
|
||||
"env": child_env,
|
||||
"stdin": subprocess.DEVNULL,
|
||||
"stdout": subprocess.PIPE,
|
||||
"stderr": subprocess.PIPE,
|
||||
"text": True,
|
||||
}
|
||||
if os.name == "nt":
|
||||
popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
else:
|
||||
popen_kwargs["start_new_session"] = True
|
||||
process = subprocess.Popen(
|
||||
[executable, *command[1:]],
|
||||
**popen_kwargs,
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
stdout, stderr = process.communicate(timeout = _CANCEL_POLL_SECONDS)
|
||||
break
|
||||
except subprocess.TimeoutExpired:
|
||||
if cancel_event.is_set():
|
||||
_stop_child(process)
|
||||
raise RuntimeError("The local Claude agent was cancelled.")
|
||||
except BaseException:
|
||||
if process.poll() is None:
|
||||
_stop_child(process)
|
||||
raise
|
||||
if process.returncode != 0:
|
||||
detail = stderr.strip() or stdout.strip()
|
||||
raise RuntimeError(
|
||||
_bounded(detail) or f"Local Claude exited with code {process.returncode}."
|
||||
)
|
||||
return _result_text(stdout)
|
||||
|
||||
|
||||
def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) -> dict | None:
|
||||
request_id = request.get("id")
|
||||
method = request.get("method")
|
||||
if request_id is None:
|
||||
return None
|
||||
if method == "initialize":
|
||||
protocol = (request.get("params") or {}).get("protocolVersion") or "2025-06-18"
|
||||
result = {
|
||||
"protocolVersion": protocol,
|
||||
"capabilities": {"tools": {"listChanged": False}},
|
||||
"serverInfo": {"name": "unsloth-local-agent", "version": "1.0.0"},
|
||||
}
|
||||
elif method == "ping":
|
||||
result = {}
|
||||
elif method == "tools/list":
|
||||
result = {
|
||||
"tools": [
|
||||
{
|
||||
"name": "unsloth_agent",
|
||||
"title": "Unsloth local agent",
|
||||
"description": _SUBAGENT_DESCRIPTION,
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"task": {
|
||||
"type": "string",
|
||||
"description": "The complete task for the local Unsloth agent.",
|
||||
}
|
||||
},
|
||||
"required": ["task"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"annotations": {
|
||||
"readOnlyHint": False,
|
||||
"destructiveHint": True,
|
||||
"idempotentHint": False,
|
||||
"openWorldHint": True,
|
||||
},
|
||||
"_meta": {"anthropic/maxResultSizeChars": _MAX_RESULT_CHARACTERS},
|
||||
}
|
||||
]
|
||||
}
|
||||
elif method == "tools/call":
|
||||
params = request.get("params") or {}
|
||||
arguments = params.get("arguments") or {}
|
||||
task = arguments.get("task") if params.get("name") == "unsloth_agent" else None
|
||||
if not isinstance(task, str) or not task.strip():
|
||||
result = {
|
||||
"content": [{"type": "text", "text": "A non-empty task is required."}],
|
||||
"isError": True,
|
||||
}
|
||||
else:
|
||||
try:
|
||||
text = run_agent(task.strip())
|
||||
result = {"content": [{"type": "text", "text": text}], "isError": False}
|
||||
except Exception as exc:
|
||||
result = {
|
||||
"content": [{"type": "text", "text": str(exc)}],
|
||||
"isError": True,
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"error": {"code": -32601, "message": f"Method not found: {method}"},
|
||||
}
|
||||
return {"jsonrpc": "2.0", "id": request_id, "result": result}
|
||||
|
||||
|
||||
def serve(
|
||||
stdin: Any = sys.stdin,
|
||||
stdout: Any = sys.stdout,
|
||||
run_agent: Callable[[str, threading.Event], str] = run_local_agent,
|
||||
) -> None:
|
||||
active: dict[object, threading.Event] = {}
|
||||
workers: list[threading.Thread] = []
|
||||
state_lock = threading.RLock()
|
||||
output_lock = threading.Lock()
|
||||
shutdown_started = threading.Event()
|
||||
|
||||
def cancel_active() -> None:
|
||||
with state_lock:
|
||||
pending = list(active.values())
|
||||
for cancel_event in pending:
|
||||
cancel_event.set()
|
||||
|
||||
def handle_shutdown(_signum: int, _frame: Any) -> None:
|
||||
# Claude Code sends SIGINT (possibly repeatedly) to cancel a tool call. Only
|
||||
# the first unwinds stdin; later ones must not interrupt process-tree cleanup.
|
||||
first_signal = not shutdown_started.is_set()
|
||||
shutdown_started.set()
|
||||
cancel_active()
|
||||
if first_signal:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
previous_handlers: dict[int, Any] = {}
|
||||
if threading.current_thread() is threading.main_thread():
|
||||
for signum in (signal.SIGINT, signal.SIGTERM):
|
||||
previous_handlers[signum] = signal.signal(signum, handle_shutdown)
|
||||
|
||||
def send(response: dict | None) -> None:
|
||||
if response is None:
|
||||
return
|
||||
with output_lock:
|
||||
stdout.write(json.dumps(response, separators = (",", ":")) + "\n")
|
||||
stdout.flush()
|
||||
|
||||
def call_tool(request: dict, request_id: object, cancel_event: threading.Event) -> None:
|
||||
try:
|
||||
response = _response(
|
||||
request,
|
||||
run_agent = lambda task: run_agent(task, cancel_event),
|
||||
)
|
||||
if not cancel_event.is_set():
|
||||
send(response)
|
||||
finally:
|
||||
with state_lock:
|
||||
if active.get(request_id) is cancel_event:
|
||||
active.pop(request_id, None)
|
||||
|
||||
try:
|
||||
for line in stdin:
|
||||
try:
|
||||
request = json.loads(line)
|
||||
if not isinstance(request, dict):
|
||||
response = None
|
||||
elif request.get("method") == "notifications/cancelled":
|
||||
request_id = (request.get("params") or {}).get("requestId")
|
||||
with state_lock:
|
||||
cancel_event = active.get(request_id)
|
||||
if cancel_event is not None:
|
||||
cancel_event.set()
|
||||
response = None
|
||||
elif request.get("method") == "tools/call" and request.get("id") is not None:
|
||||
request_id = request["id"]
|
||||
cancel_event = threading.Event()
|
||||
with state_lock:
|
||||
active[request_id] = cancel_event
|
||||
worker = threading.Thread(
|
||||
target = call_tool,
|
||||
args = (request, request_id, cancel_event),
|
||||
name = f"unsloth-agent-{request_id}",
|
||||
)
|
||||
workers.append(worker)
|
||||
worker.start()
|
||||
response = None
|
||||
else:
|
||||
response = _response(request)
|
||||
except Exception as exc:
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": None,
|
||||
"error": {"code": -32603, "message": str(exc)},
|
||||
}
|
||||
send(response)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
cancel_active()
|
||||
for worker in workers:
|
||||
if worker.ident is not None:
|
||||
worker.join()
|
||||
for signum, handler in previous_handlers.items():
|
||||
signal.signal(signum, handler)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
serve()
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -106,6 +106,13 @@ API_KEY_PBKDF2_SALT_KEY = "api_key_pbkdf2_salt"
|
|||
DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"
|
||||
DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at"
|
||||
PBKDF2_ITERATIONS = 100_000
|
||||
_START_API_KEY_MARKER_ENV = "_UNSLOTH_START_API_KEY_MARKER"
|
||||
|
||||
|
||||
def _consume_start_api_key_marker_env() -> bool:
|
||||
"""Consume the one-shot readiness marker passed across a Studio re-exec."""
|
||||
return os.environ.pop(_START_API_KEY_MARKER_ENV, None) == "1"
|
||||
|
||||
|
||||
# __file__ is unsloth_cli/commands/studio.py -- two parents up is the package root
|
||||
# (either site-packages or the repo root for editable installs).
|
||||
|
|
@ -1760,6 +1767,12 @@ def run(
|
|||
"decode speed, MoE usually don't."
|
||||
),
|
||||
),
|
||||
start_api_key_marker: bool = typer.Option(
|
||||
False,
|
||||
"--start-api-key-marker",
|
||||
hidden = True,
|
||||
help = "Emit an early API key marker for the unsloth start parent process.",
|
||||
),
|
||||
password: str = typer.Option(
|
||||
"",
|
||||
"--password",
|
||||
|
|
@ -1786,6 +1799,12 @@ def run(
|
|||
unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja
|
||||
unsloth studio run --model unsloth/Qwen3-27B-GGUF --gguf-variant Q8_0 --tensor-parallel
|
||||
"""
|
||||
# A newer outer CLI can re-exec into an older Studio venv. Pass this
|
||||
# internal signal through the environment so an older child ignores it
|
||||
# instead of treating an unknown CLI option as a llama-server argument.
|
||||
inherited_start_api_key_marker = _consume_start_api_key_marker_env()
|
||||
start_api_key_marker = start_api_key_marker or inherited_start_api_key_marker
|
||||
|
||||
# Back-compat: --not-secure is a deprecated alias for --no-secure.
|
||||
secure = _resolve_secure(secure, not_secure)
|
||||
extra_llama_args: List[str] = list(ctx.args) if ctx.args else []
|
||||
|
|
@ -1991,15 +2010,22 @@ def run(
|
|||
if extra_llama_args:
|
||||
args.extend(extra_llama_args)
|
||||
|
||||
if sys.platform == "win32":
|
||||
proc = subprocess.Popen(args)
|
||||
try:
|
||||
rc = proc.wait()
|
||||
except KeyboardInterrupt:
|
||||
rc = proc.wait()
|
||||
raise typer.Exit(rc)
|
||||
else:
|
||||
os.execvp(str(studio_bin), args)
|
||||
if start_api_key_marker:
|
||||
os.environ[_START_API_KEY_MARKER_ENV] = "1"
|
||||
try:
|
||||
if sys.platform == "win32":
|
||||
proc = subprocess.Popen(args)
|
||||
try:
|
||||
rc = proc.wait()
|
||||
except KeyboardInterrupt:
|
||||
rc = proc.wait()
|
||||
raise typer.Exit(rc)
|
||||
else:
|
||||
os.execvp(str(studio_bin), args)
|
||||
finally:
|
||||
# execvp does not return on success. Restore the parent environment
|
||||
# after Windows waits for the child, or if launch fails.
|
||||
os.environ.pop(_START_API_KEY_MARKER_ENV, None)
|
||||
|
||||
# ── 2. Start server (always suppress built-in banner) ─────────────
|
||||
run_mod = _load_run_module()
|
||||
|
|
@ -2045,6 +2071,11 @@ def run(
|
|||
|
||||
# 4. Create API key in-process.
|
||||
api_key = _create_api_key_inprocess(api_key_name)
|
||||
if start_api_key_marker:
|
||||
# `unsloth start` redirects this process to a private 0600 log and
|
||||
# uses the key to authenticate download-progress polling before the
|
||||
# blocking load returns. The normal `unsloth run` output is unchanged.
|
||||
typer.echo(f"UNSLOTH_START_API_KEY: {api_key}")
|
||||
|
||||
# 5. Load model via HTTP.
|
||||
if not silent:
|
||||
|
|
|
|||
241
unsloth_cli/pi_subagent.ts
Normal file
241
unsloth_cli/pi_subagent.ts
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
import { spawn, type ChildProcess } from "node:child_process";
|
||||
import * as fs from "node:fs";
|
||||
import * as path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
||||
import { Type } from "typebox";
|
||||
|
||||
const provider = "unsloth";
|
||||
const maxResultCharacters = 100_000;
|
||||
const cancelGraceMilliseconds = 2_000;
|
||||
const configPath = process.env.UNSLOTH_PI_SUBAGENT_CONFIG || "";
|
||||
delete process.env.UNSLOTH_PI_SUBAGENT_CONFIG;
|
||||
let config: Record<string, unknown> = {};
|
||||
if (configPath) {
|
||||
try {
|
||||
const parsed = JSON.parse(fs.readFileSync(configPath, "utf8"));
|
||||
if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
|
||||
throw new Error("expected a JSON object");
|
||||
}
|
||||
config = parsed;
|
||||
} catch (error) {
|
||||
throw new Error(`Could not read Unsloth subagent configuration: ${error}`);
|
||||
}
|
||||
}
|
||||
const model = typeof config.model === "string" ? config.model : "";
|
||||
const baseUrl = typeof config.baseUrl === "string" ? config.baseUrl : "";
|
||||
const apiKey = typeof config.apiKey === "string" ? config.apiKey : "";
|
||||
const contextWindow = positiveInt(config.contextWindow, 32768);
|
||||
const maxTokens = positiveInt(config.maxTokens, Math.min(Math.floor(contextWindow / 4), 8192));
|
||||
|
||||
function positiveInt(value: unknown, fallback: number): number {
|
||||
const parsed = Number.parseInt(typeof value === "string" ? value : String(value || ""), 10);
|
||||
return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;
|
||||
}
|
||||
|
||||
function finalText(message: any): string {
|
||||
if (message?.role !== "assistant" || !Array.isArray(message.content)) return "";
|
||||
return message.content
|
||||
.filter((part: any) => part?.type === "text" && typeof part.text === "string")
|
||||
.map((part: any) => part.text)
|
||||
.join("\n")
|
||||
.trim();
|
||||
}
|
||||
|
||||
function boundedResult(text: string): string {
|
||||
if (text.length <= maxResultCharacters) return text;
|
||||
return `${text.slice(0, maxResultCharacters)}\n\n[Local agent output truncated]`;
|
||||
}
|
||||
|
||||
function piInvocation(args: string[]): { command: string; args: string[] } {
|
||||
const currentScript = process.argv[1];
|
||||
const bunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
|
||||
if (currentScript && !bunVirtualScript && fs.existsSync(currentScript)) {
|
||||
return { command: process.execPath, args: [currentScript, ...args] };
|
||||
}
|
||||
const executable = path.basename(process.execPath).toLowerCase();
|
||||
if (!/^(node|bun)(\.exe)?$/.test(executable)) return { command: process.execPath, args };
|
||||
return { command: "pi", args };
|
||||
}
|
||||
|
||||
function signalProcessGroup(child: ChildProcess, signal: NodeJS.Signals): void {
|
||||
if (!child.pid) return;
|
||||
try {
|
||||
process.kill(-child.pid, signal);
|
||||
} catch {
|
||||
try {
|
||||
child.kill(signal);
|
||||
} catch {
|
||||
// The process tree already exited.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function stopChildTree(child: ChildProcess): Promise<void> {
|
||||
if (!child.pid) return;
|
||||
if (process.platform === "win32") {
|
||||
await new Promise<void>((resolve) => {
|
||||
const killer = spawn("taskkill", ["/PID", String(child.pid), "/T", "/F"], {
|
||||
shell: false,
|
||||
stdio: "ignore",
|
||||
windowsHide: true,
|
||||
});
|
||||
killer.once("error", () => {
|
||||
try {
|
||||
child.kill("SIGKILL");
|
||||
} catch {
|
||||
// The child already exited.
|
||||
}
|
||||
resolve();
|
||||
});
|
||||
killer.once("close", (code) => {
|
||||
if (code !== 0) {
|
||||
try {
|
||||
child.kill("SIGKILL");
|
||||
} catch {
|
||||
// The child already exited.
|
||||
}
|
||||
}
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
signalProcessGroup(child, "SIGTERM");
|
||||
await new Promise((resolve) => setTimeout(resolve, cancelGraceMilliseconds));
|
||||
signalProcessGroup(child, "SIGKILL");
|
||||
}
|
||||
|
||||
export default function unslothSubagent(pi: ExtensionAPI): void {
|
||||
if (!model || !baseUrl || !apiKey || !configPath) {
|
||||
throw new Error("Unsloth subagent configuration is incomplete.");
|
||||
}
|
||||
|
||||
pi.registerProvider(provider, {
|
||||
name: "Unsloth Studio",
|
||||
baseUrl,
|
||||
apiKey,
|
||||
api: "openai-completions",
|
||||
authHeader: true,
|
||||
models: [
|
||||
{
|
||||
id: model,
|
||||
name: `${model} via Unsloth`,
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow,
|
||||
maxTokens,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
if (process.env.UNSLOTH_PI_SUBAGENT_CHILD === "1") return;
|
||||
|
||||
pi.registerTool({
|
||||
name: "unsloth_agent",
|
||||
label: "Unsloth agent",
|
||||
description:
|
||||
"Local coding subagent powered by Unsloth for debugging, implementation, and codebase research. Use when the user asks to spawn an Unsloth or local agent.",
|
||||
parameters: Type.Object({
|
||||
task: Type.String({ description: "The complete task for the local Unsloth agent." }),
|
||||
}),
|
||||
async execute(_toolCallId, params, signal, _onUpdate, ctx) {
|
||||
const extension = fileURLToPath(import.meta.url);
|
||||
const args = [
|
||||
"--mode",
|
||||
"json",
|
||||
"--print",
|
||||
"--no-session",
|
||||
"--provider",
|
||||
provider,
|
||||
"--model",
|
||||
model,
|
||||
"--no-extensions",
|
||||
"--extension",
|
||||
extension,
|
||||
`Task: ${params.task}`,
|
||||
];
|
||||
const invocation = piInvocation(args);
|
||||
let output = "";
|
||||
let stderr = "";
|
||||
let lastResponse = "";
|
||||
let childError = "";
|
||||
let aborted = false;
|
||||
const processLine = (line: string) => {
|
||||
try {
|
||||
const event = JSON.parse(line);
|
||||
if (event.type !== "message_end") return;
|
||||
const message = event.message;
|
||||
// Pi reports model/API failures as message_end events while still
|
||||
// exiting 0, so the exit status alone cannot surface them.
|
||||
if (message?.stopReason === "error" || message?.stopReason === "aborted") {
|
||||
childError =
|
||||
(typeof message.errorMessage === "string" && message.errorMessage) ||
|
||||
`The local Unsloth agent stopped: ${message.stopReason}.`;
|
||||
return;
|
||||
}
|
||||
const response = finalText(message);
|
||||
if (response) {
|
||||
lastResponse = boundedResult(response);
|
||||
childError = "";
|
||||
}
|
||||
} catch {
|
||||
// Ignore non-JSON diagnostic lines. The exit status still reports failures.
|
||||
}
|
||||
};
|
||||
|
||||
const exitCode = await new Promise<number>((resolve, reject) => {
|
||||
const child = spawn(invocation.command, invocation.args, {
|
||||
cwd: ctx.cwd,
|
||||
detached: process.platform !== "win32",
|
||||
shell: false,
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
env: {
|
||||
...process.env,
|
||||
UNSLOTH_PI_SUBAGENT_CHILD: "1",
|
||||
UNSLOTH_PI_SUBAGENT_CONFIG: configPath,
|
||||
},
|
||||
});
|
||||
let cleanup: Promise<void> | undefined;
|
||||
const cancel = () => {
|
||||
if (aborted) return;
|
||||
aborted = true;
|
||||
cleanup = stopChildTree(child);
|
||||
};
|
||||
child.on("error", (error) => {
|
||||
signal?.removeEventListener("abort", cancel);
|
||||
reject(error);
|
||||
});
|
||||
child.stdout.on("data", (chunk) => {
|
||||
output += chunk.toString();
|
||||
const lines = output.split("\n");
|
||||
output = lines.pop() || "";
|
||||
for (const line of lines) processLine(line);
|
||||
});
|
||||
child.stderr.on("data", (chunk) => {
|
||||
stderr = (stderr + chunk.toString()).slice(-100_000);
|
||||
});
|
||||
child.on("close", async (code) => {
|
||||
signal?.removeEventListener("abort", cancel);
|
||||
await cleanup;
|
||||
if (output.trim()) processLine(output);
|
||||
resolve(code ?? 1);
|
||||
});
|
||||
signal?.addEventListener("abort", cancel, { once: true });
|
||||
if (signal?.aborted) cancel();
|
||||
});
|
||||
|
||||
if (aborted) throw new Error("The local Unsloth agent was cancelled.");
|
||||
if (exitCode !== 0) {
|
||||
throw new Error(stderr.trim() || `The local Unsloth agent exited with code ${exitCode}.`);
|
||||
}
|
||||
if (childError) throw new Error(boundedResult(childError));
|
||||
return {
|
||||
content: [{ type: "text", text: lastResponse || "The local agent returned no text." }],
|
||||
details: { provider, model },
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
338
unsloth_cli/tests/test_claude_subagent_mcp.py
Normal file
338
unsloth_cli/tests/test_claude_subagent_mcp.py
Normal file
|
|
@ -0,0 +1,338 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import unsloth_cli.claude_subagent_mcp as bridge
|
||||
|
||||
|
||||
def test_protocol_lists_and_calls_local_agent():
|
||||
initialized = bridge._response(
|
||||
{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
|
||||
)
|
||||
assert initialized["result"]["serverInfo"]["name"] == "unsloth-local-agent"
|
||||
|
||||
listed = bridge._response({"jsonrpc": "2.0", "id": 2, "method": "tools/list"})
|
||||
tool = listed["result"]["tools"][0]
|
||||
assert tool["name"] == "unsloth_agent"
|
||||
assert "spawn an Unsloth or local agent" in tool["description"]
|
||||
assert tool["inputSchema"]["required"] == ["task"]
|
||||
assert tool["_meta"]["anthropic/maxResultSizeChars"] == 100_000
|
||||
|
||||
called = bridge._response(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 3,
|
||||
"method": "tools/call",
|
||||
"params": {"name": "unsloth_agent", "arguments": {"task": " inspect this "}},
|
||||
},
|
||||
run_agent = lambda task: f"completed: {task}",
|
||||
)
|
||||
assert called["result"] == {
|
||||
"content": [{"type": "text", "text": "completed: inspect this"}],
|
||||
"isError": False,
|
||||
}
|
||||
|
||||
|
||||
def test_protocol_returns_tool_errors_to_parent():
|
||||
response = bridge._response(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {"name": "unsloth_agent", "arguments": {"task": "test"}},
|
||||
},
|
||||
run_agent = lambda task: (_ for _ in ()).throw(RuntimeError("local failure")),
|
||||
)
|
||||
assert response["result"]["isError"] is True
|
||||
assert response["result"]["content"][0]["text"] == "local failure"
|
||||
|
||||
|
||||
def test_stdio_server_ignores_notifications_and_answers_requests():
|
||||
requests = "\n".join(
|
||||
[
|
||||
json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}),
|
||||
json.dumps({"jsonrpc": "2.0", "id": 4, "method": "ping"}),
|
||||
]
|
||||
)
|
||||
output = io.StringIO()
|
||||
bridge.serve(io.StringIO(requests), output)
|
||||
assert json.loads(output.getvalue()) == {"jsonrpc": "2.0", "id": 4, "result": {}}
|
||||
|
||||
|
||||
def test_stdio_cancellation_reaches_the_running_local_agent():
|
||||
requests = "\n".join(
|
||||
[
|
||||
json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "call-1",
|
||||
"method": "tools/call",
|
||||
"params": {"name": "unsloth_agent", "arguments": {"task": "wait"}},
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/cancelled",
|
||||
"params": {"requestId": "call-1", "reason": "user cancelled"},
|
||||
}
|
||||
),
|
||||
]
|
||||
)
|
||||
output = io.StringIO()
|
||||
cancelled = []
|
||||
|
||||
def run_agent(task, cancel_event):
|
||||
assert task == "wait"
|
||||
assert cancel_event.wait(timeout = 1)
|
||||
cancelled.append(task)
|
||||
raise RuntimeError("The local Claude agent was cancelled.")
|
||||
|
||||
bridge.serve(io.StringIO(requests), output, run_agent = run_agent)
|
||||
assert cancelled == ["wait"]
|
||||
assert output.getvalue() == ""
|
||||
|
||||
|
||||
def test_stdio_sigint_stops_the_running_local_agent(monkeypatch):
|
||||
request = json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "call-1",
|
||||
"method": "tools/call",
|
||||
"params": {"name": "unsloth_agent", "arguments": {"task": "wait"}},
|
||||
}
|
||||
)
|
||||
handlers = {}
|
||||
started = bridge.threading.Event()
|
||||
cancelled = []
|
||||
|
||||
def set_handler(signum, handler):
|
||||
previous = handlers.get(signum, bridge.signal.SIG_DFL)
|
||||
handlers[signum] = handler
|
||||
return previous
|
||||
|
||||
monkeypatch.setattr(bridge.signal, "signal", set_handler)
|
||||
|
||||
class InterruptingInput:
|
||||
def __init__(self):
|
||||
self.sent = False
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if not self.sent:
|
||||
self.sent = True
|
||||
return request + "\n"
|
||||
assert started.wait(timeout = 1)
|
||||
handlers[bridge.signal.SIGINT](bridge.signal.SIGINT, None)
|
||||
raise AssertionError("SIGINT handler must unwind the stdin loop")
|
||||
|
||||
def run_agent(task, cancel_event):
|
||||
assert task == "wait"
|
||||
started.set()
|
||||
assert cancel_event.wait(timeout = 1)
|
||||
# Real Claude Code sends SIGINT twice. The second one must not abort cleanup.
|
||||
handlers[bridge.signal.SIGINT](bridge.signal.SIGINT, None)
|
||||
cancelled.append(task)
|
||||
raise RuntimeError("The local Claude agent was cancelled.")
|
||||
|
||||
output = io.StringIO()
|
||||
bridge.serve(InterruptingInput(), output, run_agent = run_agent)
|
||||
assert cancelled == ["wait"]
|
||||
assert output.getvalue() == ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("bypass", "permission"),
|
||||
[("0", "acceptEdits"), ("1", "bypassPermissions")],
|
||||
)
|
||||
def test_local_child_uses_unsloth_without_overwriting_parent_auth(
|
||||
monkeypatch, tmp_path, bypass, permission
|
||||
):
|
||||
captured = {}
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888")
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test")
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M")
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW", "32768")
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS", bypass)
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "cloud-key")
|
||||
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "cloud-oauth")
|
||||
monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude")
|
||||
monkeypatch.setattr(bridge, "_claude_flags", lambda model: ["--settings", "{}"])
|
||||
|
||||
class Process:
|
||||
pid = 1234
|
||||
returncode = 0
|
||||
|
||||
def communicate(self, timeout):
|
||||
captured["timeout"] = timeout
|
||||
return json.dumps({"is_error": False, "result": "LOCAL_OK"}), ""
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def popen(command, **kwargs):
|
||||
captured["command"] = command
|
||||
captured.update(kwargs)
|
||||
return Process()
|
||||
|
||||
monkeypatch.setattr(bridge.subprocess, "Popen", popen)
|
||||
assert bridge.run_local_agent("reply exactly LOCAL_OK") == "LOCAL_OK"
|
||||
command = captured["command"]
|
||||
assert command[:3] == ["/usr/local/bin/claude", "--model", "unsloth/model-GGUF:Q4_K_M"]
|
||||
assert command[command.index("--permission-mode") + 1] == permission
|
||||
assert "--no-session-persistence" in command
|
||||
assert captured["cwd"] == str(tmp_path)
|
||||
assert captured["stdin"] is bridge.subprocess.DEVNULL
|
||||
assert captured["stdout"] is bridge.subprocess.PIPE
|
||||
assert captured["stderr"] is bridge.subprocess.PIPE
|
||||
if os.name == "nt":
|
||||
assert captured["creationflags"] == bridge.subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
else:
|
||||
assert captured["start_new_session"] is True
|
||||
child_env = captured["env"]
|
||||
assert child_env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8888"
|
||||
assert child_env["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-test"
|
||||
assert child_env["ANTHROPIC_MODEL"] == "unsloth/model-GGUF:Q4_K_M"
|
||||
assert child_env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] == "32768"
|
||||
assert child_env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] == "90"
|
||||
assert "ANTHROPIC_API_KEY" not in child_env
|
||||
assert "CLAUDE_CODE_OAUTH_TOKEN" not in child_env
|
||||
|
||||
|
||||
def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888")
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test")
|
||||
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M")
|
||||
monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude")
|
||||
monkeypatch.setattr(bridge, "_claude_flags", lambda model: [])
|
||||
cancel_event = bridge.threading.Event()
|
||||
stopped = []
|
||||
|
||||
class Process:
|
||||
pid = 1234
|
||||
returncode = None
|
||||
|
||||
def communicate(self, timeout):
|
||||
cancel_event.set()
|
||||
raise bridge.subprocess.TimeoutExpired("claude", timeout)
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
process = Process()
|
||||
monkeypatch.setattr(bridge.subprocess, "Popen", lambda *args, **kwargs: process)
|
||||
|
||||
def stop(child):
|
||||
stopped.append(child)
|
||||
child.returncode = -15
|
||||
|
||||
monkeypatch.setattr(bridge, "_stop_child", stop)
|
||||
with pytest.raises(RuntimeError, match = "cancelled"):
|
||||
bridge.run_local_agent("wait", cancel_event)
|
||||
assert stopped == [process]
|
||||
|
||||
|
||||
def test_windows_cancellation_stops_the_child_process_tree(monkeypatch):
|
||||
monkeypatch.setattr(bridge.os, "name", "nt")
|
||||
captured = {}
|
||||
|
||||
class Process:
|
||||
pid = 4321
|
||||
returncode = None
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def wait(self, timeout = None):
|
||||
captured["wait_timeout"] = timeout
|
||||
self.returncode = 1
|
||||
|
||||
def terminate(self):
|
||||
raise AssertionError("taskkill should handle the process tree")
|
||||
|
||||
def run(command, **kwargs):
|
||||
captured["command"] = command
|
||||
captured.update(kwargs)
|
||||
return bridge.subprocess.CompletedProcess(command, 0)
|
||||
|
||||
monkeypatch.setattr(bridge.subprocess, "run", run)
|
||||
bridge._stop_child(Process())
|
||||
|
||||
assert captured["command"] == ["taskkill", "/PID", "4321", "/T", "/F"]
|
||||
assert captured["capture_output"] is True
|
||||
assert captured["check"] is False
|
||||
assert captured["wait_timeout"] == bridge._CANCEL_GRACE_SECONDS
|
||||
|
||||
|
||||
def test_windows_failed_taskkill_still_terminates_the_child(monkeypatch):
|
||||
monkeypatch.setattr(bridge.os, "name", "nt")
|
||||
captured = {}
|
||||
|
||||
class Process:
|
||||
pid = 4321
|
||||
returncode = None
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def wait(self, timeout = None):
|
||||
self.returncode = 1
|
||||
|
||||
def terminate(self):
|
||||
captured["terminated"] = True
|
||||
self.returncode = 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
bridge.subprocess,
|
||||
"run",
|
||||
lambda command, **kwargs: bridge.subprocess.CompletedProcess(command, 1),
|
||||
)
|
||||
bridge._stop_child(Process())
|
||||
|
||||
assert captured.get("terminated") is True
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason = "POSIX process groups")
|
||||
def test_stop_child_kills_survivors_after_leader_exit(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(bridge, "_CANCEL_GRACE_SECONDS", 0.2)
|
||||
marker = tmp_path / "grandchild-survived"
|
||||
grandchild = (
|
||||
"import pathlib, sys, time; time.sleep(1.0); "
|
||||
"pathlib.Path(sys.argv[1]).write_text('alive')"
|
||||
)
|
||||
process = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import subprocess, sys; "
|
||||
"subprocess.Popen([sys.executable, '-c', sys.argv[1], sys.argv[2]])",
|
||||
grandchild,
|
||||
str(marker),
|
||||
],
|
||||
start_new_session = True,
|
||||
)
|
||||
process.wait()
|
||||
|
||||
bridge._stop_child(process)
|
||||
|
||||
time.sleep(1.2)
|
||||
assert not marker.exists()
|
||||
|
||||
|
||||
def test_result_parser_accepts_diagnostics_before_json():
|
||||
output = "connector warning\n" + json.dumps({"is_error": False, "result": "OK"})
|
||||
assert bridge._result_text(output) == "OK"
|
||||
191
unsloth_cli/tests/test_pi_subagent.py
Normal file
191
unsloth_cli/tests/test_pi_subagent.py
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason = "POSIX process group regression test")
|
||||
def test_pi_cancel_kills_child_process_group(tmp_path):
|
||||
bun = shutil.which("bun")
|
||||
if bun is None:
|
||||
pytest.skip("Bun is required to execute the bundled Pi extension")
|
||||
|
||||
ready = tmp_path / "grandchild-ready"
|
||||
marker = tmp_path / "grandchild-survived"
|
||||
config = tmp_path / "subagent.json"
|
||||
config.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"baseUrl": "http://127.0.0.1:8000/v1",
|
||||
"apiKey": "private-token",
|
||||
"model": "local-model",
|
||||
"contextWindow": 32768,
|
||||
"maxTokens": 8192,
|
||||
}
|
||||
),
|
||||
encoding = "utf-8",
|
||||
)
|
||||
driver = tmp_path / "pi-driver.js"
|
||||
driver.write_text(
|
||||
"""
|
||||
import { spawn } from "node:child_process";
|
||||
|
||||
spawn(
|
||||
process.execPath,
|
||||
[
|
||||
"-e",
|
||||
`
|
||||
const fs = require("node:fs");
|
||||
process.on("SIGTERM", () => {});
|
||||
fs.writeFileSync(process.env.PI_CHILD_READY, "ready");
|
||||
setTimeout(() => fs.writeFileSync(process.env.PI_CANCEL_MARKER, "alive"), 3000);
|
||||
setInterval(() => {}, 1000);
|
||||
`,
|
||||
],
|
||||
{ stdio: "inherit" },
|
||||
);
|
||||
process.on("SIGTERM", () => {});
|
||||
setInterval(() => {}, 1000);
|
||||
""",
|
||||
encoding = "utf-8",
|
||||
)
|
||||
extension = Path(__file__).parents[1] / "pi_subagent.ts"
|
||||
test_file = tmp_path / "pi-cancel.test.ts"
|
||||
test_file.write_text(
|
||||
f"""
|
||||
import {{ expect, mock, test }} from "bun:test";
|
||||
import {{ existsSync }} from "node:fs";
|
||||
import {{ pathToFileURL }} from "node:url";
|
||||
|
||||
mock.module("typebox", () => ({{
|
||||
Type: {{ Object: (value) => value, String: (value) => value }},
|
||||
}}));
|
||||
|
||||
test("cancellation stops the Pi child process group", async () => {{
|
||||
process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r};
|
||||
process.env.PI_CHILD_READY = {str(ready)!r};
|
||||
process.env.PI_CANCEL_MARKER = {str(marker)!r};
|
||||
process.argv[1] = {str(driver)!r};
|
||||
|
||||
const loaded = await import(pathToFileURL({str(extension)!r}).href);
|
||||
let tool;
|
||||
let provider;
|
||||
loaded.default({{
|
||||
registerProvider(_name, value) {{ provider = value; }},
|
||||
registerTool(value) {{ tool = value; }},
|
||||
}});
|
||||
expect(process.env.UNSLOTH_PI_SUBAGENT_CONFIG).toBeUndefined();
|
||||
expect(process.env.UNSLOTH_PI_SUBAGENT_API_KEY).toBeUndefined();
|
||||
expect(provider.apiKey).toBe("private-token");
|
||||
|
||||
const controller = new AbortController();
|
||||
const execution = tool.execute(
|
||||
"call",
|
||||
{{ task: "wait" }},
|
||||
controller.signal,
|
||||
undefined,
|
||||
{{ cwd: {str(tmp_path)!r} }},
|
||||
);
|
||||
for (let attempt = 0; attempt < 100 && !existsSync({str(ready)!r}); attempt++) {{
|
||||
await Bun.sleep(20);
|
||||
}}
|
||||
expect(existsSync({str(ready)!r})).toBe(true);
|
||||
controller.abort();
|
||||
await expect(execution).rejects.toThrow("cancelled");
|
||||
await Bun.sleep(3200);
|
||||
expect(existsSync({str(marker)!r})).toBe(false);
|
||||
}}, 10_000);
|
||||
""",
|
||||
encoding = "utf-8",
|
||||
)
|
||||
|
||||
completed = subprocess.run(
|
||||
[bun, "test", str(test_file)],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 15,
|
||||
)
|
||||
|
||||
assert completed.returncode == 0, completed.stdout + completed.stderr
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason = "POSIX driver script")
|
||||
def test_pi_child_error_events_fail_the_tool_call(tmp_path):
|
||||
bun = shutil.which("bun")
|
||||
if bun is None:
|
||||
pytest.skip("Bun is required to execute the bundled Pi extension")
|
||||
|
||||
config = tmp_path / "subagent.json"
|
||||
config.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"baseUrl": "http://127.0.0.1:8000/v1",
|
||||
"apiKey": "private-token",
|
||||
"model": "local-model",
|
||||
"contextWindow": 32768,
|
||||
"maxTokens": 8192,
|
||||
}
|
||||
),
|
||||
encoding = "utf-8",
|
||||
)
|
||||
# Pi reports model/API failures as message_end events while exiting 0.
|
||||
driver = tmp_path / "pi-driver.js"
|
||||
driver.write_text(
|
||||
"""
|
||||
const event = {
|
||||
type: "message_end",
|
||||
message: { role: "assistant", stopReason: "error", errorMessage: "backend unreachable", content: [] },
|
||||
};
|
||||
console.log(JSON.stringify(event));
|
||||
""",
|
||||
encoding = "utf-8",
|
||||
)
|
||||
extension = Path(__file__).parents[1] / "pi_subagent.ts"
|
||||
test_file = tmp_path / "pi-error.test.ts"
|
||||
test_file.write_text(
|
||||
f"""
|
||||
import {{ expect, mock, test }} from "bun:test";
|
||||
import {{ pathToFileURL }} from "node:url";
|
||||
|
||||
mock.module("typebox", () => ({{
|
||||
Type: {{ Object: (value) => value, String: (value) => value }},
|
||||
}}));
|
||||
|
||||
test("child error events fail the tool call", async () => {{
|
||||
process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r};
|
||||
process.argv[1] = {str(driver)!r};
|
||||
|
||||
const loaded = await import(pathToFileURL({str(extension)!r}).href);
|
||||
let tool;
|
||||
loaded.default({{
|
||||
registerProvider() {{}},
|
||||
registerTool(value) {{ tool = value; }},
|
||||
}});
|
||||
|
||||
const execution = tool.execute(
|
||||
"call",
|
||||
{{ task: "fail" }},
|
||||
undefined,
|
||||
undefined,
|
||||
{{ cwd: {str(tmp_path)!r} }},
|
||||
);
|
||||
await expect(execution).rejects.toThrow("backend unreachable");
|
||||
}}, 10_000);
|
||||
""",
|
||||
encoding = "utf-8",
|
||||
)
|
||||
|
||||
completed = subprocess.run(
|
||||
[bun, "test", str(test_file)],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 15,
|
||||
)
|
||||
|
||||
assert completed.returncode == 0, completed.stdout + completed.stderr
|
||||
|
|
@ -20,6 +20,7 @@ if str(_REPO_ROOT) not in sys.path:
|
|||
|
||||
|
||||
import pytest
|
||||
import typer
|
||||
from typer.testing import CliRunner
|
||||
|
||||
import unsloth_cli.commands.start as start
|
||||
|
|
@ -618,6 +619,104 @@ def test_write_codex_config_omits_catalog_for_old_codex(tmp_path, monkeypatch):
|
|||
assert not (tmp_path / "model-catalog.json").exists()
|
||||
|
||||
|
||||
def test_write_codex_subagent_config_keeps_parent_model_out(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True)
|
||||
local = {**MODEL, "id": MODEL["id"] + ":UD-Q4_K_XL"}
|
||||
path = start.write_codex_subagent_config(BASE, "private-token", local, tmp_path)
|
||||
agent = _parse_toml(path.read_text())
|
||||
assert agent["name"] == "unsloth"
|
||||
assert "local agent" in agent["description"].lower()
|
||||
assert agent["model_provider"] == start._CODEX_PROFILE
|
||||
assert agent["model"] == local["id"]
|
||||
assert agent["model_context_window"] == MODEL["context_length"]
|
||||
assert agent["model_providers"][start._CODEX_PROFILE] == {
|
||||
"name": "Unsloth Studio",
|
||||
"base_url": f"{BASE}/v1",
|
||||
"wire_api": "responses",
|
||||
"auth": {
|
||||
"command": sys.executable,
|
||||
"args": [
|
||||
"-c",
|
||||
"import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])",
|
||||
str(tmp_path / "unsloth-auth.json"),
|
||||
],
|
||||
"timeout_ms": 5000,
|
||||
},
|
||||
}
|
||||
assert json.loads((tmp_path / "unsloth-auth.json").read_text()) == {"token": "private-token"}
|
||||
catalog = json.loads((tmp_path / agent["model_catalog_json"]).read_text())
|
||||
assert catalog["models"][0]["slug"] == local["id"]
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
|
||||
def test_codex_subagent_auth_uses_wsl_for_windows_codex(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
|
||||
monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: False)
|
||||
monkeypatch.setattr(
|
||||
start.shutil,
|
||||
"which",
|
||||
lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/codex.exe",
|
||||
)
|
||||
|
||||
path = start.write_codex_subagent_config(BASE, "private-token", MODEL, tmp_path)
|
||||
auth = _parse_toml(path.read_text())["model_providers"][start._CODEX_PROFILE]["auth"]
|
||||
|
||||
assert auth["command"] == "wsl.exe"
|
||||
assert auth["args"][:5] == ["-d", "Ubuntu", "--", sys.executable, "-c"]
|
||||
assert auth["args"][-1] == str(tmp_path / "unsloth-auth.json")
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
|
||||
def test_agent_config_path_translates_for_windows_agent(monkeypatch, tmp_path):
|
||||
windows_path = r"\\wsl.localhost\Ubuntu\tmp\unsloth.toml"
|
||||
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
|
||||
monkeypatch.setattr(
|
||||
start.shutil,
|
||||
"which",
|
||||
lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/codex",
|
||||
)
|
||||
monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: windows_path)
|
||||
|
||||
assert start._agent_config_path(tmp_path / "unsloth.toml", ["codex"]) == windows_path
|
||||
|
||||
|
||||
def test_subagent_model_id_preserves_explicit_variant(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
start,
|
||||
"_http_json",
|
||||
lambda *args, **kwargs: pytest.fail("explicit variant should not need status"),
|
||||
)
|
||||
assert (
|
||||
start._subagent_model_id(BASE, "key", MODEL, MODEL["id"], "UD-Q4_K_XL")
|
||||
== MODEL["id"] + ":UD-Q4_K_XL"
|
||||
)
|
||||
|
||||
|
||||
def test_subagent_model_id_uses_loaded_variant(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
start,
|
||||
"_http_json",
|
||||
lambda *args, **kwargs: {"is_gguf": True, "gguf_variant": "Q5_K_M"},
|
||||
)
|
||||
assert start._subagent_model_id(BASE, "key", MODEL, None, None) == MODEL["id"] + ":Q5_K_M"
|
||||
|
||||
|
||||
def test_subagent_model_id_warns_when_status_unavailable(monkeypatch, capsys):
|
||||
def raise_error(*args, **kwargs):
|
||||
raise OSError("connection refused")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", raise_error)
|
||||
assert start._subagent_model_id(BASE, "key", MODEL, None, None) == MODEL["id"]
|
||||
assert "could not verify the loaded GGUF variant" in capsys.readouterr().err
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent", ["openclaw", "hermes"])
|
||||
def test_unsupported_agents_reject_as_subagent(agent):
|
||||
result = CliRunner().invoke(start.start_app, [agent, "--as-subagent"])
|
||||
assert result.exit_code == 1
|
||||
assert f"--as-subagent is not supported for {agent}." in result.output
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fake_studio(tmp_path, monkeypatch):
|
||||
calls = []
|
||||
|
|
@ -639,8 +738,13 @@ def fake_studio(tmp_path, monkeypatch):
|
|||
if url.endswith("/api/auth/api-keys"):
|
||||
return {"key": "sk-unsloth-feedfacefeedface"}
|
||||
if url.endswith("/api/inference/load"):
|
||||
already_loaded = state["models"][0]["id"] == payload["model_path"]
|
||||
state["models"] = [{"id": payload["model_path"], "context_length": 4096}]
|
||||
return {}
|
||||
return {
|
||||
"status": "already_loaded" if already_loaded else "loaded",
|
||||
"model": payload["model_path"],
|
||||
"display_name": payload["model_path"],
|
||||
}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "find_studio_server", lambda: BASE)
|
||||
|
|
@ -684,6 +788,82 @@ def test_connect_claude_no_launch(fake_studio):
|
|||
assert ".claude/settings.json" not in result.output
|
||||
|
||||
|
||||
def test_connect_claude_as_subagent_preserves_cloud_parent(fake_studio, tmp_path):
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
[
|
||||
"claude",
|
||||
"--as-subagent",
|
||||
"--no-launch",
|
||||
"--model",
|
||||
MODEL["id"] + ":UD-Q4_K_XL",
|
||||
"hello",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
command = _launch_command(result.output)
|
||||
plugin = tmp_path / "agents" / "claude-subagent" / "unsloth-local-agent"
|
||||
assert command == [
|
||||
"claude",
|
||||
"--plugin-dir",
|
||||
str(plugin),
|
||||
"--allowedTools",
|
||||
start._CLAUDE_SUBAGENT_TOOL,
|
||||
"hello",
|
||||
]
|
||||
assert "--model" not in command
|
||||
parent_base = "$env:ANTHROPIC_BASE_URL" if os.name == "nt" else "export ANTHROPIC_BASE_URL="
|
||||
parent_token = (
|
||||
"$env:ANTHROPIC_AUTH_TOKEN" if os.name == "nt" else "export ANTHROPIC_AUTH_TOKEN="
|
||||
)
|
||||
assert parent_base not in result.output
|
||||
assert parent_token not in result.output
|
||||
assert "unset ANTHROPIC_API_KEY" not in result.output
|
||||
assert "UNSLOTH_CLAUDE_SUBAGENT_API_KEY" not in result.output
|
||||
assert "sk-unsloth-feedfacefeedface" not in result.output
|
||||
assert json.loads((plugin / ".claude-plugin" / "plugin.json").read_text())["name"] == (
|
||||
"unsloth-local-agent"
|
||||
)
|
||||
mcp = json.loads((plugin / ".mcp.json").read_text())["mcpServers"]["unsloth"]
|
||||
assert mcp["command"] == sys.executable
|
||||
assert mcp["args"] == ["-m", start._CLAUDE_SUBAGENT_MCP_MODULE]
|
||||
assert mcp["env"] == {
|
||||
"UNSLOTH_CLAUDE_SUBAGENT_BASE_URL": BASE,
|
||||
"UNSLOTH_CLAUDE_SUBAGENT_API_KEY": "sk-unsloth-feedfacefeedface",
|
||||
"UNSLOTH_CLAUDE_SUBAGENT_MODEL": MODEL["id"] + ":UD-Q4_K_XL",
|
||||
"UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS": "0",
|
||||
"UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW": "4096",
|
||||
}
|
||||
skill = (plugin / "skills" / "local-agent" / "SKILL.md").read_text()
|
||||
assert "spawn an Unsloth agent or local agent" in skill
|
||||
assert "Ask Claude to spawn an Unsloth or local agent." in result.output
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
|
||||
def test_claude_subagent_plugin_uses_wsl_for_windows_claude(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
|
||||
monkeypatch.setenv("WSLENV", "EXISTING")
|
||||
monkeypatch.setattr(
|
||||
start.shutil,
|
||||
"which",
|
||||
lambda _: "/mnt/c/Users/x/AppData/Local/Programs/claude.exe",
|
||||
)
|
||||
server_env = {"UNSLOTH_CLAUDE_SUBAGENT_API_KEY": "secret"}
|
||||
plugin = start.write_claude_subagent_plugin(tmp_path, server_env)
|
||||
mcp = json.loads((plugin / ".mcp.json").read_text())["mcpServers"]["unsloth"]
|
||||
assert mcp["command"] == "wsl.exe"
|
||||
assert mcp["args"] == [
|
||||
"-d",
|
||||
"Ubuntu",
|
||||
"--",
|
||||
sys.executable,
|
||||
"-m",
|
||||
start._CLAUDE_SUBAGENT_MCP_MODULE,
|
||||
]
|
||||
assert mcp["env"]["UNSLOTH_CLAUDE_SUBAGENT_API_KEY"] == "secret"
|
||||
assert mcp["env"]["WSLENV"].split(":") == ["EXISTING", "UNSLOTH_CLAUDE_SUBAGENT_API_KEY"]
|
||||
|
||||
|
||||
def test_connect_claude_compact_window_omitted_without_context(fake_studio, monkeypatch):
|
||||
# A model that doesn't report a context length -> leave Claude's default window
|
||||
# rather than guessing one.
|
||||
|
|
@ -808,6 +988,38 @@ def test_connect_codex_no_launch(fake_studio, tmp_path):
|
|||
assert (home / "unsloth_api.config.toml").exists()
|
||||
|
||||
|
||||
def test_connect_codex_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True)
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
[
|
||||
"codex",
|
||||
"--as-subagent",
|
||||
"--no-launch",
|
||||
"--model",
|
||||
MODEL["id"] + ":UD-Q4_K_XL",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
command = _launch_command(result.output)
|
||||
assert command[0] == "codex"
|
||||
assert command[1:3] == ["--enable", "multi_agent"]
|
||||
assert "agents.max_depth=1" in command
|
||||
assert "--oss" not in command
|
||||
assert "--profile" not in command
|
||||
assert "--model" not in command
|
||||
assert "CODEX_HOME" not in result.output
|
||||
assert start._CODEX_ENV_KEY not in result.output
|
||||
assert "sk-unsloth-feedfacefeedface" not in result.output
|
||||
home = tmp_path / "agents" / "codex-subagent"
|
||||
agent_path = home / "unsloth.toml"
|
||||
agent = _parse_toml(agent_path.read_text())
|
||||
assert agent["model"] == MODEL["id"] + ":UD-Q4_K_XL"
|
||||
assert "env_key" not in agent["model_providers"][start._CODEX_PROFILE]
|
||||
assert f"agents.unsloth.config_file={json.dumps(str(agent_path))}" in command
|
||||
assert "Ask Codex to spawn an Unsloth or local agent." in result.output
|
||||
|
||||
|
||||
def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, tmp_path):
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
|
|
@ -824,7 +1036,7 @@ def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, t
|
|||
assert profile["model"] == MODEL["id"]
|
||||
|
||||
|
||||
def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch):
|
||||
def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch, capsys):
|
||||
calls = []
|
||||
state = {"loaded": False}
|
||||
|
||||
|
|
@ -862,6 +1074,8 @@ def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch):
|
|||
|
||||
assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF"
|
||||
assert any(c[1].endswith("/api/inference/load") for c in calls)
|
||||
output = capsys.readouterr().out
|
||||
assert "please wait" not in output
|
||||
|
||||
|
||||
def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
|
||||
|
|
@ -903,6 +1117,35 @@ def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
|
|||
assert any(u.endswith("/api/inference/load") for _, u in calls)
|
||||
|
||||
|
||||
def test_resolve_model_does_not_attach_if_catalog_stays_unloaded(monkeypatch):
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
if url.endswith("/v1/models"):
|
||||
return {
|
||||
"data": [
|
||||
{
|
||||
"id": "unsloth/Gemma-4-GGUF",
|
||||
"loaded": False,
|
||||
"context_length": 131072,
|
||||
}
|
||||
]
|
||||
}
|
||||
if url.endswith("/api/inference/load"):
|
||||
return {"status": "loaded", "model": "unsloth/Gemma-4-GGUF"}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
|
||||
with pytest.raises(typer.Exit):
|
||||
start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf")
|
||||
|
||||
|
||||
def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch):
|
||||
# The mirror case: a loaded entry (loaded == True) that case-matches attaches with
|
||||
# no /api/inference/load call.
|
||||
|
|
@ -931,6 +1174,25 @@ def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch
|
|||
assert not any(u.endswith("/api/inference/load") for _, u in calls)
|
||||
|
||||
|
||||
def test_resolve_model_without_request_rejects_unloaded_catalog(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
start,
|
||||
"_http_json",
|
||||
lambda *a, **k: {
|
||||
"data": [
|
||||
{
|
||||
"id": "unsloth/Gemma-4-GGUF",
|
||||
"loaded": False,
|
||||
"context_length": 131072,
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(typer.Exit):
|
||||
start._resolve_model(BASE, "sk-test", None)
|
||||
|
||||
|
||||
def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch):
|
||||
# Against a remote Unsloth the local existence probe cannot see server-side paths,
|
||||
# so a case-variant loaded id must NOT attach without a load: it could be a distinct
|
||||
|
|
@ -1213,6 +1475,9 @@ def test_connect_model_flag_loads_on_server(fake_studio):
|
|||
assert loads == [
|
||||
("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"})
|
||||
]
|
||||
assert result.output.index(
|
||||
f"Switching the Unsloth server from {MODEL['id']} to unsloth/Qwen3.5-35B-A3B.\n"
|
||||
) < result.output.index("This unloads the current model for every attached session.\n")
|
||||
_assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B")
|
||||
|
||||
|
||||
|
|
@ -1303,6 +1568,7 @@ def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio):
|
|||
assert result.exit_code == 0, result.output
|
||||
loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")]
|
||||
assert loads == []
|
||||
assert f"Reusing loaded model: {MODEL['id']}\n" in result.output
|
||||
_assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"])
|
||||
|
||||
|
||||
|
|
@ -1324,6 +1590,7 @@ def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio):
|
|||
{"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"},
|
||||
)
|
||||
]
|
||||
assert f"Reusing loaded model: {MODEL['id']}:UD-Q4_K_XL\n" in result.output
|
||||
_assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"])
|
||||
|
||||
|
||||
|
|
@ -1730,8 +1997,9 @@ def _reset_auto_served():
|
|||
start._auto_served_server = None
|
||||
|
||||
|
||||
def test_start_studio_server_builds_command_and_waits(monkeypatch):
|
||||
def test_start_studio_server_builds_command_and_waits(monkeypatch, capsys):
|
||||
captured = {}
|
||||
monkeypatch.setenv(start._START_API_KEY_MARKER_ENV, "parent")
|
||||
|
||||
class FakePopen:
|
||||
def __init__(self, command, **kwargs):
|
||||
|
|
@ -1761,13 +2029,144 @@ def test_start_studio_server_builds_command_and_waits(monkeypatch):
|
|||
assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL"
|
||||
assert cmd[cmd.index("--context-length") + 1] == "8192"
|
||||
assert "--tensor-parallel" in cmd
|
||||
assert "--start-api-key-marker" not in cmd
|
||||
assert captured["kwargs"]["env"][start._START_API_KEY_MARKER_ENV] == "1"
|
||||
assert start.os.environ[start._START_API_KEY_MARKER_ENV] == "parent"
|
||||
assert cmd[cmd.index("-p") + 1] == "8888"
|
||||
assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd
|
||||
assert captured["kwargs"].get("start_new_session") is True # own process group
|
||||
assert server.pid == 4321
|
||||
output = capsys.readouterr().out
|
||||
assert "Starting Unsloth server\n" in output
|
||||
assert "Model: unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL\n" in output
|
||||
assert "No Unsloth server at" not in output
|
||||
assert "server ready" not in output
|
||||
|
||||
|
||||
def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch):
|
||||
def test_start_studio_server_polls_progress_from_early_key(monkeypatch):
|
||||
class FakePopen:
|
||||
pid = 4321
|
||||
|
||||
def poll(self):
|
||||
return None
|
||||
|
||||
tails = iter(
|
||||
[
|
||||
"UNSLOTH_START_API_KEY: sk-unsloth-early\nLoading model...",
|
||||
"UNSLOTH_START_API_KEY: sk-unsloth-early\nModel loaded: owner/model",
|
||||
]
|
||||
)
|
||||
created = []
|
||||
|
||||
class FakeProgress:
|
||||
def __init__(self, base, key, model, variant):
|
||||
created.append((base, key, model, variant, "created"))
|
||||
|
||||
def poll(self):
|
||||
created.append("poll")
|
||||
|
||||
def close(self):
|
||||
created.append("close")
|
||||
|
||||
def complete(self):
|
||||
created.append("complete")
|
||||
|
||||
monkeypatch.setattr(start.subprocess, "Popen", lambda *a, **k: FakePopen())
|
||||
monkeypatch.setattr(start, "_studio_healthy", lambda *a, **k: True)
|
||||
monkeypatch.setattr(start, "_log_tail", lambda *a, **k: next(tails))
|
||||
monkeypatch.setattr(start, "_ModelDownloadProgress", FakeProgress)
|
||||
monkeypatch.setattr(start.time, "sleep", lambda _s: None)
|
||||
monkeypatch.setattr(
|
||||
start.typer,
|
||||
"echo",
|
||||
lambda message = "", **_kwargs: created.append(("echo", message)),
|
||||
)
|
||||
|
||||
server = start._start_studio_server(
|
||||
BASE,
|
||||
"owner/model-GGUF",
|
||||
start.LoadOptions(gguf_variant = "Q4_K_M"),
|
||||
)
|
||||
|
||||
assert server.pid == 4321
|
||||
assert (BASE, "sk-unsloth-early", "owner/model-GGUF", "Q4_K_M", "created") in created
|
||||
assert created.count("poll") == 2
|
||||
assert created[-2:] == ["complete", "close"]
|
||||
assert not any(isinstance(event, tuple) and "server ready" in event[-1] for event in created)
|
||||
|
||||
|
||||
def test_load_model_with_progress_uses_selected_gguf_size(monkeypatch, capsys):
|
||||
release = start.threading.Event()
|
||||
calls = []
|
||||
|
||||
def http_json(
|
||||
method,
|
||||
url,
|
||||
token,
|
||||
payload = None,
|
||||
timeout = 30,
|
||||
error = None,
|
||||
):
|
||||
calls.append((method, url, payload))
|
||||
if url.endswith("/api/inference/load"):
|
||||
assert release.wait(timeout = 2)
|
||||
return {"model": "owner/model-GGUF"}
|
||||
if "/api/hub/gguf-variants?" in url:
|
||||
return {
|
||||
"default_variant": "Q8_0",
|
||||
"variants": [
|
||||
{
|
||||
"quant": "UD-Q4_K_XL",
|
||||
"filename": "model-UD-Q4_K_XL.gguf",
|
||||
"size_bytes": 4 * 1024**3,
|
||||
"download_size_bytes": 4 * 1024**3,
|
||||
}
|
||||
],
|
||||
}
|
||||
if "/api/hub/gguf-download-progress?" in url:
|
||||
release.set()
|
||||
return {
|
||||
"downloaded_bytes": 2 * 1024**3,
|
||||
"expected_bytes": 4 * 1024**3,
|
||||
"progress": 0.5,
|
||||
}
|
||||
raise AssertionError(f"unexpected request: {method} {url}")
|
||||
|
||||
monkeypatch.setattr(start, "_http_json", http_json)
|
||||
monkeypatch.setattr(start, "_DOWNLOAD_POLL_INTERVAL_S", 0.001)
|
||||
result = start._load_model_with_progress(
|
||||
BASE,
|
||||
"sk-test",
|
||||
"owner/model-GGUF",
|
||||
start.LoadOptions(gguf_variant = "UD-Q4_K_XL"),
|
||||
{"model_path": "owner/model-GGUF", "gguf_variant": "UD-Q4_K_XL"},
|
||||
)
|
||||
|
||||
assert result == {"model": "owner/model-GGUF"}
|
||||
output = capsys.readouterr().out
|
||||
assert "Downloading model" in output
|
||||
assert "100%" in output
|
||||
progress_url = next(url for method, url, _ in calls if "gguf-download-progress" in url)
|
||||
assert "variant=UD-Q4_K_XL" in progress_url
|
||||
assert f"expected_bytes={4 * 1024**3}" in progress_url
|
||||
|
||||
|
||||
def test_download_progress_ignores_fully_cached_bytes(capsys):
|
||||
display = start._DownloadProgressDisplay()
|
||||
display.update(
|
||||
{
|
||||
"downloaded_bytes": 4 * 1024**3,
|
||||
"completed_bytes": 4 * 1024**3,
|
||||
"expected_bytes": 4 * 1024**3,
|
||||
"progress": 0.99,
|
||||
}
|
||||
)
|
||||
display.close()
|
||||
|
||||
assert capsys.readouterr().out == ""
|
||||
|
||||
|
||||
def test_auto_serves_when_no_server_then_keeps_server(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start, "find_studio_server", lambda: None)
|
||||
started = {}
|
||||
fake = SimpleNamespace(pid = 999, poll = lambda: None)
|
||||
|
|
@ -1793,8 +2192,83 @@ def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch):
|
|||
assert started["model"] == "unsloth/Qwen3-1.7B-GGUF"
|
||||
assert started["load"].gguf_variant == "UD-Q4_K_XL"
|
||||
assert started["base"] == BASE
|
||||
# Torn down after the agent session ended.
|
||||
assert started.get("down") is fake
|
||||
# A successful agent exit releases ownership and leaves the server available
|
||||
# for another terminal. Explicit startup failures still use the cleanup path.
|
||||
assert "down" not in started
|
||||
assert start._auto_served_server is None
|
||||
assert "is still running" in result.output
|
||||
assert "unsloth studio stop" in result.output
|
||||
|
||||
|
||||
def test_auto_served_agent_launch_failure_stops_server(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start, "find_studio_server", lambda: None)
|
||||
stopped = []
|
||||
fake = SimpleNamespace(pid = 999, poll = lambda: None)
|
||||
|
||||
def fake_start(*_args):
|
||||
start._auto_served_server = fake
|
||||
return fake
|
||||
|
||||
monkeypatch.setattr(start, "_start_studio_server", fake_start)
|
||||
monkeypatch.setattr(start, "_shutdown_server", stopped.append)
|
||||
monkeypatch.setattr(
|
||||
start,
|
||||
"_launch",
|
||||
lambda *a, **k: (_ for _ in ()).throw(RuntimeError("agent launch failed")),
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert stopped == [fake]
|
||||
assert "is still running" not in result.output
|
||||
|
||||
|
||||
def test_auto_served_server_exit_is_not_reported_as_running(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start, "find_studio_server", lambda: None)
|
||||
fake = SimpleNamespace(pid = 999, poll = lambda: 1)
|
||||
|
||||
def fake_start(*_args):
|
||||
start._auto_served_server = fake
|
||||
return fake
|
||||
|
||||
monkeypatch.setattr(start, "_start_studio_server", fake_start)
|
||||
monkeypatch.setattr(start, "_launch", lambda *a, **k: 0)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "stopped during the session" in result.output
|
||||
assert "is still running" not in result.output
|
||||
|
||||
|
||||
def test_attached_server_prints_stop_hint_after_agent_exits(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
|
||||
monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
|
||||
monkeypatch.setattr(
|
||||
start.subprocess,
|
||||
"run",
|
||||
lambda command, env: SimpleNamespace(returncode = 0),
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(start.start_app, ["claude"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert f"Unsloth ready at {BASE} · model {MODEL['id']}\n" in result.output
|
||||
assert f"Unsloth Studio is still running at {BASE}." in result.output
|
||||
assert "Stop it with: unsloth studio stop\n" in result.output
|
||||
|
||||
|
||||
def test_no_launch_recipe_does_not_print_stop_hint(fake_studio):
|
||||
result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "is still running" not in result.output
|
||||
|
||||
|
||||
def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch):
|
||||
|
|
@ -2092,8 +2566,7 @@ def test_write_opencode_config_fresh(tmp_path):
|
|||
MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}}
|
||||
}
|
||||
assert config["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"
|
||||
# The overlay never writes disabled_providers; the dedicated provider id is one a
|
||||
# user's disable list would not target, so nothing needs re-enabling.
|
||||
# Provider filters belong to the launch-time inline overlay, not this config writer.
|
||||
assert "disabled_providers" not in config
|
||||
# Compaction buffer scaled to ~10% of the window (compact near 90%).
|
||||
assert config["compaction"] == {"auto": True, "reserved": 131072 // 10}
|
||||
|
|
@ -2134,6 +2607,109 @@ def test_write_opencode_config_keeps_foreign_disabled_providers(tmp_path):
|
|||
assert config["disabled_providers"] == ["openai", "gemini"]
|
||||
|
||||
|
||||
def test_write_opencode_config_as_subagent_preserves_parent_model(tmp_path):
|
||||
path = tmp_path / "opencode.json"
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"small_model": "anthropic/claude-haiku-4-5",
|
||||
"compaction": {"auto": False},
|
||||
}
|
||||
)
|
||||
)
|
||||
local = {**MODEL, "id": MODEL["id"] + ":UD-Q4_K_XL"}
|
||||
start.write_opencode_config(
|
||||
BASE,
|
||||
"sk-unsloth-abc",
|
||||
local,
|
||||
path,
|
||||
as_subagent = True,
|
||||
)
|
||||
config = json.loads(path.read_text())
|
||||
assert config["model"] == "anthropic/claude-sonnet-4-5"
|
||||
assert config["small_model"] == "anthropic/claude-haiku-4-5"
|
||||
assert config["compaction"] == {"auto": False}
|
||||
agent = config["agent"]["unsloth"]
|
||||
assert agent["mode"] == "subagent"
|
||||
assert agent["model"] == f"{start._OPENCODE_PROVIDER}/{local['id']}"
|
||||
assert "local agent" in agent["description"].lower()
|
||||
assert local["id"] in config["provider"][start._OPENCODE_PROVIDER]["models"]
|
||||
|
||||
|
||||
def test_opencode_subagent_inline_keeps_parent_provider_filters(monkeypatch, tmp_path):
|
||||
config_path = tmp_path / "opencode.json"
|
||||
inherited = {"theme": "tokyonight"}
|
||||
monkeypatch.setenv("OPENCODE_CONFIG_CONTENT", json.dumps(inherited))
|
||||
monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: "/usr/bin/opencode")
|
||||
captured = {}
|
||||
|
||||
def run(command, **kwargs):
|
||||
captured["command"] = command
|
||||
captured.update(kwargs)
|
||||
return SimpleNamespace(
|
||||
returncode = 0,
|
||||
stdout = json.dumps(
|
||||
{
|
||||
"enabled_providers": ["opencode-go"],
|
||||
"disabled_providers": ["ollama", start._OPENCODE_PROVIDER],
|
||||
"subagent_depth": 0,
|
||||
}
|
||||
),
|
||||
stderr = "",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(start.subprocess, "run", run)
|
||||
permission = {"edit": "allow"}
|
||||
inline = start._opencode_subagent_inline_config(config_path, permission)
|
||||
|
||||
assert captured["command"] == ["/usr/bin/opencode", "debug", "config"]
|
||||
assert captured["env"]["OPENCODE_CONFIG"] == str(config_path)
|
||||
assert inline == {
|
||||
"theme": "tokyonight",
|
||||
"enabled_providers": ["opencode-go", start._OPENCODE_PROVIDER],
|
||||
"disabled_providers": ["ollama"],
|
||||
"subagent_depth": 1,
|
||||
"permission": permission,
|
||||
}
|
||||
|
||||
|
||||
def test_opencode_subagent_inline_preserves_positive_depth(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: "/usr/bin/opencode")
|
||||
monkeypatch.setattr(
|
||||
start.subprocess,
|
||||
"run",
|
||||
lambda *args, **kwargs: SimpleNamespace(
|
||||
returncode = 0,
|
||||
stdout = json.dumps({"subagent_depth": 3}),
|
||||
stderr = "",
|
||||
),
|
||||
)
|
||||
|
||||
inline = start._opencode_subagent_inline_config(tmp_path / "opencode.json", {})
|
||||
|
||||
assert inline["subagent_depth"] == 3
|
||||
|
||||
|
||||
def test_opencode_subagent_inline_merges_inherited_filters_without_binary(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv(
|
||||
"OPENCODE_CONFIG_CONTENT",
|
||||
json.dumps(
|
||||
{
|
||||
"enabled_providers": ["opencode-go"],
|
||||
"disabled_providers": ["ollama", start._OPENCODE_PROVIDER],
|
||||
}
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: None)
|
||||
|
||||
inline = start._opencode_subagent_inline_config(tmp_path / "opencode.json", {})
|
||||
|
||||
assert inline["enabled_providers"] == ["opencode-go", start._OPENCODE_PROVIDER]
|
||||
assert inline["disabled_providers"] == ["ollama"]
|
||||
assert inline["subagent_depth"] == 1
|
||||
|
||||
|
||||
def _opencode_inline_config(output: str) -> dict:
|
||||
# --no-launch prints OPENCODE_CONFIG_CONTENT as a POSIX `export NAME=<shell-quoted>`
|
||||
# line on Unix/WSL and a PowerShell `$env:NAME = "<escaped>"` line on native Windows;
|
||||
|
|
@ -2222,6 +2798,130 @@ def test_connect_opencode_no_launch(fake_studio, tmp_path):
|
|||
assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
|
||||
|
||||
|
||||
def test_connect_opencode_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(start, "_opencode_subagent_inline_config", lambda path, permission: {})
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
[
|
||||
"opencode",
|
||||
"--as-subagent",
|
||||
"--no-launch",
|
||||
"--model",
|
||||
MODEL["id"] + ":UD-Q4_K_XL",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert _launch_command(result.output) == ["opencode"]
|
||||
expected_model = f"{start._OPENCODE_PROVIDER}/{MODEL['id']}:UD-Q4_K_XL"
|
||||
# The agent rides in the inline overlay; nothing else comes from the empty base.
|
||||
assert _opencode_inline_config(result.output) == {
|
||||
"agent": {
|
||||
"unsloth": {
|
||||
"description": start._SUBAGENT_DESCRIPTION,
|
||||
"mode": "subagent",
|
||||
"model": expected_model,
|
||||
"prompt": start._SUBAGENT_INSTRUCTIONS,
|
||||
}
|
||||
}
|
||||
}
|
||||
path = tmp_path / "agents" / "opencode-subagent" / "opencode.json"
|
||||
config = json.loads(path.read_text())
|
||||
assert "model" not in config
|
||||
assert "small_model" not in config
|
||||
assert "compaction" not in config
|
||||
agent = config["agent"]["unsloth"]
|
||||
assert agent["model"] == expected_model
|
||||
assert "Unsloth is available as @unsloth and in /models." in result.output
|
||||
|
||||
|
||||
def test_claude_subagent_allowed_tools_precede_forwarded_delimiter(fake_studio):
|
||||
# A forwarded `--` makes everything after it positional; the tool pre-approval
|
||||
# must be parsed as an option, so it rides before ctx.args.
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
["claude", "--as-subagent", "--no-launch", "--", "--resume", "abc123"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
command = _launch_command(result.output)
|
||||
assert command.index("--allowedTools") < command.index("--resume")
|
||||
|
||||
|
||||
def test_opencode_subagent_installs_binary_before_filter_inspection(fake_studio, monkeypatch):
|
||||
# The effective-config inspection needs the opencode binary; a first launch must
|
||||
# offer the install before building the overlay, or a global allowlist read only
|
||||
# after _launch installs OpenCode would filter out the new provider.
|
||||
installed = {}
|
||||
monkeypatch.setattr(
|
||||
start,
|
||||
"_which_with_install_dirs",
|
||||
lambda name: "/usr/local/bin/opencode" if installed.get("done") else None,
|
||||
)
|
||||
|
||||
def install(name, hint):
|
||||
installed["done"] = True
|
||||
installed["name"] = name
|
||||
return "/usr/local/bin/opencode"
|
||||
|
||||
monkeypatch.setattr(start, "_install_agent", install)
|
||||
inspected = {}
|
||||
|
||||
def inline(path, permission):
|
||||
inspected["binary"] = start._which_with_install_dirs("opencode")
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr(start, "_opencode_subagent_inline_config", inline)
|
||||
monkeypatch.setattr(start, "_run", lambda *a, **k: None)
|
||||
|
||||
result = CliRunner().invoke(start.start_app, ["opencode", "--as-subagent"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert installed["name"] == "opencode"
|
||||
assert inspected["binary"] == "/usr/local/bin/opencode"
|
||||
|
||||
|
||||
def test_opencode_subagent_pins_agent_in_inline_overlay(fake_studio, monkeypatch):
|
||||
# A project opencode.json outranks the session file, so the agent must ride in
|
||||
# OPENCODE_CONFIG_CONTENT where a repo's own agent.unsloth cannot field-merge over it.
|
||||
monkeypatch.setattr(start, "_opencode_subagent_inline_config", lambda path, permission: {})
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
["opencode", "--as-subagent", "--no-launch", "--model", MODEL["id"] + ":UD-Q4_K_XL"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
agent = _opencode_inline_config(result.output)["agent"]["unsloth"]
|
||||
assert agent["mode"] == "subagent"
|
||||
assert agent["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}:UD-Q4_K_XL"
|
||||
assert agent["prompt"] == start._SUBAGENT_INSTRUCTIONS
|
||||
assert agent["description"] == start._SUBAGENT_DESCRIPTION
|
||||
|
||||
|
||||
def test_connect_opencode_subagent_yolo_no_launch_stays_append_safe(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start, "_opencode_supports_native_auto", lambda: True)
|
||||
captured = {}
|
||||
|
||||
def inline(path, permission):
|
||||
captured["permission"] = permission
|
||||
return {"permission": permission}
|
||||
|
||||
monkeypatch.setattr(start, "_opencode_subagent_inline_config", inline)
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
["opencode", "--as-subagent", "--no-launch", "--yolo"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert _launch_command(result.output) == ["opencode"]
|
||||
assert "--auto" not in result.output
|
||||
assert captured["permission"] == {
|
||||
"edit": "allow",
|
||||
"bash": "allow",
|
||||
"webfetch": "allow",
|
||||
"task": "allow",
|
||||
"external_directory": {"*": "allow"},
|
||||
}
|
||||
assert _opencode_inline_config(result.output)["permission"] == captured["permission"]
|
||||
|
||||
|
||||
# ── Hermes (OpenAI /v1/chat/completions, key via env) ────────────────
|
||||
|
||||
|
||||
|
|
@ -2364,6 +3064,39 @@ def test_connect_pi_no_launch(fake_studio, tmp_path):
|
|||
assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
|
||||
|
||||
|
||||
def test_connect_pi_as_subagent_preserves_cloud_parent(fake_studio, tmp_path):
|
||||
result = CliRunner().invoke(
|
||||
start.start_app,
|
||||
[
|
||||
"pi",
|
||||
"--as-subagent",
|
||||
"--no-launch",
|
||||
"--model",
|
||||
MODEL["id"] + ":UD-Q4_K_XL",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
command = _launch_command(result.output)
|
||||
assert command[:2] == ["pi", "--extension"]
|
||||
assert command[2].endswith("unsloth_cli/pi_subagent.ts")
|
||||
assert "--provider" not in command
|
||||
assert "--model" not in command
|
||||
assert "PI_CODING_AGENT_DIR" not in result.output
|
||||
assert "export HOME=" not in result.output
|
||||
assert "UNSLOTH_PI_SUBAGENT_API_KEY" not in result.output
|
||||
assert "sk-unsloth-feedfacefeedface" not in result.output
|
||||
config_path = tmp_path / "agents" / "pi-subagent" / "subagent.json"
|
||||
_assert_env_set(result.output, "UNSLOTH_PI_SUBAGENT_CONFIG", str(config_path))
|
||||
assert json.loads(config_path.read_text()) == {
|
||||
"baseUrl": f"{BASE}/v1",
|
||||
"apiKey": "sk-unsloth-feedfacefeedface",
|
||||
"model": MODEL["id"] + ":UD-Q4_K_XL",
|
||||
"contextWindow": 4096,
|
||||
"maxTokens": 1024,
|
||||
}
|
||||
assert "Ask Pi to spawn an Unsloth or local agent." in result.output
|
||||
|
||||
|
||||
def test_connect_pi_no_launch_windows_relocates_userprofile(fake_studio, tmp_path, monkeypatch):
|
||||
# On native Windows Node resolves ~/.pi via USERPROFILE, not HOME, so the session
|
||||
# must point USERPROFILE at the relocated home or Pi reads the user's real ~/.pi.
|
||||
|
|
@ -2907,6 +3640,27 @@ def test_opencode_non_yolo_flips_only_explicit_allow(tmp_path):
|
|||
assert session == {} # a non-yolo session carries no permission inline
|
||||
|
||||
|
||||
def test_opencode_subagent_non_yolo_clears_yolo_task_permission(tmp_path):
|
||||
path = tmp_path / "opencode.json"
|
||||
start.write_opencode_config(
|
||||
BASE,
|
||||
"sk-unsloth-abc",
|
||||
MODEL,
|
||||
path,
|
||||
yolo = True,
|
||||
as_subagent = True,
|
||||
)
|
||||
start.write_opencode_config(
|
||||
BASE,
|
||||
"sk-unsloth-abc",
|
||||
MODEL,
|
||||
path,
|
||||
as_subagent = True,
|
||||
)
|
||||
|
||||
assert json.loads(path.read_text())["permission"]["task"] == "ask"
|
||||
|
||||
|
||||
def test_opencode_non_yolo_leaves_string_permission(tmp_path):
|
||||
# A global string rule ("deny") is a user-managed catch-all; leave it untouched and
|
||||
# carry no inline override.
|
||||
|
|
|
|||
|
|
@ -170,13 +170,24 @@ def _install_reexec_capture(monkeypatch, *, platform):
|
|||
|
||||
monkeypatch.setattr(sys, "platform", platform)
|
||||
|
||||
def capture(kind, argv):
|
||||
captured.append(
|
||||
{
|
||||
"kind": kind,
|
||||
"argv": list(argv),
|
||||
"start_api_key_marker": studio_mod.os.environ.get(
|
||||
studio_mod._START_API_KEY_MARKER_ENV
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
def fake_execvp(file, argv):
|
||||
captured.append({"kind": "execvp", "argv": list(argv)})
|
||||
capture("execvp", argv)
|
||||
raise _ExecCaptured(argv)
|
||||
|
||||
class _FakePopen:
|
||||
def __init__(self, argv, *a, **kw):
|
||||
captured.append({"kind": "popen", "argv": list(argv)})
|
||||
capture("popen", argv)
|
||||
self._argv = argv
|
||||
|
||||
def wait(self):
|
||||
|
|
@ -235,6 +246,30 @@ def test_reexec_forwards_parallel_all_aliases(monkeypatch, flag, value):
|
|||
), f"{flag} {value} was dropped on re-exec; argv = {argv}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("platform", ["linux", "darwin", "win32"])
|
||||
def test_reexec_hands_off_start_api_key_marker_out_of_band(monkeypatch, platform):
|
||||
"""A new child receives the marker while an old child sees no unknown flag."""
|
||||
result, captured = _invoke_run(
|
||||
monkeypatch,
|
||||
_BASE + ["--start-api-key-marker"],
|
||||
platform = platform,
|
||||
)
|
||||
assert len(captured) == 1, result.output
|
||||
assert "--start-api-key-marker" not in captured[0]["argv"]
|
||||
assert captured[0]["start_api_key_marker"] == "1"
|
||||
|
||||
|
||||
def test_reexeced_child_consumes_start_api_key_marker_env(monkeypatch):
|
||||
"""A supported child consumes the handoff before starting descendants."""
|
||||
studio_mod = _load_run_command()
|
||||
monkeypatch.setenv(studio_mod._START_API_KEY_MARKER_ENV, "1")
|
||||
|
||||
inherited = studio_mod._consume_start_api_key_marker_env()
|
||||
|
||||
assert inherited is True
|
||||
assert studio_mod._START_API_KEY_MARKER_ENV not in studio_mod.os.environ
|
||||
|
||||
|
||||
@pytest.mark.parametrize("platform", ["linux", "darwin", "win32"])
|
||||
def test_reexec_argv_is_consistent_across_platforms(monkeypatch, platform):
|
||||
"""Linux/Darwin (execvp) and Windows (Popen) must build the same argv."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue