Bundle Needle 2 as a built-in, default AI assistant

Ships the official Needle 2 CLI binary (Apache-2.0) baked into the ferrum
binary itself for windows/amd64, linux/amd64, linux/arm64, and
darwin/arm64 via go:embed behind per-platform build tags. No download,
no FERRUM_NEEDLE_BIN, no manual provider setup on those platforms.

- internal/needle: resolveBinPath prefers an explicit FERRUM_NEEDLE_BIN,
  otherwise extracts the embedded binary to a cache file on first use.
- Fixed a port collision: Needle's --serve defaulted to :8080, the same
  default as ferrum's own server; it now runs on a dedicated port.
- Fixed the real 'assistant times out' bug: Needle's tool-retrieval does
  a one-time embedding pass on its first request once the tool catalog
  exceeds 5 tools (ferrum declares 10), which routinely took longer than
  the old readiness check's 500ms-per-attempt retry loop allowed. Each
  cancelled attempt kept occupying Needle's single-threaded request
  loop, so retries piled up and never let a real response through.
  Replaced with two phases: cheap/retryable raw TCP dials until the
  socket accepts a connection, then exactly one real request allowed to
  run for the full startup budget.
  Subprocess stdout/stderr are now captured so a future failure surfaces
  a real reason instead of a bare timeout.
- seedBuiltinNeedleProvider now runs on every startup (not just fresh
  installs), idempotently upserting the built-in provider and always
  reasserting its model as the global default AI assistant.
- .gitignore: carved out an exception for the bundled Windows binary,
  which the blanket *.exe rule would otherwise have silently excluded.
