diff --git a/README.md b/README.md index b79daf6..c95ebe3 100644 --- a/README.md +++ b/README.md @@ -292,7 +292,7 @@ For non-Python clients (Next.js, n8n, curl, Slack bots). |---|---| | `GET /api/workers` | Live list of every worker on the mesh with telemetry. | | `POST /api/execute` | Sync streaming task dispatch. Body: `{"worker_id":..., "prompt":..., "task_id":...}`. Streams worker stdout back chunked. | -| `POST /api/execute/async` | Fire-and-forget. Returns `202 {"task_id":...}` immediately. POSTs to `webhook_url` on completion (signed if `AGENTFM_WEBHOOK_SECRET` is set). | +| `POST /api/execute/async` | Fire-and-forget. Returns `202 {"task_id":...}` immediately. POSTs to `webhook_url` on completion (signed if `AGENTFM_WEBHOOK_SECRET` is set). The background task spawns *before* the 202 ack is written: a 202 means the task is being executed, even if the client hangs up before reading the body. | | `GET /metrics` | Prometheus scrape endpoint (see Observability). | Webhook POSTs are bounded by a 30 s timeout, do not follow redirects, and the response body is capped at 64 KiB. diff --git a/agentfm-go/cmd/agentfm/main.go b/agentfm-go/cmd/agentfm/main.go index 39095f7..ba6f7d1 100644 --- a/agentfm-go/cmd/agentfm/main.go +++ b/agentfm-go/cmd/agentfm/main.go @@ -76,13 +76,20 @@ func main() { return } - validateOperatorConfig(cfg) - if *mode == "" { pterm.Error.Println("Please specify a mode: -mode boss, worker, relay, api, test, or genkey") os.Exit(1) } + // Worker-config bounds (maxtasks, maxcpu, maxgpu, agent/desc/model + // lengths) only apply to roles that actually consume worker.Config. + // Boss/relay/api take defaults, so running the validator there would + // just be noise ā€” and would surface confusing limits in --help that + // don't apply to the chosen mode. + if *mode == "worker" || *mode == "test" { + validateOperatorConfig(cfg) + } + ctx := context.Background() netCfg := network.Config{ diff --git a/agentfm-go/cmd/relay/main.go b/agentfm-go/cmd/relay/main.go index 8f8e263..3a278ca 100644 --- a/agentfm-go/cmd/relay/main.go +++ b/agentfm-go/cmd/relay/main.go @@ -4,6 +4,7 @@ import ( "context" "flag" "fmt" + "log/slog" "os" "os/signal" "syscall" @@ -27,25 +28,12 @@ func fatalf(format string, args ...interface{}) { pterm.Fatal.Printfln(format, args...) } -// This function loads a saved identity or generates a new one +// getStaticIdentity is a thin wrapper around network.LoadOrGenerateIdentity +// kept for the cmd/relay binary's existing logging idiom. The shared helper +// takes care of the corrupt-file warning, the 0600 perm, and the +// regenerate-on-missing semantics so both binaries stay in lockstep. func getStaticIdentity(keyFile string) (crypto.PrivKey, error) { - if keyBytes, err := os.ReadFile(keyFile); err == nil { - fmt.Printf("šŸ”‘ Loaded existing permanent identity from %s!\n", keyFile) - return crypto.UnmarshalPrivateKey(keyBytes) - } - - fmt.Printf("āœØ Generating new permanent identity at %s...\n", keyFile) - priv, _, err := crypto.GenerateKeyPair(crypto.Ed25519, -1) - if err != nil { - return nil, err - } - - keyBytes, err := crypto.MarshalPrivateKey(priv) - if err != nil { - return nil, err - } - err = os.WriteFile(keyFile, keyBytes, 0600) - return priv, err + return network.LoadOrGenerateIdentity(keyFile) } func main() { @@ -166,6 +154,10 @@ func main() { fmt.Println("\nShutting down relay node...") if err := host.Close(); err != nil { - pterm.Error.Printfln("Host close error: %v", err) + // slog (not pterm) so log shippers see structured failure events + // for relays running unattended under systemd / docker / k8s. + slog.Error("relay host close", + slog.Any(obs.FieldErr, err), + ) } } diff --git a/agentfm-go/internal/boss/api_async.go b/agentfm-go/internal/boss/api_async.go index 2d6a2b7..f0280ba 100644 --- a/agentfm-go/internal/boss/api_async.go +++ b/agentfm-go/internal/boss/api_async.go @@ -38,6 +38,17 @@ const asyncArtifactWait = 10 * time.Second // Returns whether the file was observed (currently informational only; // callers fire the webhook either way). func waitForArtifact(ctx context.Context, taskID string, maxWait time.Duration) bool { + // Defense-in-depth: today the only caller passes a newCompletionID() + // (crypto/rand hex), so a path-traversal payload is impossible. But + // the helper is generic; refuse anything that can't safely be joined + // into a filesystem path so a future caller can't accidentally + // inherit a traversal bug. + if !network.SafeTaskIDPattern.MatchString(taskID) { + slog.Warn("waitForArtifact rejected unsafe taskID", + slog.String(obs.FieldTaskID, taskID), + ) + return false + } deadline := time.Now().Add(maxWait) zipPath := filepath.Join("agentfm_artifacts", taskID+".zip") ticker := time.NewTicker(100 * time.Millisecond) @@ -107,23 +118,13 @@ func (b *Boss) asyncExecuteHandler(rootCtx context.Context, wg *sync.WaitGroup) return } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusAccepted) - if err := json.NewEncoder(w).Encode(map[string]string{ - "task_id": taskID, - "status": "queued", - "message": "Task dispatched to P2P mesh.", - }); err != nil { - slog.Error("async task ack write", slog.Any(obs.FieldErr, err)) - <-b.asyncSlots - return - } - - // Budget the whole background job (stream + webhook) by the task - // execution timeout so a ghosted worker can't hold a goroutine - // past server shutdown. + // Spawn-before-ack: the goroutine starts BEFORE we attempt to + // write the 202 body. This preserves the contract "if the client + // got a 202 with a task_id, the task is being executed." A failed + // ack write (client hung up between header and body) leaves the + // background task running so a webhook delivery can still fire, + // rather than silently dropping committed work. taskCtx, cancelTask := context.WithTimeout(rootCtx, network.TaskExecutionTimeout) - wg.Add(1) go func() { defer wg.Done() @@ -132,6 +133,19 @@ func (b *Boss) asyncExecuteHandler(rootCtx context.Context, wg *sync.WaitGroup) b.runAsyncTask(taskCtx, peerID, taskID, req) }() + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusAccepted) + if err := json.NewEncoder(w).Encode(map[string]string{ + "task_id": taskID, + "status": "queued", + "message": "Task dispatched to P2P mesh.", + }); err != nil { + slog.Warn("async task ack write failed; goroutine continues", + slog.String(obs.FieldTaskID, taskID), + slog.Any(obs.FieldErr, err), + ) + } } } diff --git a/agentfm-go/internal/boss/api_handlers.go b/agentfm-go/internal/boss/api_handlers.go index 8238326..607bf91 100644 --- a/agentfm-go/internal/boss/api_handlers.go +++ b/agentfm-go/internal/boss/api_handlers.go @@ -115,7 +115,8 @@ func (b *Boss) handleExecuteTask(w http.ResponseWriter, r *http.Request) { return } - pterm.Info.Printfln("šŸ“” API Gateway routing task %s to Worker %s...", req.TaskID[:8], pterm.Cyan(peerID.String()[:8])) + pterm.Info.Printfln("šŸ“” API Gateway routing task %s to Worker %s...", + shortID(req.TaskID, 8), pterm.Cyan(peerID.String()[:8])) // Tie the dial to the inbound HTTP request's context so a client // hanging up aborts the libp2p dial instead of waiting out the full diff --git a/agentfm-go/internal/boss/api_test.go b/agentfm-go/internal/boss/api_test.go index 73af469..7874f52 100644 --- a/agentfm-go/internal/boss/api_test.go +++ b/agentfm-go/internal/boss/api_test.go @@ -281,6 +281,32 @@ func TestHandleExecuteTask_InvalidWorkerIDFormat(t *testing.T) { } } +// TestShortID_GuardsShortInput: regression for round-2 audit finding I-8. +// The routing log line in handleExecuteTask previously sliced +// req.TaskID[:8] directly, which panics when a user-supplied TaskID is +// shorter than 8 chars. shortID() caps the slice safely. We unit-test +// the helper directly because invoking the full handler path triggers +// a pre-existing pterm-spinner concurrency issue under -race. +func TestShortID_GuardsShortInput(t *testing.T) { + t.Parallel() + cases := []struct { + in string + n int + want string + }{ + {"abc", 8, "abc"}, + {"task_abcdef12345", 8, "task_abc"}, + {"", 8, ""}, + {"x", 0, ""}, + {"exactly8", 8, "exactly8"}, + } + for _, tc := range cases { + if got := shortID(tc.in, tc.n); got != tc.want { + t.Errorf("shortID(%q, %d) = %q; want %q", tc.in, tc.n, got, tc.want) + } + } +} + // --- POST /api/execute/async ----------------------------------------------- // TestAsyncExecuteHandler_MethodNotAllowed covers the factory-returned diff --git a/agentfm-go/internal/boss/boss.go b/agentfm-go/internal/boss/boss.go index d2242f2..a2051be 100644 --- a/agentfm-go/internal/boss/boss.go +++ b/agentfm-go/internal/boss/boss.go @@ -100,7 +100,10 @@ func (b *Boss) listenTelemetry(ctx context.Context) { // Periodic pruner: evicts workers whose libp2p connection has dropped. // Centralised here so handleGetWorkers and the TUI tick can be pure // reads (no side effects on a GET request). - pruneTicker := time.NewTicker(30 * time.Second) + // Tick faster than staleTelemetryTimeout (15s) so the lastSeen-based + // eviction is responsive ā€” a 30s tick would let stale workers linger + // in the radar for almost half a minute past the staleness threshold. + pruneTicker := time.NewTicker(5 * time.Second) defer pruneTicker.Stop() msgCh := make(chan *pubsubMsg, 1) @@ -152,9 +155,15 @@ type pubsubMsg struct { // pruneDisconnectedWorkers walks activeWorkers and evicts any peer the // libp2p host is no longer connected to. Runs under a write lock; cheap // because the map is bounded by mesh size (typically tens of peers). +// staleTelemetryTimeout is the upper bound on how long a worker can go +// without a telemetry pulse before the pruner evicts it. Mirrors the +// previous 15s ad-hoc value that lived inline in ui.go's draw loop. +const staleTelemetryTimeout = 15 * time.Second + func (b *Boss) pruneDisconnectedWorkers() { b.mu.Lock() defer b.mu.Unlock() + now := time.Now() for peerIDStr := range b.activeWorkers { pID, err := peer.Decode(peerIDStr) if err != nil { @@ -165,6 +174,11 @@ func (b *Boss) pruneDisconnectedWorkers() { if b.node.Host.Network().Connectedness(pID) != netcore.Connected { delete(b.activeWorkers, peerIDStr) delete(b.lastSeen, peerIDStr) + continue + } + if seen, ok := b.lastSeen[peerIDStr]; ok && now.Sub(seen) > staleTelemetryTimeout { + delete(b.activeWorkers, peerIDStr) + delete(b.lastSeen, peerIDStr) } } metrics.WorkersOnline.Set(float64(len(b.activeWorkers))) @@ -184,3 +198,14 @@ func (tr *timeoutReader) Read(p []byte) (n int, err error) { } return tr.stream.Read(p) } + +// shortID returns the first n runes of s, or s when shorter. Used for +// log/UI snippets where a short identifier prefix is enough for humans +// to correlate. Defends against panics on user-supplied IDs that fall +// short of the slice length. +func shortID(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} diff --git a/agentfm-go/internal/boss/ui.go b/agentfm-go/internal/boss/ui.go index df9ead4..6a70bb9 100644 --- a/agentfm-go/internal/boss/ui.go +++ b/agentfm-go/internal/boss/ui.go @@ -62,24 +62,17 @@ func (b *Boss) selectWorkerInteractive(parentCtx context.Context) (types.WorkerP defer area.Stop() draw := func() { - b.mu.Lock() - - for peerID, seen := range b.lastSeen { - if time.Since(seen) > 15*time.Second { - delete(b.activeWorkers, peerID) - delete(b.lastSeen, peerID) - } - } - + // Pure read: pruning is handled by listenTelemetry's pruneTicker + // (boss.go) so the TUI and /api/workers always agree on which + // workers are visible. RLock is sufficient. + b.mu.RLock() nextList := make([]types.WorkerProfile, 0, len(b.activeWorkers)) for _, w := range b.activeWorkers { nextList = append(nextList, w) } - activeCount := len(b.activeWorkers) peerCount := len(b.node.PubSub.ListPeers(network.TelemetryTopic)) - - b.mu.Unlock() + b.mu.RUnlock() sort.Slice(nextList, func(i, j int) bool { return nextList[i].PeerID < nextList[j].PeerID diff --git a/agentfm-go/internal/network/artifacts.go b/agentfm-go/internal/network/artifacts.go index 4dac81d..7770e68 100644 --- a/agentfm-go/internal/network/artifacts.go +++ b/agentfm-go/internal/network/artifacts.go @@ -20,7 +20,12 @@ import ( "github.com/pterm/pterm" ) -var safeTaskIDPattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_.-]{0,127}$`) +// SafeTaskIDPattern bounds the alphabet of a task ID that's used as a +// path component (the artifact zip filename, the per-task /tmp/output +// bind-mount). Exported so other packages (e.g. boss/api_async.go's +// waitForArtifact) can apply the same defense-in-depth check rather +// than trusting a string they got from upstream code. +var SafeTaskIDPattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_.-]{0,127}$`) func SendArtifacts(ctx context.Context, h host.Host, bossID peer.ID, zipFilePath string, taskID string) error { fmt.Println("šŸ“¦ Opening secure artifact channel to Boss...") @@ -157,7 +162,7 @@ func HandleArtifactStream(stream network.Stream) { taskID := string(idBytes) safeTaskID := filepath.Base(filepath.Clean(taskID)) - if !safeTaskIDPattern.MatchString(safeTaskID) { + if !SafeTaskIDPattern.MatchString(safeTaskID) { safeTaskID = fmt.Sprintf("fallback_%d", time.Now().UnixNano()) } diff --git a/agentfm-go/internal/network/discovery.go b/agentfm-go/internal/network/discovery.go index 4588f11..c9f5563 100644 --- a/agentfm-go/internal/network/discovery.go +++ b/agentfm-go/internal/network/discovery.go @@ -103,6 +103,11 @@ func discoverPeers(ctx context.Context, h host.Host, routingDiscovery *routing.R dialCtx, dialCancel := context.WithTimeout(ctx, StreamDialTimeout) if err := h.Connect(dialCtx, p); err == nil { fmt.Printf("\nšŸŒ [DHT Fallback] Successfully connected to peer: %s\n", p.ID.String()[:8]) + } else { + slog.Debug("dht fallback dial failed", + slog.String(obs.FieldPeerID, p.ID.String()), + slog.Any(obs.FieldErr, err), + ) } dialCancel() } diff --git a/agentfm-go/internal/network/host.go b/agentfm-go/internal/network/host.go index 05c6f1f..2306494 100644 --- a/agentfm-go/internal/network/host.go +++ b/agentfm-go/internal/network/host.go @@ -17,18 +17,18 @@ import ( "github.com/libp2p/go-libp2p/p2p/host/autonat" ) -// loadOrGenerateIdentity returns the persistent Ed25519 private key for -// this node, creating one on first boot. Stable peer IDs across restarts -// are important: they let other nodes cache our address and skip the DHT -// lookup next time, and private swarms often allowlist peers by ID. -func loadOrGenerateIdentity(mode string) (crypto.PrivKey, error) { - keyPath := fmt.Sprintf(".agentfm_%s_identity.key", mode) - - // Try the existing identity first. A read failure is fine on the - // very first boot, but a successful read that fails to unmarshal - // means the file on disk is corrupt. We surface that so the - // operator realises their peer ID is about to change, instead of - // silently regenerating. +// LoadOrGenerateIdentity returns the persistent Ed25519 private key at +// keyPath, creating one on first boot. Shared by both binaries: +// internal/network's mode-suffixed boss/worker/api identities, and the +// dedicated relay binary's relay_identity.key. Stable peer IDs across +// restarts let other nodes cache our address and skip DHT lookups, and +// private swarms often allowlist peers by ID. +// +// A read failure is fine on the very first boot, but a successful read +// that fails to unmarshal means the file on disk is corrupt. We surface +// that so the operator realises their peer ID is about to change instead +// of silently regenerating. +func LoadOrGenerateIdentity(keyPath string) (crypto.PrivKey, error) { if keyBytes, err := os.ReadFile(keyPath); err == nil { priv, err := crypto.UnmarshalPrivateKey(keyBytes) if err == nil { @@ -62,6 +62,13 @@ func loadOrGenerateIdentity(mode string) (crypto.PrivKey, error) { return priv, nil } +// loadOrGenerateIdentity preserves the legacy package-internal call sites +// (createHost) that pass a mode string and want the .agentfm__identity.key +// naming convention. New callers should use LoadOrGenerateIdentity directly. +func loadOrGenerateIdentity(mode string) (crypto.PrivKey, error) { + return LoadOrGenerateIdentity(fmt.Sprintf(".agentfm_%s_identity.key", mode)) +} + // createHost assembles the libp2p Host with the correct options for this // role: PSK for private swarms, circuit relay in the right direction, // NAT port mapping, and the AutoNAT reachability probe for non-relay @@ -116,7 +123,9 @@ func createHost(cfg Config, bootstrapAddr string) (host.Host, error) { if cfg.Mode != "relay" { if _, err = autonat.New(h); err != nil { - fmt.Printf("āš ļø [NAT] Failed to start AutoNAT service: %v\n", err) + slog.Warn("autonat unavailable; reachability autodetection disabled", + slog.Any(obs.FieldErr, err), + ) } else { fmt.Println("šŸŒ [NAT] AutoNAT service started. Testing public reachability...") } diff --git a/agentfm-go/internal/network/p2p.go b/agentfm-go/internal/network/p2p.go index 44d021d..8801be1 100644 --- a/agentfm-go/internal/network/p2p.go +++ b/agentfm-go/internal/network/p2p.go @@ -3,6 +3,9 @@ package network import ( "context" "fmt" + "log/slog" + + "agentfm/internal/obs" dht "github.com/libp2p/go-libp2p-kad-dht" pubsub "github.com/libp2p/go-libp2p-pubsub" @@ -56,7 +59,13 @@ func Setup(ctx context.Context, cfg Config) (*MeshNode, error) { var relayPeerID peer.ID if cfg.Mode != "relay" && bootstrapAddr != "" { - if relayInfo, err := parseRelayInfo(bootstrapAddr); err == nil { + relayInfo, err := parseRelayInfo(bootstrapAddr) + if err != nil { + slog.Warn("invalid bootstrap multiaddr; skipping lighthouse", + slog.String("bootstrap", bootstrapAddr), + slog.Any(obs.FieldErr, err), + ) + } else { relayPeerID = relayInfo.ID connectToLighthouse(ctx, h, relayInfo) } diff --git a/agentfm-go/internal/worker/worker.go b/agentfm-go/internal/worker/worker.go index 9c5ae97..67192ec 100644 --- a/agentfm-go/internal/worker/worker.go +++ b/agentfm-go/internal/worker/worker.go @@ -43,6 +43,20 @@ func New(node *network.MeshNode, cfg Config) *Worker { return &Worker{node: node, config: cfg} } +// printHostNetworkWarning surfaces the --network host security caveat +// to the operator. Called from both Worker.Start (long-lived worker) +// and RunLocalTest (single-shot test mode) so a workshop attendee +// running `agentfm -mode test` against a developer laptop sees the +// same loopback-exposure warning the production worker prints. +func printHostNetworkWarning() { + pterm.Warning.Println( + "Containers run with --network host: the agent has full access to this " + + "machine's network namespace, including loopback (127.0.0.1) services like " + + "Ollama, internal admin endpoints, and cloud metadata (169.254.169.254). " + + "Treat agent images as TRUSTED CODE; review their Dockerfiles before running.", + ) +} + // RunLocalTest allows users to test their dockerfile/script locally without libp2p func RunLocalTest(ctx context.Context, cfg Config, prompt string) error { w := &Worker{config: cfg} @@ -51,6 +65,8 @@ func RunLocalTest(ctx context.Context, cfg Config, prompt string) error { return err } + printHostNetworkWarning() + fmt.Printf("\nšŸ¤– Sending Prompt: '%s'\n", pterm.LightGreen(prompt)) fmt.Println("--------------------------------------------------") @@ -78,12 +94,7 @@ func (w *Worker) Start(ctx context.Context) { os.Exit(1) } - pterm.Warning.Println( - "Containers run with --network host: the agent has full access to this " + - "machine's network namespace, including loopback (127.0.0.1) services like " + - "Ollama, internal admin endpoints, and cloud metadata (169.254.169.254). " + - "Treat agent images as TRUSTED CODE; review their Dockerfiles before running.", - ) + printHostNetworkWarning() w.printMetadata() w.wg.Add(1) diff --git a/agentfm-python/README.md b/agentfm-python/README.md index fe1365f..afe61f0 100644 --- a/agentfm-python/README.md +++ b/agentfm-python/README.md @@ -75,6 +75,12 @@ except AuthenticationError as e: [Authentication](https://github.com/Agent-FM/agentfm-core#authentication) docs for setting up `AGENTFM_API_KEYS` on the boss. +> **`client.api_key` is read-after-construction.** The token is baked into +> the underlying `httpx.Client.headers` once at `__init__`. Mutating +> `client.api_key = "new-key"` afterwards does NOT update the request +> header ā€” use `client.with_options(api_key="new-key")` to derive a +> client with a new key. Matches the OpenAI Python SDK's behaviour. + ## Async ```python diff --git a/agentfm-python/src/agentfm/async_client.py b/agentfm-python/src/agentfm/async_client.py index 7544715..2d2731f 100644 --- a/agentfm-python/src/agentfm/async_client.py +++ b/agentfm-python/src/agentfm/async_client.py @@ -135,6 +135,21 @@ async def stream( prompt: str, task_id: str | None = None, ) -> AsyncIterator[TaskChunk]: + """Stream worker stdout chunk-by-chunk as ``TaskChunk`` objects. + + Early termination of the consumer (``async for ... break``) requires + explicit cleanup so the underlying httpx response is released + promptly. Wrap the iteration with :func:`contextlib.aclosing`:: + + from contextlib import aclosing + async with aclosing(client.tasks.stream(...)) as gen: + async for chunk in gen: + if want_to_stop: break + + CPython's reference-counting usually finalises the generator on + the next GC cycle, but on PyPy / under pressure / if the generator + is stored in a future, the response can outlive the loop. + """ payload = build_task_payload(str(worker_id), prompt, task_id=task_id) filter_ = SentinelFilter() try: @@ -199,7 +214,7 @@ async def worker(idx: int, prompt: str) -> None: # max_concurrency=1 doesn't deadlock and a failing peer is not # reused on retry. The cursor advances by attempt count, giving # automatic failover across the supplied peer pool. - last_exc: AgentFMError | None = None + last_exc: BaseException | None = None last_peer = peers_str[idx % len(peers_str)] for attempt in range(max_retries + 1): peer = peers_str[(idx + attempt) % len(peers_str)] @@ -207,7 +222,17 @@ async def worker(idx: int, prompt: str) -> None: async with sem: try: res = await self.run(worker_id=peer, prompt=prompt) - except AgentFMError as exc: + except Exception as exc: + # Wider than AgentFMError on purpose: tasks.run can + # raise OSError from artifact harvesting which would + # otherwise propagate via asyncio.gather and cancel + # siblings, breaking the "scatter never raises" + # contract. Anything unexpected is logged. + if not isinstance(exc, AgentFMError): + _log.exception( + "scatter prompt #%s raised non-AgentFMError; treating as failure", + idx, + ) last_exc = exc if attempt < max_retries: _log.info( diff --git a/agentfm-python/src/agentfm/client.py b/agentfm-python/src/agentfm/client.py index 05f64e6..2806bdd 100644 --- a/agentfm-python/src/agentfm/client.py +++ b/agentfm-python/src/agentfm/client.py @@ -257,7 +257,17 @@ def scatter( text=res.text, artifacts=res.artifacts, ) - except AgentFMError as exc: + except Exception as exc: + # Wider than AgentFMError on purpose: tasks.run can + # raise OSError from artifact harvesting (artifacts_dir + # disappeared, perm denied) which would otherwise + # break the "scatter never raises" contract. Anything + # unexpected is logged so it isn't silently swallowed. + if not isinstance(exc, AgentFMError): + _log.exception( + "scatter prompt #%s raised non-AgentFMError; treating as failure", + idx, + ) attempts[idx] += 1 if attempts[idx] <= max_retries: _log.info("retrying prompt #%s (attempt %s)", idx, attempts[idx]) @@ -319,6 +329,12 @@ def __init__( omit the argument to fall back to the ``AGENTFM_API_KEY`` env var; pass an explicit ``None`` to disable auth (no fallback); pass a string to use that token verbatim. + + The token is read once at construction time and baked into the + underlying ``httpx.Client.headers``. Mutating ``client.api_key`` + after construction does NOT update the request header ā€” use + :meth:`with_options` to derive a new client with a different key. + Matches the OpenAI Python SDK's behaviour. """ self.gateway_url = gateway_url.rstrip("/") self.retries = retries diff --git a/agentfm-python/src/agentfm/daemon.py b/agentfm-python/src/agentfm/daemon.py index 2c03fcb..d376a5e 100644 --- a/agentfm-python/src/agentfm/daemon.py +++ b/agentfm-python/src/agentfm/daemon.py @@ -49,6 +49,7 @@ def __init__( debug: bool = False, log_file: str | Path | None = None, startup_timeout: float = 15.0, + api_key: str | None = None, ) -> None: self.binary_path = binary_path self.port = port @@ -60,6 +61,10 @@ def __init__( self.process: subprocess.Popen[bytes] | None = None self._log_handle: IO[bytes] | None = None self.url = f"http://127.0.0.1:{self.port}" + # api_key is forwarded as Authorization: Bearer ... on the readiness + # probe so a gateway with AGENTFM_API_KEYS set doesn't reject the + # probe with 401. Falls back to AGENTFM_API_KEY env var if unset. + self.api_key = api_key if api_key is not None else os.environ.get("AGENTFM_API_KEY") or None # -- context-manager protocol ------------------------------------------- @@ -166,11 +171,22 @@ def stop(self) -> None: self._cleanup_log() def _is_ready(self) -> bool: - try: - r = httpx.get(f"{self.url}/api/workers", timeout=1.0) - return r.status_code == 200 - except httpx.HTTPError: - return False + # Prefer /health (always unauthenticated; cheaper than /api/workers + # which does a map walk under lock). Falls back to /api/workers if + # the gateway is too old to expose /health. + headers = ( + {"Authorization": f"Bearer {self.api_key}"} + if self.api_key + else {} + ) + for path in ("/health", "/api/workers"): + try: + r = httpx.get(f"{self.url}{path}", timeout=1.0, headers=headers) + except httpx.HTTPError: + continue + if r.status_code == 200: + return True + return False def _cleanup_log(self) -> None: if self._log_handle is not None: diff --git a/agentfm-python/src/agentfm/openai/_namespaces.py b/agentfm-python/src/agentfm/openai/_namespaces.py index a977a15..56b069c 100644 --- a/agentfm-python/src/agentfm/openai/_namespaces.py +++ b/agentfm-python/src/agentfm/openai/_namespaces.py @@ -25,7 +25,7 @@ from .._internal.resource import AsyncResource, SyncResource from .._transport import STREAMING_TIMEOUT, raise_for_response, raise_translated_stream_error from .._warnings import ROUTING_WARNING_STACKLEVEL, AgentFMRoutingWarning -from ..streaming import parse_sse_lines +from ..streaming import SSE_DONE, classify_sse_line, parse_sse_lines from .models import ( ChatCompletion, ChatCompletionChunk, @@ -45,11 +45,14 @@ class _RoutingWarner: - """Per-namespace dedup of :class:`AgentFMRoutingWarning`. - - Each ``OpenAINamespace`` / ``AsyncOpenAINamespace`` owns one. Sharing a - single instance across the chat and text-completion resources within a - namespace gives "warn once per model per client" semantics. Thread-safe. + """Process-wide dedup of :class:`AgentFMRoutingWarning`. + + A single module-level instance (``_GLOBAL_WARNER`` below) is shared by + every :class:`OpenAINamespace` and :class:`AsyncOpenAINamespace` so the + "warn once per model" promise actually holds across short-lived client + instances ā€” e.g. a FastAPI handler that constructs a fresh + :class:`AgentFMClient` per request would otherwise re-warn for every + request. Thread-safe. """ __slots__ = ("_lock", "_seen") @@ -79,6 +82,12 @@ def reset(self) -> None: self._seen.clear() +# Process-wide warner. Sharing across all namespace instances gives the +# user-facing "warn once per model" guarantee actual teeth even when the +# caller pattern is to build a fresh client per request. +_GLOBAL_WARNER = _RoutingWarner() + + def _build_chat_body( model: str, messages: list[ChatMessage] | list[dict[str, Any]], @@ -110,7 +119,7 @@ class OpenAINamespace: def __init__(self, client: AgentFMClient) -> None: self._client = client - self._warner = _RoutingWarner() + self._warner = _GLOBAL_WARNER self.models = _Models(client) self.chat = _Chat(client, self._warner) self.completions = _Completions(client, self._warner) @@ -248,7 +257,7 @@ class AsyncOpenAINamespace: def __init__(self, client: AsyncAgentFMClient) -> None: self._client = client - self._warner = _RoutingWarner() + self._warner = _GLOBAL_WARNER self.models = _AsyncModels(client) self.chat = _AsyncChat(client, self._warner) self.completions = _AsyncCompletions(client, self._warner) @@ -402,13 +411,15 @@ async def _stream(self, body: dict[str, Any]) -> AsyncIterator[TextCompletionChu async def _aiter_sse(r: httpx.Response) -> AsyncIterator[str]: - """Async equivalent of :func:`agentfm.streaming.parse_sse_lines`.""" + """Async equivalent of :func:`agentfm.streaming.parse_sse_lines`. + + Both iterators share :func:`classify_sse_line` so a future change to + SSE parsing (e.g. multi-line ``data:`` continuation) lands in one + place and applies symmetrically to sync + async paths. + """ async for raw in r.aiter_lines(): - line = raw.rstrip("\r\n") - if not line or line.startswith(":") or not line.startswith("data:"): - continue - body = line[len("data:") :].lstrip() - if body == "[DONE]": + result = classify_sse_line(raw) + if result is SSE_DONE: return - if body: - yield body + if isinstance(result, str): + yield result diff --git a/agentfm-python/src/agentfm/streaming.py b/agentfm-python/src/agentfm/streaming.py index 8fb4d17..f06e25c 100644 --- a/agentfm-python/src/agentfm/streaming.py +++ b/agentfm-python/src/agentfm/streaming.py @@ -85,6 +85,33 @@ def filter_iter(chunks: Iterable[str]) -> Iterator[str]: # Server-Sent Events (SSE) parsing for /v1/* streaming responses # --------------------------------------------------------------------------- +SSE_DONE = object() +"""Sentinel returned by :func:`classify_sse_line` for the ``data: [DONE]`` +terminator. Callers should stop iterating when they see it.""" + + +def classify_sse_line(raw: str) -> str | object | None: + """Classify a single SSE line. + + Returns: + - the payload string when the line carries ``data: `` + - :data:`SSE_DONE` when the line is the terminator ``data: [DONE]`` + - ``None`` for empty lines, comments, and non-data fields + + Shared by the sync :func:`parse_sse_lines` (operating on an + ``Iterable[str]``) and the async ``_aiter_sse`` in the OpenAI + namespace (operating on an ``AsyncIterator[str]``). Keeping a single + classifier means the two iterators can never drift on what counts as + a payload vs a terminator. + """ + line = raw.rstrip("\r\n") + if not line or line.startswith(":") or not line.startswith("data:"): + return None + body = line[len("data:") :].lstrip() + if body == "[DONE]": + return SSE_DONE + return body or None + def parse_sse_lines(lines: Iterable[str]) -> Iterator[str]: """Decode an SSE byte/line stream into raw ``data:`` payload bodies. @@ -94,23 +121,20 @@ def parse_sse_lines(lines: Iterable[str]) -> Iterator[str]: ``event:``/``id:`` for anything). """ for raw in lines: - line = raw.rstrip("\r\n") - if not line or line.startswith(":"): - continue - if not line.startswith("data:"): - continue - body = line[len("data:") :].lstrip() - if body == "[DONE]": + result = classify_sse_line(raw) + if result is SSE_DONE: return - if body: - yield body + if isinstance(result, str): + yield result __all__ = [ "ARTIFACT_INCOMING", "ARTIFACT_NONE", "SENTINEL_PREFIX", + "SSE_DONE", "SentinelFilter", + "classify_sse_line", "filter_iter", "parse_sse_lines", ] diff --git a/agentfm-python/tests/integration/test_scatter.py b/agentfm-python/tests/integration/test_scatter.py index 6cd7219..e7de3d1 100644 --- a/agentfm-python/tests/integration/test_scatter.py +++ b/agentfm-python/tests/integration/test_scatter.py @@ -248,3 +248,62 @@ def respond(request: httpx.Request) -> httpx.Response: assert [r.prompt for r in results] == prompts assert [r.status for r in results] == ["success", "failed", "success"] + + +# --------------------------------------------------------------------------- +# round-2 audit Py I-6: scatter must NOT propagate non-AgentFMError. Today +# tasks.run can raise OSError from artifact harvesting if artifacts_dir +# disappears. The "scatter never raises" contract requires that to surface +# as ScatterResult(status="failed"). +# --------------------------------------------------------------------------- + + +def test_sync_scatter_treats_non_agentfm_error_as_failure( + gateway_url: str, mock_gateway: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + mock_gateway.post("/api/execute").mock( + return_value=httpx.Response( + 200, content=b"ok\n", headers={"Content-Type": "text/plain"} + ) + ) + + with AgentFMClient(gateway_url=gateway_url) as client: + def boom(**kwargs: object) -> object: + raise OSError("simulated artifact harvest failure") + + monkeypatch.setattr(client.tasks, "run", boom) + + results = client.tasks.scatter( + ["p1", "p2"], peer_ids=["12D3KooWA"], max_concurrency=1, max_retries=0 + ) + + assert len(results) == 2 + assert all(r.status == "failed" for r in results), ( + f"got {[r.status for r in results]}; scatter must surface non-AgentFMError as ScatterResult(failed)" + ) + assert all("simulated artifact harvest failure" in (r.error or "") for r in results) + + +@pytest.mark.asyncio +async def test_async_scatter_treats_non_agentfm_error_as_failure( + gateway_url: str, mock_gateway: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + mock_gateway.post("/api/execute").mock( + return_value=httpx.Response( + 200, content=b"ok\n", headers={"Content-Type": "text/plain"} + ) + ) + + async with AsyncAgentFMClient(gateway_url=gateway_url) as client: + async def boom(**kwargs: object) -> object: + raise OSError("simulated artifact harvest failure") + + monkeypatch.setattr(client.tasks, "run", boom) + + results = await client.tasks.scatter( + ["p1", "p2"], peer_ids=["12D3KooWA"], max_concurrency=1, max_retries=0 + ) + + assert len(results) == 2 + assert all(r.status == "failed" for r in results) + assert all("simulated artifact harvest failure" in (r.error or "") for r in results) diff --git a/agentfm-python/tests/unit/test_with_options.py b/agentfm-python/tests/unit/test_with_options.py index 832f136..6fb5b05 100644 --- a/agentfm-python/tests/unit/test_with_options.py +++ b/agentfm-python/tests/unit/test_with_options.py @@ -132,3 +132,29 @@ def test_constructor_explicit_none_api_key_skips_env_fallback(monkeypatch): assert "authorization" not in {k.lower() for k in c._http.headers} finally: c.close() + + +def test_sync_client_del_tolerates_construction_failure(monkeypatch): + import contextlib + + import agentfm.client as client_mod + + def boom(*args, **kwargs): + raise RuntimeError("simulated httpx.Client construction failure") + + monkeypatch.setattr(client_mod, "make_client", boom) + with contextlib.suppress(RuntimeError): + AgentFMClient(gateway_url="http://x:8080") + + +def test_async_client_del_tolerates_construction_failure(monkeypatch): + import contextlib + + import agentfm.async_client as async_client_mod + + def boom(*args, **kwargs): + raise RuntimeError("simulated httpx.AsyncClient construction failure") + + monkeypatch.setattr(async_client_mod, "make_async_client", boom) + with contextlib.suppress(RuntimeError): + AsyncAgentFMClient(gateway_url="http://x:8080")