This commit is contained in:
Anand
2026-09-06 20:26:00 +05:30
parent 0bcd15fe6a
commit 846cb8ce63
17 changed files with 605 additions and 42 deletions
+1
View File
@@ -9,6 +9,7 @@ config.yaml
# Go build artifacts
/ferrum
*.exe
!internal/needle/bundled/**/*.exe
*.exe~
*.dll
*.so
+6 -5
View File
@@ -155,13 +155,14 @@ Everything else — notifications, SSO details, security policy, system settings
The AI Assistant and MCP tool-calling loop can use any OpenAI-chat-completions-compatible provider (OpenAI, Ollama, LM Studio, LocalAI, OpenRouter, ...) configured under **Settings > AI Providers**. There's also an optional zero-config, no-API-key, fully local option backed by [Needle 2](https://huggingface.co/Cactus-Compute/needle2) — a small (45M-parameter) tool-calling model that runs as a self-contained CLI binary with no GPU and no network access required at inference time.
Ferrum does **not** download or bundle this binary itself — it's a third-party artifact only distributed from Hugging Face, and Ferrum never fetches executable content from the network on its own. To enable it:
Needle 2 is Apache-2.0 licensed, so on **Windows, Linux, and macOS (amd64 or arm64)** Ferrum ships its official CLI binary baked into the `ferrum` binary itself (`internal/needle/bundled_*.go`, one per platform via `go:embed`) — nothing to download, nothing to configure. On a fresh install (no AI provider configured yet), Ferrum extracts it to a cache file and registers it automatically as the default assistant the first time it starts — no manual "Add provider" step needed. If you've already configured a provider, or want to add/re-add it yourself, use **Settings > AI Providers** > "Add provider" > the **Needle 2 (built-in, local)** preset.
1. Download the `needle` CLI binary for your platform from the [Needle 2 files](https://huggingface.co/Cactus-Compute/needle2/tree/main) (the `linux/`, `macos/`, or `windows/` directory).
2. Point Ferrum at it: set `FERRUM_NEEDLE_BIN=/path/to/needle` (or `needleBinPath` in `config.yaml`) before starting Ferrum.
3. In **Settings > AI Providers**, click "Add provider" and choose the **Needle 2 (built-in, local)** preset, then save.
On any other platform (32-bit, RISC-V, Windows/ARM64, ...) there's no bundled binary — Ferrum still never fetches executable content from the network on its own. To enable it there:
Ferrum starts the binary itself (as a local subprocess, `127.0.0.1`-only) the first time it's used, and stops it on shutdown. If `FERRUM_NEEDLE_BIN` isn't set, or the file doesn't exist, this provider simply isn't usable — every other provider is unaffected.
1. Download the `needle` CLI binary for your platform from the [Needle 2 files](https://huggingface.co/Cactus-Compute/needle2/tree/main).
2. Point Ferrum at it: set `FERRUM_NEEDLE_BIN=/path/to/needle` (or `needleBinPath` in `config.yaml`) before starting Ferrum. This also overrides the bundled binary on a supported platform, if you'd rather run a different build.
Ferrum starts the binary itself (as a local subprocess, `127.0.0.1`-only) the first time it's used, and stops it on shutdown. If no binary is bundled for the platform and `FERRUM_NEEDLE_BIN` isn't set (or doesn't exist), this provider simply isn't usable — every other provider is unaffected.
## API access, MCP, and audit logging
+69
View File
@@ -7,6 +7,7 @@ import (
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"sort"
"strings"
@@ -18,6 +19,74 @@ import (
"ferrum/internal/needle"
)
// seedBuiltinNeedleProvider makes the built-in Needle 2 provider Ferrum's
// default AI assistant, unconditionally, on every startup where a binary is
// available for the running platform (bundled — see internal/needle's
// go:embed — or FERRUM_NEEDLE_BIN) — not just on a fresh install. It is
// idempotent (matches the existing row by base_url = needle.BaseURL rather
// than inserting a duplicate on every restart) and always (re)asserts its
// model as THE global default, demoting whatever else held that spot: the
// point of a *built-in* assistant is that it's always there and always
// selected, with zero setup, on any install that has a binary for its
// platform — that's a deliberate, standing product decision, not a
// one-time fallback for an empty database.
func (s *Server) seedBuiltinNeedleProvider(ctx context.Context) {
if !s.needle.Available() {
return
}
now := time.Now().UTC().Format(time.RFC3339)
var providerID string
err := s.db.QueryRowContext(ctx, `SELECT id FROM ai_providers WHERE base_url = ?`, needle.BaseURL).Scan(&providerID)
switch {
case err == sql.ErrNoRows:
providerID = uuid.NewString()
if _, err := s.db.ExecContext(ctx, `
INSERT INTO ai_providers (id, name, base_url, api_key_enc, model, is_enabled, is_default, created_at, updated_at)
VALUES (?, 'Needle 2 (built-in, local)', ?, NULL, '', 1, 0, ?, ?)`,
providerID, needle.BaseURL, now, now,
); err != nil {
slog.Error("seeding built-in needle provider", "error", err)
return
}
case err != nil:
slog.Error("looking up built-in needle provider", "error", err)
return
default:
// Already registered from a previous run — make sure it's still
// enabled even if an operator had switched it off.
if _, err := s.db.ExecContext(ctx, `UPDATE ai_providers SET is_enabled = 1, updated_at = ? WHERE id = ?`, now, providerID); err != nil {
slog.Error("re-enabling built-in needle provider", "error", err)
}
}
var modelID string
err = s.db.QueryRowContext(ctx, `SELECT id FROM ai_provider_models WHERE provider_id = ?`, providerID).Scan(&modelID)
switch {
case err == sql.ErrNoRows:
modelID = uuid.NewString()
if _, err := s.db.ExecContext(ctx, `
INSERT INTO ai_provider_models (id, provider_id, label, model_id, is_default, created_at)
VALUES (?, ?, 'needle2', 'needle2', 1, ?)`,
modelID, providerID, now,
); err != nil {
slog.Error("seeding built-in needle model", "error", err)
return
}
case err != nil:
slog.Error("looking up built-in needle model", "error", err)
return
}
if err := s.clearOtherDefaultModels(ctx, modelID); err != nil {
slog.Error("clearing other default AI models", "error", err)
return
}
if _, err := s.db.ExecContext(ctx, `UPDATE ai_provider_models SET is_default = 1 WHERE id = ?`, modelID); err != nil {
slog.Error("setting built-in needle model as default", "error", err)
}
}
// aiModelDTO is one selectable model under a provider. label is the
// human-facing name shown in pickers; modelID is the exact identifier sent
// to the provider's API — kept distinct since they're frequently different
+70
View File
@@ -0,0 +1,70 @@
package api
import (
"context"
"testing"
"ferrum/internal/needle"
)
// TestSeedBuiltinNeedleProviderIsDefaultAndSticky exercises the "built-in
// assistant is always there and always the default" contract: New() already
// seeds it once; this checks that a competing provider set as default gets
// demoted back on the next seed pass (e.g. server restart), and that
// re-seeding never creates a second built-in provider row.
func TestSeedBuiltinNeedleProviderIsDefaultAndSticky(t *testing.T) {
env := newTestEnv(t)
if !env.server.needle.Available() {
t.Skip("no needle binary bundled/configured for this platform")
}
ctx := context.Background()
var providerCount int
if err := env.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM ai_providers WHERE base_url = ?`, needle.BaseURL).Scan(&providerCount); err != nil {
t.Fatalf("counting needle providers: %v", err)
}
if providerCount != 1 {
t.Fatalf("expected exactly one built-in provider row after New(), got %d", providerCount)
}
// Simulate an operator adding another provider and making it the
// default — this must not survive the next seed pass.
if _, err := env.db.ExecContext(ctx, `
INSERT INTO ai_providers (id, name, base_url, model, is_enabled, is_default, created_at, updated_at)
VALUES ('other-provider', 'Other', 'https://example.com/v1', '', 1, 0, 'now', 'now')`); err != nil {
t.Fatalf("inserting competing provider: %v", err)
}
if _, err := env.db.ExecContext(ctx, `
INSERT INTO ai_provider_models (id, provider_id, label, model_id, is_default, created_at)
VALUES ('other-model', 'other-provider', 'Other model', 'other', 1, 'now')`); err != nil {
t.Fatalf("inserting competing model: %v", err)
}
env.server.seedBuiltinNeedleProvider(ctx)
var defaultCount int
if err := env.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM ai_provider_models WHERE is_default = 1`).Scan(&defaultCount); err != nil {
t.Fatalf("counting default models: %v", err)
}
if defaultCount != 1 {
t.Fatalf("expected exactly one default model, got %d", defaultCount)
}
var defaultProviderBaseURL string
if err := env.db.QueryRowContext(ctx, `
SELECT p.base_url FROM ai_provider_models m
JOIN ai_providers p ON p.id = m.provider_id
WHERE m.is_default = 1`).Scan(&defaultProviderBaseURL); err != nil {
t.Fatalf("finding default provider: %v", err)
}
if defaultProviderBaseURL != needle.BaseURL {
t.Fatalf("expected the built-in Needle provider to be the default, got %q", defaultProviderBaseURL)
}
if err := env.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM ai_providers WHERE base_url = ?`, needle.BaseURL).Scan(&providerCount); err != nil {
t.Fatalf("re-counting needle providers: %v", err)
}
if providerCount != 1 {
t.Fatalf("re-seeding must not duplicate the built-in provider row, got %d", providerCount)
}
}
+1
View File
@@ -162,6 +162,7 @@ func New(db *store.DB, authSvc *auth.Service, secretBox *secrets.Box, opts Serve
}
return out
})
s.seedBuiltinNeedleProvider(context.Background())
return s
}
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+10
View File
@@ -0,0 +1,10 @@
//go:build darwin && arm64
package needle
import _ "embed"
//go:embed bundled/darwin-arm64/needle
var bundledBinary []byte
const bundledName = "needle"
+10
View File
@@ -0,0 +1,10 @@
//go:build linux && amd64
package needle
import _ "embed"
//go:embed bundled/linux-amd64/needle
var bundledBinary []byte
const bundledName = "needle"
+10
View File
@@ -0,0 +1,10 @@
//go:build linux && arm64
package needle
import _ "embed"
//go:embed bundled/linux-arm64/needle
var bundledBinary []byte
const bundledName = "needle"
+12
View File
@@ -0,0 +1,12 @@
//go:build !((windows && amd64) || (linux && amd64) || (linux && arm64) || (darwin && arm64))
// This file backs every platform Ferrum doesn't ship a bundled Needle 2
// binary for (32-bit, RISC-V, Windows/ARM64, ...) — see the README "Built-in
// LLM (Needle 2)" section for the full platform list Needle itself supports.
// An operator on one of those can still point FERRUM_NEEDLE_BIN at a binary
// they downloaded themselves; only the auto-bundled fallback is unavailable.
package needle
var bundledBinary []byte
const bundledName = ""
+10
View File
@@ -0,0 +1,10 @@
//go:build windows && amd64
package needle
import _ "embed"
//go:embed bundled/windows-amd64/needle.exe
var bundledBinary []byte
const bundledName = "needle.exe"
+170 -35
View File
@@ -3,15 +3,17 @@
// self-contained CLI binary — as a zero-config, no-API-key "built-in" AI
// provider for Ferrum's AI Assistant and MCP tool-calling loop.
//
// Ferrum does not download, bundle, or execute this binary automatically:
// it is a proprietary third-party artifact (CLI binary or WASM component)
// distributed only from Hugging Face, and Ferrum never fetches executable
// content from the network on its own. An operator who wants this provider
// must download the CLI binary for their platform themselves and point
// FERRUM_NEEDLE_BIN (or the default ./data/needle/needle[.exe]) at it — see
// README "Built-in LLM (Needle 2)" for the exact steps. Available() reports
// false, and every call fails with a clear "not installed" error, until
// that file exists.
// Needle 2 is Apache-2.0, so its official CLI binary (bundled_*.go, one
// per platform, go:embed'd behind build tags) ships baked into the Ferrum
// binary itself for Windows/Linux/macOS on amd64 or arm64 — no download, no
// FERRUM_NEEDLE_BIN, no manual step; resolveBinPath extracts it to a cache
// file on first use. On any other platform (32-bit, RISC-V, Windows/ARM64,
// ...) Ferrum still never fetches executable content from the network on
// its own: an operator there must download the CLI binary themselves and
// point FERRUM_NEEDLE_BIN at it — see README "Built-in LLM (Needle 2)" for
// the exact steps. Available() reports false, and every call fails with a
// clear "not installed" error, until a binary (bundled or configured) is in
// place.
//
// Needle's own HTTP server (`needle --serve`) is NOT OpenAI-compatible: it
// takes a single `{"input": "..."}` string per request (no message history,
@@ -28,6 +30,7 @@ import (
"context"
"encoding/json"
"fmt"
"net"
"net/http"
"os"
"os/exec"
@@ -46,16 +49,25 @@ const BaseURL = "needle://local"
// IsBuiltin reports whether baseURL names the built-in Needle provider.
func IsBuiltin(baseURL string) bool { return baseURL == BaseURL }
// listenAddr is where the Needle CLI's --serve mode listens. The README
// excerpt Ferrum was built against documents no --port flag (just "runs on
// localhost:8080"), so this is fixed rather than guessed at — a future
// Needle build that adds one can have this promoted to a Manager field.
const listenAddr = "127.0.0.1:8080"
// listenPort/listenAddr is where the Needle CLI's --serve mode listens.
// Needle defaults --serve to :8080, which collides with Ferrum's own
// default server.addr (also :8080, see config.example.yaml) — a Needle
// subprocess started after Ferrum would fail to bind and never become
// ready. `--port` (undocumented in the model card's README, but present in
// the actual CLI's --help) moves it out of the way; the chosen value just
// needs to avoid Ferrum's own port and other common local dev ports.
const listenPort = "58211"
const listenAddr = "127.0.0.1:" + listenPort
// startTimeout bounds how long ensureRunning waits for a freshly spawned
// process to answer its first request before giving up and reporting the
// binary as unusable this run.
const startTimeout = 10 * time.Second
// binary as unusable this run. Generous because it isn't just a TCP-accept
// wait: Needle's "tool retrieval" (README) does a one-time embedding pass
// over every declared tool on the process's first request once the catalog
// exceeds 5 tools — true for Ferrum's real catalog — so cold start plus
// first-inference warmup can run a few seconds even though the listener
// itself comes up almost immediately.
const startTimeout = 20 * time.Second
// Manager owns the lifecycle of at most one Needle CLI subprocess — "each
// component instance owns one conversation" per the model's own docs, and
@@ -74,15 +86,76 @@ func NewManager(binPath string) *Manager {
return &Manager{binPath: binPath}
}
// Available reports whether the configured binary exists — used to hide/
// disable the built-in provider in the UI and to fail fast with a clear
// message instead of an opaque connection error.
// syncBuffer is bytes.Buffer plus a mutex, so it's safe as an exec.Cmd's
// Stdout/Stderr (written from the subprocess-reading goroutines the os/exec
// package spawns internally) while ensureRunning's own goroutine reads it
// via String() to build an error message.
type syncBuffer struct {
mu sync.Mutex
buf bytes.Buffer
}
func (b *syncBuffer) Write(p []byte) (int, error) {
b.mu.Lock()
defer b.mu.Unlock()
return b.buf.Write(p)
}
func (b *syncBuffer) String() string {
b.mu.Lock()
defer b.mu.Unlock()
return strings.TrimSpace(b.buf.String())
}
// Available reports whether a usable binary exists — either the one an
// operator configured, or (see resolveBinPath) the one bundled for this
// platform — used to hide/disable the built-in provider in the UI and to
// fail fast with a clear message instead of an opaque connection error.
func (m *Manager) Available() bool {
if m.binPath == "" {
return false
_, err := m.resolveBinPath()
return err == nil
}
// resolveBinPath returns the binary path this Manager will actually run: an
// operator-configured FERRUM_NEEDLE_BIN wins outright (an explicit choice
// should never be silently overridden); otherwise, if this platform has one
// baked in via go:embed (see bundled_*.go — Windows/Linux/macOS on amd64 or
// arm64), it's written out to a cache file once and reused, so Ferrum works
// with zero manual download/config on the platforms it ships a binary for.
func (m *Manager) resolveBinPath() (string, error) {
if m.binPath != "" {
info, err := os.Stat(m.binPath)
if err != nil || info.IsDir() {
return "", fmt.Errorf("configured needle binary not found at %q", m.binPath)
}
return m.binPath, nil
}
info, err := os.Stat(m.binPath)
return err == nil && !info.IsDir()
if len(bundledBinary) == 0 {
return "", fmt.Errorf("no needle binary bundled for this platform")
}
return extractBundled()
}
// extractBundled writes the embedded binary to a stable cache path once,
// skipping the write if a file of the exact same size is already there
// (cheap enough to check every call, avoids re-extracting ~15MB on every
// startup). Not content-hashed — ponytail: a corrupted cache file that
// happens to match the size would be reused as-is; delete the cache
// directory by hand if that's ever suspected, a byte-for-byte checksum isn't
// worth it for a file this binary only ever writes itself.
func extractBundled() (string, error) {
dir := filepath.Join(os.TempDir(), "ferrum-needle-bin")
if err := os.MkdirAll(dir, 0o755); err != nil {
return "", fmt.Errorf("preparing needle cache dir: %w", err)
}
path := filepath.Join(dir, bundledName)
if info, err := os.Stat(path); err == nil && info.Size() == int64(len(bundledBinary)) {
return path, nil
}
if err := os.WriteFile(path, bundledBinary, 0o755); err != nil {
return "", fmt.Errorf("extracting bundled needle binary: %w", err)
}
return path, nil
}
// Close stops the subprocess, if running. Safe to call even if it was never
@@ -113,8 +186,9 @@ func (m *Manager) closeLocked() {
// process is spawned. Add real supervision if Needle proves flaky in
// practice; nothing observed in its docs suggests it will be.
func (m *Manager) ensureRunning(ctx context.Context) error {
if !m.Available() {
return fmt.Errorf("needle binary not found at %q — see README \"Built-in LLM (Needle 2)\" for install steps", m.binPath)
binPath, err := m.resolveBinPath()
if err != nil {
return fmt.Errorf("%w — see README \"Built-in LLM (Needle 2)\" for install steps", err)
}
m.mu.Lock()
@@ -129,7 +203,16 @@ func (m *Manager) ensureRunning(ctx context.Context) error {
return fmt.Errorf("writing needle tools file: %w", err)
}
cmd := exec.Command(m.binPath, "--tools", toolsPath, "--serve")
cmd := exec.Command(binPath, "--tools", toolsPath, "--serve", "--port", listenPort)
// Captured rather than discarded: without this, a subprocess that
// starts and immediately exits (bad args, missing runtime dep, port
// already taken from outside our own bookkeeping) fails completely
// silently — ensureRunning just spins until startTimeout with no clue
// why. logBuf is small (Needle logs one line on startup) so keeping it
// in memory for the process's lifetime is fine.
var logBuf syncBuffer
cmd.Stdout = &logBuf
cmd.Stderr = &logBuf
if err := cmd.Start(); err != nil {
return fmt.Errorf("starting needle: %w", err)
}
@@ -139,8 +222,10 @@ func (m *Manager) ensureRunning(ctx context.Context) error {
// so the next ensureRunning call notices and respawns instead of trying
// to reuse a dead process — Signal(0)-style liveness checks aren't
// reliably supported cross-platform (notably on Windows), but Wait() is.
exited := make(chan struct{})
go func() {
_ = cmd.Wait()
close(exited)
m.mu.Lock()
if m.cmd == cmd {
m.running = false
@@ -149,21 +234,71 @@ func (m *Manager) ensureRunning(ctx context.Context) error {
m.mu.Unlock()
}()
deadline := time.Now().Add(startTimeout)
for time.Now().Before(deadline) {
// Any completed HTTP round-trip — even a Needle-side error response
// — proves the server is accepting connections; only a transport
// failure (connection refused, still starting up) means "not ready".
pingCtx, cancel := context.WithTimeout(ctx, 500*time.Millisecond)
// Two phases, deliberately not one poll loop of full requests: dialing a
// bare TCP connection (no bytes sent) is cheap and safe to retry while
// the listener is still coming up, but a real /complete round-trip is
// not — Needle serves one request at a time, and a real request that
// gets cancelled client-side keeps running server-side, so retrying
// *those* on a short timeout just piles up work that never finishes
// (see startTimeout's doc comment on why that first request is slow).
// So: poll the bare socket until something answers, then send exactly
// one real request and let it run.
select {
case <-waitForListener(ctx, listenAddr, startTimeout):
case <-exited:
return fmt.Errorf("needle exited immediately: %s", logBuf.String())
}
pingCtx, cancel := context.WithTimeout(ctx, startTimeout)
pingErr := make(chan error, 1)
go func() {
_, _, err := doComplete(pingCtx, "ping")
pingErr <- err
}()
select {
case err := <-pingErr:
cancel()
if err == nil {
return nil
}
time.Sleep(150 * time.Millisecond)
m.closeLocked()
return fmt.Errorf("needle did not become ready within %s (%v): %s", startTimeout, err, logBuf.String())
case <-exited:
cancel()
// Crashed (or exited) before ever answering a request — no point
// waiting out the rest of startTimeout. logBuf carries whatever it
// printed (its own error, a missing dependency, license/arch
// mismatch, ...) so this isn't a bare "not ready".
return fmt.Errorf("needle exited immediately: %s", logBuf.String())
}
m.closeLocked()
return fmt.Errorf("needle did not become ready within %s", startTimeout)
}
// waitForListener returns a channel that closes as soon as addr accepts a
// bare TCP connection, or when timeout elapses (whichever first) — the
// caller distinguishes the two via the connection's own read/request
// afterward, so this never needs to report which happened. Only a raw dial
// is retried here; see ensureRunning for why a real request isn't.
func waitForListener(ctx context.Context, addr string, timeout time.Duration) <-chan struct{} {
ready := make(chan struct{})
go func() {
defer close(ready)
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
d := net.Dialer{Timeout: 200 * time.Millisecond}
conn, err := d.DialContext(ctx, "tcp", addr)
if err == nil {
conn.Close()
return
}
select {
case <-ctx.Done():
return
case <-time.After(50 * time.Millisecond):
}
}
}()
return ready
}
// --- request/response translation ---
+34 -2
View File
@@ -14,18 +14,49 @@ func TestIsBuiltin(t *testing.T) {
}
}
// withNoBundledBinary simulates running on a platform Ferrum doesn't embed
// a Needle binary for (see bundled_other.go), so tests of the "nothing
// available" path aren't at the mercy of which OS/arch actually runs them —
// every supported CI platform (windows/amd64, linux/amd64, linux/arm64,
// darwin/arm64) now has one baked in via go:embed.
func withNoBundledBinary(t *testing.T) {
t.Helper()
orig := bundledBinary
bundledBinary = nil
t.Cleanup(func() { bundledBinary = orig })
}
func TestAvailableFalseWhenBinaryMissing(t *testing.T) {
withNoBundledBinary(t)
m := NewManager("")
if m.Available() {
t.Fatal("expected Available() to be false with no configured path")
t.Fatal("expected Available() to be false with no configured path and no bundled binary")
}
m2 := NewManager("/does/not/exist/needle")
if m2.Available() {
t.Fatal("expected Available() to be false for a nonexistent path")
t.Fatal("expected Available() to be false for a nonexistent configured path")
}
}
func TestAvailableTrueWhenBundled(t *testing.T) {
if len(bundledBinary) == 0 {
t.Skip("no binary bundled for this platform")
}
m := NewManager("")
if !m.Available() {
t.Fatal("expected Available() to be true when a binary is bundled for this platform and no override is configured")
}
}
func TestExplicitBinPathWinsOverBundled(t *testing.T) {
m := NewManager("/does/not/exist/needle")
if m.Available() {
t.Fatal("an explicitly configured (but missing) path must not silently fall back to the bundled binary")
}
}
func TestChatCompletionFailsClearlyWhenNotInstalled(t *testing.T) {
withNoBundledBinary(t)
m := NewManager("")
_, status, err := m.ChatCompletion(context.Background(), []map[string]any{{"role": "user", "content": "hi"}}, nil)
if err == nil {
@@ -37,6 +68,7 @@ func TestChatCompletionFailsClearlyWhenNotInstalled(t *testing.T) {
}
func TestTestConnectionFailsClearlyWhenNotInstalled(t *testing.T) {
withNoBundledBinary(t)
m := NewManager("")
if _, err := m.TestConnection(context.Background()); err == nil {
t.Fatal("expected an error when the binary isn't installed")