diff --git a/README.md b/README.md index ccf05f3..4636b9b 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,9 @@ From source: jump to [Install from source](#install-from-source). `,` and `?` are **conversational** — each terminal pane keeps its own thread. Type `, the same but with json output` and the model knows what "the same" means. After 30 min idle the thread auto-rotates so a forgotten tab doesn't bleed stale context. -`?` has tools: read files (filesystem-gated), DuckDuckGo search, fetch URL. Force web-first with `? --web `. +`,,` (alias for `, --fix`) repairs whatever just failed: it sends the previous command, its exit status, and — at higher [terminal context](#terminal-context) levels — recent output, and proposes corrected commands through the same picker. + +`?` has tools: read files (filesystem-gated), DuckDuckGo search, fetch URL. Force web-first with `? --web `. It's also pipe-friendly — `make 2>&1 | ? what broke` turns the piped output into context. `???` is the durable layer. Bare queries hit the archive of every past session (BM25; semantic if you've added embeddings). Flags pin long-term facts that get injected into every `?` system prompt. @@ -85,6 +87,18 @@ huggingface-cli download unsloth/Qwen3-Coder-Next-GGUF `??` finds the GGUF inside your HuggingFace cache — no path config required. +Each tier binds to its own port (`fast` :8091, `balanced` :8080, `smart` :8093), so tiers can serve **side by side** — and a single call can be routed to whichever fits: + +```sh +?? --start fast # fast tier up, alongside balanced +, --fast rename all .jpeg to .jpg # this one call uses the fast tier +? --smart why is my Makefile rebuilding everything # this one, the smart tier +?? --status # which tiers are up +?? --stop fast # stop just that one +``` + +Don't want to babysit the server at all? `export SHELLLM_AUTOSTART=1` and `,` / `?` bring the default tier up on demand the first time you use them. + ## Sessions Each pane × command gets a sticky JSONL session at `~/.cache/shellllm/sessions/`. Pane identity is `TERM_SESSION_ID` (Terminal.app / iTerm) → `TMUX_PANE` → `WINDOWID` → `$PPID`, first one that resolves. @@ -102,6 +116,35 @@ When the conversation crosses 80% of `SHELLLM_CTX`, older turns are auto-summari Every expired or `--new`'d session flows into `~/.cache/shellllm/archive.db` (sqlite + FTS5) so `???` can search across panes and days. +## Terminal context + +The terminal knows what just happened — shellllm uses it, at a level you control. One env var, a ladder of levels: + +```sh +export SHELLLM_SHELL_CONTEXT=off # capture nothing +export SHELLLM_SHELL_CONTEXT=cmd # previous command + exit status (default) +export SHELLLM_SHELL_CONTEXT=history # + last 10 commands +export SHELLLM_SHELL_CONTEXT=output # + recent pane output (tmux only) +``` + +References resolve themselves: + +```sh +$ git push origin amin +error: src refspec amin does not match any +$ ,, # proposes: git push origin main +$ ? why did that fail # the model sees the command and its exit status +``` + +How it stays private: + +- **Local-first**: with a local model, nothing leaves the machine anyway. The ladder matters when you point `SHELLLM_BASE_URL` at a hosted API. +- **Redaction**: captures pass through a secret scrubber (`KEY=…` assignments, `Bearer` headers, AWS/GitHub/Slack/Stripe-shaped tokens, JWTs) before the model sees them. Git SHAs survive — they're useful. +- **Per-call env, never exported**: the zsh layer passes captures as one-shot environment for the single invocation; nothing lingers in your shell. +- **Both sides enforce the ladder**: zsh won't capture above your level, and the Python side independently re-checks it. +- **Ephemeral**: context blocks are rebuilt per turn and never persisted into sessions or the archive. +- `--no-ctx` skips injection for one call; piped stdin (`cmd | ? …`) is its own explicit consent and works regardless of the ladder. + ## Semantic recall (optional) Recall works in BM25-only mode out of the box. Adding a tiny embedding server upgrades it to **hybrid semantic + BM25** (fused with Reciprocal Rank Fusion): @@ -146,21 +189,32 @@ Every file read through `?` goes through `safe_fs.safe_read`. Four rules, all en Reads cap at 1 MB and use `O_NOFOLLOW` as a belt against a resolve-then-open symlink race. ```sh -pytest -v # 197 tests; 38 dedicated to symlinks, traversal, denylist, lookalikes, truncation +pytest -v # 257 tests; 38 dedicated to symlinks, traversal, denylist, lookalikes, truncation ``` ## Configuration | Variable | Default | Purpose | |---|---|---| -| `SHELLLM_BASE_URL` | `http://127.0.0.1:8080` | llama-server endpoint | -| `SHELLLM_PORT` | `8080` | Server port | +| `SHELLLM_BASE_URL` | `http://127.0.0.1:8080` | llama-server (or hosted) endpoint | +| `SHELLLM_API_KEY` | — | Bearer auth for chat — set when pointing at a hosted API | +| `SHELLLM_MODEL` | `local` | Model id passed in chat requests (override per provider) | +| `SHELLLM_PORT` | `8080` | Server port (default route + `balanced` tier) | +| `SHELLLM_PORT_FAST` | `8091` | `fast` tier port | +| `SHELLLM_PORT_BALANCED` | `$SHELLLM_PORT` | `balanced` tier port | +| `SHELLLM_PORT_SMART` | `8093` | `smart` tier port | +| `SHELLLM_SHELL_CONTEXT` | `cmd` (zsh layer) | `off` / `cmd` / `history` / `output` — terminal-context ladder | +| `SHELLLM_AUTOSTART` | unset | `1` to auto-start the default tier when `,` / `?` find it down | | `SHELLLM_NGL` | `99` | GPU offload layers | | `SHELLLM_CTX` | `32768` | Context window (tokens) | | `SHELLLM_TIMEOUT` | `120` | HTTP timeout (seconds) | | `SHELLLM_LLAMA_MODEL` | — | Explicit GGUF path, overrides tier | | `SHELLLM_LOG` | `~/.cache/shellllm/llama-server.log` | Server log path | -| `SHELLLM_EMBED_URL` | `http://127.0.0.1:8081` | Local embedding server endpoint | +| `SHELLLM_RENDER_URL` | — | Firecrawl-compatible JS-render endpoint (enables) | +| `SHELLLM_RENDER_API_KEY` | — | Bearer auth for the render service | +| `SHELLLM_RENDER_TIMEOUT` | `30` | Render-service HTTP timeout (seconds) | +| `SHELLLM_EMBED_URL` | `http://127.0.0.1:8081` | Local (or hosted) embedding endpoint | +| `SHELLLM_EMBED_API_KEY` | — | Bearer auth for embeddings (falls back to `SHELLLM_API_KEY`) | | `SHELLLM_EMBED_PORT` | `8081` | Embedding server port | | `SHELLLM_EMBED_CTX` | `2048` | Embedding context window | | `SHELLLM_EMBED_MODEL` | `local-embed` | Model name passed to `/v1/embeddings` | @@ -192,10 +246,53 @@ You don't need to ship every "what does git stash do" question to a frontier mod The bet: a local 27B model is roughly equivalent to a frontier model for the questions you ask between `git commit` and `make test`. The wins — privacy, latency, offline availability, $0 per question — are real, every day. +## Use a hosted API instead of llama-server + +shellllm's chat and embedding paths are both OpenAI-compatible. Point them at any provider, BYOK: + +```sh +# OpenAI +export SHELLLM_BASE_URL="https://api.openai.com" +export SHELLLM_API_KEY="sk-..." +export SHELLLM_MODEL="gpt-4o-mini" + +# OpenRouter +export SHELLLM_BASE_URL="https://openrouter.ai/api" +export SHELLLM_API_KEY="sk-or-..." +export SHELLLM_MODEL="anthropic/claude-3.5-sonnet" + +# Groq (very fast) +export SHELLLM_BASE_URL="https://api.groq.com/openai" +export SHELLLM_API_KEY="gsk-..." +export SHELLLM_MODEL="llama-3.3-70b-versatile" +``` + +When `SHELLLM_API_KEY` is set, every chat request carries `Authorization: Bearer …`. Without it, requests stay anonymous (which is what the local `llama-server` wants). The same fall-through applies to embeddings: set `SHELLLM_EMBED_API_KEY` for a separate provider, or let it inherit `SHELLLM_API_KEY` when both endpoints share auth. + +Mix and match: local chat + hosted embeddings, hosted chat + local embeddings, or both hosted. The model decides; shellllm just plumbs. + +## JS rendering for `fetch_url` + +Static HTML works for most pages. SPAs (React/Vue/Svelte sites that paint after the initial response) come back empty. To handle them, point `fetch_url` at a Firecrawl-compatible scraper — hosted or self-hosted: + +```sh +# Hosted: https://firecrawl.dev +export SHELLLM_RENDER_URL="https://api.firecrawl.dev" +export SHELLLM_RENDER_API_KEY="fc-..." + +# Self-hosted: https://github.com/mendableai/firecrawl +export SHELLLM_RENDER_URL="http://localhost:3002" +export SHELLLM_RENDER_API_KEY="any-string-if-disabled" +``` + +When configured, every `fetch_url` call tries `POST {url}/v1/scrape` first (Bearer auth, asks for markdown). On any failure — connection refused, 4xx/5xx, timeout — it falls back transparently to the static fetcher. Without the env vars, behavior is unchanged. + +Any service speaking the same `/v1/scrape` shape works. Other vendors can be bridged with a tiny proxy that translates between their API and Firecrawl's. + ## What's deliberately not built - **GBNF prefix grammar for `,`.** JSON schema is enough for v1; the system prompt forbids the obvious destructive commands. -- **JS rendering for `fetch_url`.** Pages are fetched as static HTML and reduced to text — SPAs that need JS to populate content will look empty. +- **Local browser-based JS rendering.** Skipped on purpose — a 300MB Chromium dependency doesn't match the offline-by-default story. The BYOK `SHELLLM_RENDER_URL` path above is the supported alternative. ## Development @@ -218,9 +315,11 @@ src/shellllm/ ├── embed.py client for a local llama-server in --embedding mode ├── compact.py summary-buffer compaction over the same local model ├── claude_mem.py optional adapter for claude-mem server-beta +├── render.py optional Firecrawl-compatible JS renderer for fetch_url ├── safe_fs.py filesystem hard wall -├── client.py llama-server HTTP client (one-shot + streaming) +├── client.py OpenAI-compatible chat HTTP client (Bearer auth opt-in) ├── context.py date/OS/timezone prelude +├── shell_context.py opt-in terminal context: ladder, redaction, piped stdin └── web.py stdlib DuckDuckGo scraper + fetch_url with SSRF guard ``` diff --git a/demo.tape b/demo.tape new file mode 100644 index 0000000..7b5861c --- /dev/null +++ b/demo.tape @@ -0,0 +1,146 @@ +# shellllm — flow demo +# +# Tells a 60-second story: you're working in a directory, you hit a +# command you don't quite remember, you ask the model about it, you +# refine, you move on, and later you recall a thing you asked earlier +# — all without leaving your zsh prompt. +# +# Render +# brew install vhs # one-time, ~30s +# ?? --start balanced # llama-server must be running +# vhs demo.tape # writes demo.gif +# +# The Hide block at the top sets up a sandbox in /tmp/shellllm-demo with +# a few sized files (so `du` has something to show) and seeds the +# archive with one prior conversation about git stash (so the `???` +# recall later finds a real, pre-existing hit). Everything else is the +# tool acting in its native environment. + +Output demo.gif + +# ── Look ──────────────────────────────────────────────────────────────── +Require zsh +Require llama-server +Set Shell "zsh" +Set FontSize 18 +Set FontFamily "JetBrainsMono Nerd Font" +Set Theme "Catppuccin Mocha" +Set Width 1200 +Set Height 720 +Set Padding 28 +Set BorderRadius 10 +Set WindowBar Colorful +Set TypingSpeed 60ms +Set PlaybackSpeed 1.0 +Set CursorBlink true + +# ── Setup (hidden) ────────────────────────────────────────────────────── +# Fresh sandbox, isolated archive/memory paths so we don't touch the +# user's real cache while recording. +Hide +Type "export SHELLLM_ARCHIVE_DB=/tmp/shellllm-demo/archive.db SHELLLM_MEMORY_FILE=/tmp/shellllm-demo/memory.jsonl SHELLLM_SESSIONS_DIR=/tmp/shellllm-demo/sessions" +Enter +Sleep 100ms +Type "rm -rf /tmp/shellllm-demo && mkdir -p /tmp/shellllm-demo/sessions && cd /tmp/shellllm-demo" +Enter +Sleep 100ms +# Some real-sized files so `du -ah | sort -hr | head -5` has output. +Type "head -c 1500000 release-build.tar.gz" +Enter +Type "head -c 920000 node_modules.cache" +Enter +Type "head -c 410000 coverage-report.html" +Enter +Type "head -c 180000 webpack-stats.json" +Enter +Type "touch -d '2 hours ago' release-build.tar.gz coverage-report.html" +Enter +# Seed one prior shellllm session into the archive so `??? git` finds +# something later in the demo. We use a heredoc to keep the quoting +# manageable inside the Type command. +Type@10ms "python3 - <<'PY'" +Enter +Type@10ms "from shellllm.archive import Archive" +Enter +Type@10ms "Archive().ingest_session(cmd='ask', terminal_id='demo', created_at=0, last_used=0, last_pwd='~/code/api', last_date='2026-06-08', turn_count=1, messages=[{'role':'user','content':'how does git stash work'},{'role':'assistant','content':'git stash temporarily shelves uncommitted changes so you can switch contexts, then `git stash pop` re-applies them on top of the current working tree.'}])" +Enter +Type@10ms "PY" +Enter +Sleep 400ms +Type "clear" +Enter +Sleep 200ms +Show + +# ── Beat 1 — `,` proposes a command ───────────────────────────────────── +# You're in a directory full of build artifacts. You vaguely remember +# `du` exists. Instead of googling, you ask the model in-band. +Type ", find the five largest files in this directory" +Sleep 600ms +Enter +# Wait for the model + fzf to render the candidate picker. +Sleep 6s +# Accept the highlighted proposal — `print -z` drops it on the prompt. +Enter +Sleep 800ms +# Run it. +Enter +Sleep 3s + +# ── Beat 2 — refine via the sticky session ────────────────────────────── +# Same `,` thread in this pane: the model remembers the prior proposal. +# "the same but only ones modified today" works because it has +# context, exactly like a real conversation. +Type ", the same but only ones modified in the last hour" +Sleep 600ms +Enter +Sleep 6s +Enter +Sleep 800ms +Enter +Sleep 3s + +# ── Beat 3 — `?` answers without breaking the flow ────────────────────── +# Two files made the cut. Now you need to commit, but `git stash` came +# to mind and you blanked. Ask the model in-band, get an answer. +Type "? in markdown and 3 lines max, what does git stash do" +Sleep 600ms +Enter +# Streaming markdown render — generous sleep because turn length varies +# with the local model. +Sleep 14s + +# ── Beat 4 — pin a long-term fact ─────────────────────────────────────── +# While you're here, pin something. Now every `?` carries it. +Type "??? --add this project uses pnpm and ships via release-build.tar.gz" +Sleep 500ms +Enter +Sleep 2s + +# ── Beat 5 — start a fresh `?` thread so the prior turn lands in the archive +# `--new` archives the current session — perfect mid-flow segue. +Type "? --new" +Sleep 500ms +Enter +Sleep 2s +Type "clear" +Enter +Sleep 400ms + +# ── Beat 6 — `???` recalls across past sessions ───────────────────────── +# "What was that git thing I asked earlier?" Bare query → FTS5 search +# across every archived conversation in this pane and others. +Type "??? git stash" +Sleep 500ms +Enter +Sleep 5s + +# ── Beat 7 — `??? --archives` shows what's in the vault ───────────────── +Type "??? --archives 5" +Sleep 500ms +Enter +Sleep 5s + +# ── Outro ─────────────────────────────────────────────────────────────── +# Let the final frame breathe so a viewer can read the last row. +Sleep 3s diff --git a/src/shellllm/archive.py b/src/shellllm/archive.py index 884cf55..ac3bade 100644 --- a/src/shellllm/archive.py +++ b/src/shellllm/archive.py @@ -33,7 +33,8 @@ import re import sqlite3 import time -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -396,11 +397,23 @@ def _search_cosine( # ── Internals ------------------------------------------------------ - def _conn(self) -> sqlite3.Connection: + @contextmanager + def _conn(self) -> Iterator[sqlite3.Connection]: + """Yield a connection inside a transaction, then *close* it. + + ``with sqlite3.connect(...)`` alone only scopes the transaction — + the underlying connection (3 fds in WAL mode: db, -wal, -shm) + stays open until GC. Under the test suite that exhausted the + default macOS fd limit, so we close deterministically. + """ conn = sqlite3.connect(self.path) - conn.execute("PRAGMA journal_mode=WAL") - conn.execute("PRAGMA synchronous=NORMAL") - return conn + try: + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + with conn: + yield conn + finally: + conn.close() # ── Helpers -------------------------------------------------------------- diff --git a/src/shellllm/ask.py b/src/shellllm/ask.py index 5bd3456..7678128 100644 --- a/src/shellllm/ask.py +++ b/src/shellllm/ask.py @@ -49,6 +49,7 @@ from .memory import MemoryStore, render_memory_block from .safe_fs import WallViolation, safe_read_text from .session import SessionStore, sweep_expired +from .shell_context import MAX_PIPED_CHARS, build_piped_block, build_shell_context_block from .web import fetch_url_as_text, search_as_text MAX_ITERATIONS = 12 @@ -274,6 +275,8 @@ def run_agent( cmd_label: str = "ask", archive: Archive | None = None, auto_recall: bool = False, + shell_ctx: bool = True, + piped_input: str = "", ) -> int: """Run the tool-calling agent loop. Returns a process-style exit code. @@ -326,6 +329,22 @@ def run_agent( ): system_msgs.append({"role": "system", "content": build_prelude()}) + # Terminal context (opt-in via SHELLLM_SHELL_CONTEXT) is per-turn + # ephemeral: rebuilt fresh on every invocation and never persisted — + # the just-prepended system prefix is stripped before the session is + # written, so "that command" always means the *latest* one. + if include_context and shell_ctx: + ctx_block = build_shell_context_block() + if ctx_block: + system_msgs.append({"role": "system", "content": ctx_block}) + + # Piped stdin (`make 2>&1 | ? what broke`) — explicit consent by + # construction, so no ladder gate. Per-turn ephemeral like the rest + # of the system prefix. + piped_block = build_piped_block(piped_input) + if piped_block: + system_msgs.append({"role": "system", "content": piped_block}) + # claude-mem context injection: only on a brand-new session, so we # don't repeatedly re-paste the same prior observations as the # conversation grows. The model already has them in history after @@ -505,6 +524,19 @@ def _strip_leading_system(messages: list[dict[str, Any]], *, expected: int) -> l ] +def _read_piped_stdin() -> str: + """Return piped stdin content (capped), or "" on a TTY / no data.""" + stdin = sys.stdin + try: + if stdin is None or stdin.isatty(): + return "" + # Read one char past the cap so build_piped_block knows to mark + # the result as truncated. + return stdin.read(MAX_PIPED_CHARS + 1) + except (OSError, ValueError, UnicodeDecodeError): + return "" + + def _print_usage(label: str, *, to: Any = None) -> None: """Write the flag reference to stdout (default) or a given stream.""" out = to or sys.stdout @@ -517,6 +549,9 @@ def _print_usage(label: str, *, to: Any = None) -> None: f" {label} --compact force compaction\n" f" {label} --auto-recall inject archive hits as context\n" f" {label} --no-auto-recall skip recall this turn\n" + f" {label} --no-ctx skip terminal context this turn\n" + f" cmd 2>&1 | {label} piped input becomes context\n" + f" {label} --fast|--balanced|--smart route this call to a tier (zsh)\n" f" {label} --mem | --no-mem force claude-mem on/off for this call\n" f" {label} --help show this message\n" f"\n" @@ -589,6 +624,8 @@ def _consume_flag(flag: str) -> bool: # --auto-recall (or SHELLLM_AUTO_RECALL=1) injects archive hits as # context on the first turn of a new session. --no-auto-recall # overrides the env var for this call. + shell_ctx = not _consume_flag("--no-ctx") + no_recall = _consume_flag("--no-auto-recall") auto_recall_flag = _consume_flag("--auto-recall") env_recall = os.environ.get("SHELLLM_AUTO_RECALL", "").strip().lower() in ( @@ -646,6 +683,8 @@ def _consume_flag(flag: str) -> bool: cmd_label=cmd, archive=archive, auto_recall=auto_recall, + shell_ctx=shell_ctx, + piped_input=_read_piped_stdin(), ) except LlamaServerError as exc: sys.stderr.write(f"{_RED}{err_label} error:{_RESET} {exc}\n") diff --git a/src/shellllm/client.py b/src/shellllm/client.py index 6c0ee6a..a5cdd5a 100644 --- a/src/shellllm/client.py +++ b/src/shellllm/client.py @@ -18,6 +18,26 @@ DEFAULT_MODEL = os.environ.get("SHELLLM_MODEL", "local") DEFAULT_TIMEOUT = float(os.environ.get("SHELLLM_TIMEOUT", "120")) +# Env name read lazily so tests can monkeypatch + the same key powers +# both chat and (optionally) embeddings without an import-time race. +API_KEY_ENV = "SHELLLM_API_KEY" + + +def _auth_headers() -> dict[str, str]: + """Build request headers with optional Bearer auth. + + The local ``llama-server`` doesn't need auth, so the bearer token + is opt-in via ``SHELLLM_API_KEY``. Setting it lets you point + ``SHELLLM_BASE_URL`` at any OpenAI-compatible endpoint (OpenAI, + OpenRouter, Groq, Together, Mistral, …) and have the chat path + "just work". + """ + headers = {"Content-Type": "application/json"} + key = os.environ.get(API_KEY_ENV, "").strip() + if key: + headers["Authorization"] = f"Bearer {key}" + return headers + class LlamaServerError(RuntimeError): pass @@ -60,11 +80,14 @@ def chat( r = httpx.post( f"{base_url}/v1/chat/completions", json=payload, + headers=_auth_headers(), timeout=DEFAULT_TIMEOUT, ) except httpx.ConnectError as exc: raise LlamaServerError( - f"can't reach llama-server at {base_url}\n what to do: run `??` to start it" + f"can't reach llama-server at {base_url}\n" + " what to do: run `??` to start it " + "(or `export SHELLLM_AUTOSTART=1` to start on demand)" ) from exc except httpx.ReadTimeout as exc: raise LlamaServerError(f"llama-server timed out after {DEFAULT_TIMEOUT}s") from exc @@ -122,6 +145,7 @@ def chat_stream( "POST", f"{base_url}/v1/chat/completions", json=payload, + headers=_auth_headers(), timeout=timeout, ) as r: if r.status_code != 200: @@ -160,7 +184,9 @@ def chat_stream( finish_reason = choice["finish_reason"] except httpx.ConnectError as exc: raise LlamaServerError( - f"can't reach llama-server at {base_url}\n what to do: run `??` to start it" + f"can't reach llama-server at {base_url}\n" + " what to do: run `??` to start it " + "(or `export SHELLLM_AUTOSTART=1` to start on demand)" ) from exc except httpx.ReadTimeout as exc: raise LlamaServerError("llama-server stream timed out") from exc diff --git a/src/shellllm/comma.py b/src/shellllm/comma.py index 9d35408..2c1b8b8 100644 --- a/src/shellllm/comma.py +++ b/src/shellllm/comma.py @@ -31,6 +31,7 @@ from .client import LlamaServerError, chat from .embed import embed as embed_text from .session import SessionStore, sweep_expired +from .shell_context import build_shell_context_block SCHEMA = { "type": "object", @@ -116,9 +117,12 @@ def _print_usage(*, to: Any = None) -> None: out = to or sys.stdout out.write( "usage: , \n" + " , --fix [hint] fix the previous command (also: `,,`)\n" " , --new start a fresh session\n" " , --reset drop current session\n" " , --history print session transcript\n" + " , --no-ctx skip terminal context this turn\n" + " , --fast|--balanced|--smart … route this call to a tier (zsh)\n" " , --help show this message\n" "\n" "For facts and cross-session recall, see `?: help`.\n" @@ -229,6 +233,7 @@ def _build_messages( prompt: str, first_turn: bool, resumed: bool, + shell_ctx: bool = True, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: """Return (messages_to_send, history_to_persist_after). @@ -245,6 +250,14 @@ def _build_messages( if first_turn or resumed or session.meta.last_pwd != pwd or session.meta.last_date != date: system_msgs.append({"role": "system", "content": _context_block()}) + # Terminal context is per-turn ephemeral: rebuilt every call, never + # persisted (system messages are stripped before the session is + # written), and empty unless SHELLLM_SHELL_CONTEXT is set. + if shell_ctx: + ctx_block = build_shell_context_block() + if ctx_block: + system_msgs.append({"role": "system", "content": ctx_block}) + history = list(session.messages) user_msg: dict[str, Any] = {"role": "user", "content": prompt} return system_msgs + history + [user_msg], history + [user_msg] @@ -309,7 +322,29 @@ def _consume_flag(flag: str) -> bool: if _consume_flag("--new"): session.archive_and_reset(archive=archive, embed_fn=_safe_embed) + shell_ctx = not _consume_flag("--no-ctx") + fix_mode = _consume_flag("--fix") + prompt = " ".join(argv).strip() + + if fix_mode: + if not shell_ctx: + _err.print(f"{_RED}, error:{_RESET} --fix needs terminal context; drop --no-ctx") + return 2 + if not build_shell_context_block(): + _err.print(f"{_RED}, error:{_RESET} --fix needs terminal context, and none arrived.") + _err.print( + " if you set SHELLLM_SHELL_CONTEXT=off, re-enable it: export SHELLLM_SHELL_CONTEXT=cmd" + ) + _err.print(" otherwise re-source zsh/shellllm.zsh (older versions didn't capture).") + return 2 + repair = ( + "Diagnose the previous command using the terminal context " + "(command, exit status, output if present) and propose corrected " + "commands that do what the user wanted." + ) + prompt = f"{repair} Hint from the user: {prompt}" if prompt else repair + if not prompt: _print_usage(to=sys.stderr) return 2 @@ -323,7 +358,11 @@ def _consume_flag(flag: str) -> bool: _note(f"refining — turn {session.meta.turn_count + 1}") messages, new_history_with_user = _build_messages( - session=session, prompt=prompt, first_turn=first_turn, resumed=resumed + session=session, + prompt=prompt, + first_turn=first_turn, + resumed=resumed, + shell_ctx=shell_ctx, ) result = _ask_model(messages) diff --git a/src/shellllm/embed.py b/src/shellllm/embed.py index 260e8ff..03a5df6 100644 --- a/src/shellllm/embed.py +++ b/src/shellllm/embed.py @@ -26,6 +26,8 @@ DEFAULT_EMBED_URL = os.environ.get("SHELLLM_EMBED_URL", "http://127.0.0.1:8081") DEFAULT_EMBED_MODEL = os.environ.get("SHELLLM_EMBED_MODEL", "local-embed") DEFAULT_TIMEOUT = float(os.environ.get("SHELLLM_EMBED_TIMEOUT", "8")) +EMBED_API_KEY_ENV = "SHELLLM_EMBED_API_KEY" +FALLBACK_API_KEY_ENV = "SHELLLM_API_KEY" # Lazy `os.environ.get` re-read so tests can monkeypatch the var. @@ -33,6 +35,22 @@ def _base_url() -> str: return os.environ.get("SHELLLM_EMBED_URL", DEFAULT_EMBED_URL).rstrip("/") +def _auth_headers() -> dict[str, str]: + """Bearer auth for hosted embedding endpoints. + + ``SHELLLM_EMBED_API_KEY`` takes precedence; if unset we fall back to + ``SHELLLM_API_KEY`` so a single key powers chat + embeddings against + the same provider. The local llama-server doesn't need either. + """ + headers = {"Content-Type": "application/json"} + key = os.environ.get(EMBED_API_KEY_ENV, "").strip() + if not key: + key = os.environ.get(FALLBACK_API_KEY_ENV, "").strip() + if key: + headers["Authorization"] = f"Bearer {key}" + return headers + + def normalize(vec: list[float]) -> list[float]: """L2-normalize so retrieval reduces to a dot product at query time.""" n = math.sqrt(sum(v * v for v in vec)) @@ -85,7 +103,7 @@ def embed( url = f"{(base_url or _base_url()).rstrip('/')}/v1/embeddings" payload = {"model": model, "input": body} try: - r = httpx.post(url, json=payload, timeout=timeout) + r = httpx.post(url, json=payload, headers=_auth_headers(), timeout=timeout) r.raise_for_status() data = r.json() except Exception: # noqa: BLE001 — embedding is best-effort diff --git a/src/shellllm/render.py b/src/shellllm/render.py new file mode 100644 index 0000000..37441d9 --- /dev/null +++ b/src/shellllm/render.py @@ -0,0 +1,112 @@ +"""Optional JS-rendering for ``fetch_url`` via a hosted/self-hosted service. + +The default ``web.fetch_url`` flow fetches static HTML and reduces it +to text — fast, offline, and zero deps, but blind to anything an SPA +populates after first paint. This module adds an opt-in escape hatch: +if the user configures a Firecrawl-compatible endpoint, every +``fetch_url`` call tries the rendered path first and falls back to the +static fetcher on any failure. + +The contract is Firecrawl's `/v1/scrape` shape because it's open, +documented, self-hostable, and already supported by a handful of +adjacent tools: + +* hosted: https://api.firecrawl.dev +* self-hosted: https://github.com/mendableai/firecrawl + +Anything else can be bridged with a tiny proxy that translates between +the upstream API and Firecrawl's shape. + +Bring your own key +~~~~~~~~~~~~~~~~~~ + +:: + + export SHELLLM_RENDER_URL="https://api.firecrawl.dev" + export SHELLLM_RENDER_API_KEY="fc-..." + +Without those two env vars the integration is inert and ``fetch_url`` +behaves exactly as before. +""" + +from __future__ import annotations + +import os + +import httpx + +RENDER_URL_ENV = "SHELLLM_RENDER_URL" +RENDER_KEY_ENV = "SHELLLM_RENDER_API_KEY" +RENDER_TIMEOUT_ENV = "SHELLLM_RENDER_TIMEOUT" +DEFAULT_TIMEOUT = 30.0 + + +def _config() -> tuple[str, str, float]: + """Read env at call time so tests can monkeypatch cleanly.""" + base = os.environ.get(RENDER_URL_ENV, "").rstrip("/") + key = os.environ.get(RENDER_KEY_ENV, "") + timeout_raw = os.environ.get(RENDER_TIMEOUT_ENV, "") + try: + timeout = float(timeout_raw) if timeout_raw else DEFAULT_TIMEOUT + except ValueError: + timeout = DEFAULT_TIMEOUT + return base, key, timeout + + +def is_configured() -> bool: + """True iff both URL and API key are set.""" + base, key, _ = _config() + return bool(base and key) + + +def render_url(url: str) -> str | None: + """POST the URL to the configured renderer, return rendered text or None. + + Returns ``None`` when the integration isn't configured, when the + URL is empty, or when any error short-circuits the call. We never + raise — ``fetch_url`` keeps working through its static-HTML + fallback no matter what happens here. + """ + base, key, timeout = _config() + if not base or not key: + return None + target = url.strip() + if not target: + return None + + try: + r = httpx.post( + f"{base}/v1/scrape", + json={"url": target, "formats": ["markdown"]}, + headers={ + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + }, + timeout=timeout, + ) + r.raise_for_status() + data = r.json() + except Exception: # noqa: BLE001 — rendering is best-effort + return None + + return _extract_text(data) + + +def _extract_text(data: object) -> str | None: + """Pull rendered text out of a Firecrawl-shaped response. + + The wire format has shifted across releases. We check the common + fields, prefer markdown, and tolerate older "html only" payloads. + """ + if not isinstance(data, dict): + return None + + payload = data.get("data") if isinstance(data.get("data"), dict) else data + if not isinstance(payload, dict): + return None + + for key in ("markdown", "content", "text", "html"): + value = payload.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return None diff --git a/src/shellllm/shell_context.py b/src/shellllm/shell_context.py new file mode 100644 index 0000000..9e0a63c --- /dev/null +++ b/src/shellllm/shell_context.py @@ -0,0 +1,152 @@ +"""Opt-in terminal context: what just happened in this pane. + +The zsh wrappers capture the previous command, its exit status, recent +history, and (inside tmux) recent pane output, and pass them down via +environment variables. This module turns those into a small, redacted +system block — so `, why did that fail` and `? what does that error +mean` work without re-typing anything. + +Privacy ladder — ``SHELLLM_SHELL_CONTEXT``: + + off nothing is captured or injected + cmd previous command + exit status + history + last few commands + output + recent pane output (tmux only) + +This module treats an *absent* variable as ``off``; the stock zsh layer +defaults the ladder to ``cmd`` and passes the level down per +invocation, so out of the box `,,` and "why did that fail" just work. +``SHELLLM_SHELL_CONTEXT=off`` disables capture entirely. + +Both the zsh side (capture) and this module (injection) enforce the +ladder independently, so a stale exported variable can never leak past +the configured level. Unknown levels fail safe to ``off``. All values +pass through :func:`redact` before injection; with a local model the +data never leaves the machine anyway, but ``SHELLLM_BASE_URL`` may +point at a hosted API. +""" + +from __future__ import annotations + +import os +import re +from collections.abc import Mapping + +LEVEL_ENV = "SHELLLM_SHELL_CONTEXT" +LEVELS = ("off", "cmd", "history", "output") + +MAX_LAST_CMD_CHARS = 500 +MAX_HISTORY_LINES = 10 +MAX_HISTORY_LINE_CHARS = 200 +MAX_OUTPUT_CHARS = 4_000 +MAX_PIPED_CHARS = 16_000 + +_REDACTED = "[redacted]" + +# Keyed assignments/headers: KEY=value, key: value. The value is dropped, +# the key kept so the model still sees *what* was being set. +_KEYED = re.compile( + r"(?i)\b([A-Za-z0-9_-]*(?:api[_-]?key|access[_-]?key|secret|token|passw(?:or)?d|credential)" + r"[A-Za-z0-9_-]*)(\s*[=:]\s*)(\"[^\"]*\"|'[^']*'|\S+)" +) + +# Well-known token shapes. Deliberately NOT a generic long-blob pattern: +# 40-hex git SHAs are useful context and must survive redaction. +_TOKEN_SHAPES = [ + re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]{8,}"), + re.compile(r"\bAKIA[0-9A-Z]{16}\b"), + re.compile(r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]+"), + re.compile(r"\bgh[pousr]_[A-Za-z0-9]{20,}\b"), + re.compile(r"\bsk[-_](?:live[-_]|test[-_])?[A-Za-z0-9_-]{10,}\b"), + re.compile(r"\bxox[abprs]-[A-Za-z0-9-]{10,}\b"), +] + + +def redact(text: str) -> str: + """Strip likely secrets from a capture before it reaches the model.""" + out = _KEYED.sub(lambda m: f"{m.group(1)}{m.group(2)}{_REDACTED}", text) + for pattern in _TOKEN_SHAPES: + out = pattern.sub(_REDACTED, out) + return out + + +def _level(env: Mapping[str, str]) -> int: + """Return the ladder index for the configured level; 0 == off.""" + raw = env.get(LEVEL_ENV, "off").strip().lower() + try: + return LEVELS.index(raw) + except ValueError: + return 0 + + +def build_shell_context_block(env: Mapping[str, str] | None = None) -> str: + """Render the terminal-context system block, or "" when off / empty. + + ``env`` is injectable for tests; defaults to ``os.environ``. + """ + env = os.environ if env is None else env + level = _level(env) + if level == 0: + return "" + + parts: list[str] = [] + + last_cmd = env.get("SHELLLM_LAST_CMD", "").strip() + if last_cmd: + parts.append(f"- previous command: `{redact(last_cmd[:MAX_LAST_CMD_CHARS])}`") + + raw_status = env.get("SHELLLM_LAST_STATUS", "").strip() + if raw_status: + try: + status = int(raw_status) + except ValueError: + status = None + if status is not None: + suffix = "" if status == 0 else " (it failed)" + parts.append(f"- its exit status: {status}{suffix}") + + if level >= LEVELS.index("history"): + history = env.get("SHELLLM_RECENT_HISTORY", "").strip() + if history: + lines = [ + ln.strip()[:MAX_HISTORY_LINE_CHARS] for ln in history.splitlines() if ln.strip() + ] + lines = lines[-MAX_HISTORY_LINES:] + if lines: + joined = "\n".join(f" {redact(ln)}" for ln in lines) + parts.append(f"- recent commands (oldest first):\n{joined}") + + if level >= LEVELS.index("output"): + output = env.get("SHELLLM_PANE_OUTPUT", "").strip() + if output: + tail = output[-MAX_OUTPUT_CHARS:] + parts.append(f"- recent terminal output (most recent last):\n{redact(tail)}") + + if not parts: + return "" + + return "\n".join( + [ + "Terminal context (shared per the user's configured level; secrets are redacted):", + *parts, + "Use this to resolve references like 'that command', 'the error', or 'why did it fail'.", + ] + ) + + +def build_piped_block(text: str) -> str: + """Render piped stdin as a system block, or "" when empty. + + Piping is explicit consent — no ladder gate — but the content still + goes through :func:`redact`. Errors usually sit at the end of a + capture, so oversized input keeps the tail. + """ + text = text.strip() + if not text: + return "" + truncated = len(text) > MAX_PIPED_CHARS + if truncated: + text = text[-MAX_PIPED_CHARS:] + header = "Piped input (the user piped this into the command" + header += f"; truncated to the last {MAX_PIPED_CHARS} characters):" if truncated else "):" + return f"{header}\n{redact(text)}" diff --git a/src/shellllm/web.py b/src/shellllm/web.py index 880495d..327b5d3 100644 --- a/src/shellllm/web.py +++ b/src/shellllm/web.py @@ -293,7 +293,22 @@ def fetch_url( def fetch_url_as_text(url: str, *, max_chars: int = FETCH_MAX_CHARS) -> str: - """LLM-tool-friendly wrapper: never raises, always returns text.""" + """LLM-tool-friendly wrapper: never raises, always returns text. + + When ``SHELLLM_RENDER_URL`` + ``SHELLLM_RENDER_API_KEY`` are set, we + try the JS-rendering service first and fall back to the static + HTML fetcher on any failure. This lets the agent read SPA-rendered + pages (React/Vue/Svelte sites that paint after the initial HTML) + without forcing a heavy local browser install. + """ + from . import render + + rendered = render.render_url(url) + if rendered: + if len(rendered) > max_chars: + rendered = rendered[:max_chars] + f"\n\n(truncated at {max_chars} chars)" + return rendered + try: body = fetch_url(url) except FetchError as exc: diff --git a/tests/test_api_key.py b/tests/test_api_key.py new file mode 100644 index 0000000..5726900 --- /dev/null +++ b/tests/test_api_key.py @@ -0,0 +1,179 @@ +"""Verify ``SHELLLM_API_KEY`` / ``SHELLLM_EMBED_API_KEY`` add Bearer auth. + +Covers the BYOK path that lets shellllm point at hosted OpenAI-compatible +endpoints (OpenAI, OpenRouter, Groq, Together, …) instead of a local +llama-server. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from shellllm import client, embed + + +@dataclass +class FakeResponse: + status_code: int = 200 + body: dict[str, Any] | None = None + text: str = "" + + def json(self) -> Any: + return self.body if self.body is not None else {} + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + +# ── client._auth_headers ---------------------------------------------------- + + +def test_auth_headers_no_key(monkeypatch): + monkeypatch.delenv("SHELLLM_API_KEY", raising=False) + headers = client._auth_headers() + assert headers["Content-Type"] == "application/json" + assert "Authorization" not in headers + + +def test_auth_headers_with_key(monkeypatch): + monkeypatch.setenv("SHELLLM_API_KEY", "sk-test-123") + headers = client._auth_headers() + assert headers["Authorization"] == "Bearer sk-test-123" + + +def test_auth_headers_strips_whitespace(monkeypatch): + monkeypatch.setenv("SHELLLM_API_KEY", " sk-test-123\n") + headers = client._auth_headers() + assert headers["Authorization"] == "Bearer sk-test-123" + + +def test_auth_headers_empty_key_omits_auth(monkeypatch): + monkeypatch.setenv("SHELLLM_API_KEY", " ") + headers = client._auth_headers() + assert "Authorization" not in headers + + +# ── client.chat sends the header ------------------------------------------ + + +def test_chat_sends_bearer_when_key_set(monkeypatch): + monkeypatch.setenv("SHELLLM_API_KEY", "sk-test") + captured: list[dict] = [] + + def fake_post(url, *, json, headers, timeout): # noqa: A002 + captured.append({"url": url, "headers": headers, "json": json}) + return FakeResponse( + status_code=200, + body={"choices": [{"message": {"role": "assistant", "content": "hi"}}]}, + ) + + monkeypatch.setattr(client.httpx, "post", fake_post) + msg = client.chat( + [{"role": "user", "content": "hello"}], + base_url="https://api.example.com", + ) + assert msg["content"] == "hi" + assert captured[0]["url"] == "https://api.example.com/v1/chat/completions" + assert captured[0]["headers"]["Authorization"] == "Bearer sk-test" + + +def test_chat_no_auth_header_without_key(monkeypatch): + monkeypatch.delenv("SHELLLM_API_KEY", raising=False) + captured: list[dict] = [] + + def fake_post(url, *, json, headers, timeout): # noqa: A002 + captured.append({"headers": headers}) + return FakeResponse( + status_code=200, + body={"choices": [{"message": {"role": "assistant", "content": "hi"}}]}, + ) + + monkeypatch.setattr(client.httpx, "post", fake_post) + client.chat([{"role": "user", "content": "hello"}]) + assert "Authorization" not in captured[0]["headers"] + + +# ── embed._auth_headers -------------------------------------------------- + + +def test_embed_auth_prefers_embed_specific_key(monkeypatch): + monkeypatch.setenv("SHELLLM_API_KEY", "fallback-key") + monkeypatch.setenv("SHELLLM_EMBED_API_KEY", "embed-key") + headers = embed._auth_headers() + assert headers["Authorization"] == "Bearer embed-key" + + +def test_embed_auth_falls_back_to_shared_key(monkeypatch): + monkeypatch.delenv("SHELLLM_EMBED_API_KEY", raising=False) + monkeypatch.setenv("SHELLLM_API_KEY", "shared-key") + headers = embed._auth_headers() + assert headers["Authorization"] == "Bearer shared-key" + + +def test_embed_auth_no_key(monkeypatch): + monkeypatch.delenv("SHELLLM_API_KEY", raising=False) + monkeypatch.delenv("SHELLLM_EMBED_API_KEY", raising=False) + headers = embed._auth_headers() + assert "Authorization" not in headers + + +def test_embed_sends_bearer_when_key_set(monkeypatch): + monkeypatch.setenv("SHELLLM_API_KEY", "shared-key") + monkeypatch.setenv("SHELLLM_EMBED_URL", "https://api.example.com") + captured: list[dict] = [] + + def fake_post(url, *, json, headers=None, timeout=None): # noqa: A002 + captured.append({"url": url, "headers": headers or {}}) + return FakeResponse(body={"data": [{"embedding": [1.0, 0.0]}]}) + + monkeypatch.setattr(embed.httpx, "post", fake_post) + vec = embed.embed("hello", normalize_output=False) + assert vec == [1.0, 0.0] + assert captured[0]["headers"]["Authorization"] == "Bearer shared-key" + + +# ── web.fetch_url_as_text picks render first when configured ------------- + + +def test_fetch_url_uses_render_when_configured(monkeypatch): + from shellllm import web + + monkeypatch.setenv("SHELLLM_RENDER_URL", "https://render.test") + monkeypatch.setenv("SHELLLM_RENDER_API_KEY", "fc-key") + + rendered_called = {"hit": False} + static_called = {"hit": False} + + def fake_render(url): + rendered_called["hit"] = True + return "RENDERED markdown body" + + def fake_static(url): + static_called["hit"] = True + return "STATIC fallback" + + monkeypatch.setattr("shellllm.render.render_url", fake_render) + monkeypatch.setattr(web, "fetch_url", fake_static) + + out = web.fetch_url_as_text("https://example.com/spa") + assert "RENDERED" in out + assert rendered_called["hit"] + assert not static_called["hit"] + + +def test_fetch_url_falls_back_to_static_when_render_returns_none(monkeypatch): + from shellllm import web + + def fake_render(url): + return None + + def fake_static(url): + return "STATIC body" + + monkeypatch.setattr("shellllm.render.render_url", fake_render) + monkeypatch.setattr(web, "fetch_url", fake_static) + + out = web.fetch_url_as_text("https://example.com/static") + assert "STATIC" in out diff --git a/tests/test_ask_context.py b/tests/test_ask_context.py new file mode 100644 index 0000000..fd4f9bd --- /dev/null +++ b/tests/test_ask_context.py @@ -0,0 +1,68 @@ +"""`?` terminal-context + piped-stdin plumbing (run_cli → run_agent).""" + +from __future__ import annotations + +import io +import sys + +import pytest + +from shellllm import ask + + +@pytest.fixture +def isolated(tmp_path, monkeypatch): + monkeypatch.setenv("SHELLLM_SESSIONS_DIR", str(tmp_path / "sessions")) + monkeypatch.setenv("SHELLLM_ARCHIVE_DB", str(tmp_path / "archive.db")) + monkeypatch.setenv("SHELLLM_MEMORY_FILE", str(tmp_path / "memory.jsonl")) + monkeypatch.setenv("TERM_SESSION_ID", "test-pane-ctx") + monkeypatch.delenv("SHELLLM_SHELL_CONTEXT", raising=False) + return tmp_path + + +@pytest.fixture +def captured_agent(monkeypatch): + captured: dict = {} + + def fake_run_agent(prompt, **kwargs): + captured["prompt"] = prompt + captured.update(kwargs) + return 0 + + monkeypatch.setattr(ask, "run_agent", fake_run_agent) + return captured + + +class _TtyStdin(io.StringIO): + def isatty(self): + return True + + +def _run(argv, monkeypatch): + monkeypatch.setattr(sys, "argv", ["shellllm-ask", *argv]) + return ask.main() + + +def test_tty_stdin_means_no_piped_input(monkeypatch, isolated, captured_agent): + monkeypatch.setattr(sys, "stdin", _TtyStdin("")) + assert _run(["what", "is", "foo"], monkeypatch) == 0 + assert captured_agent["piped_input"] == "" + + +def test_piped_stdin_is_forwarded(monkeypatch, isolated, captured_agent): + monkeypatch.setattr(sys, "stdin", io.StringIO("error: segfault at line 3")) + assert _run(["what", "broke"], monkeypatch) == 0 + assert "segfault at line 3" in captured_agent["piped_input"] + + +def test_no_ctx_flag_forwarded(monkeypatch, isolated, captured_agent): + monkeypatch.setattr(sys, "stdin", _TtyStdin("")) + assert _run(["--no-ctx", "what", "is", "foo"], monkeypatch) == 0 + assert captured_agent["shell_ctx"] is False + assert captured_agent["prompt"] == "what is foo" + + +def test_shell_ctx_defaults_on(monkeypatch, isolated, captured_agent): + monkeypatch.setattr(sys, "stdin", _TtyStdin("")) + assert _run(["what", "is", "foo"], monkeypatch) == 0 + assert captured_agent["shell_ctx"] is True diff --git a/tests/test_comma_session.py b/tests/test_comma_session.py index 9470eed..5f1321b 100644 --- a/tests/test_comma_session.py +++ b/tests/test_comma_session.py @@ -182,6 +182,74 @@ def test_resume_hint_writes_raw_ansi_not_rich_markup( assert "↻ refining" in err +def _enable_shell_ctx(monkeypatch, *, status="1"): + monkeypatch.setenv("SHELLLM_SHELL_CONTEXT", "cmd") + monkeypatch.setenv("SHELLLM_LAST_CMD", "git push origin main") + monkeypatch.setenv("SHELLLM_LAST_STATUS", status) + + +def test_shell_context_injected_when_enabled(monkeypatch, capsys, isolated, fake_model, auto_pick): + _enable_shell_ctx(monkeypatch) + _run(["retry", "that"], monkeypatch) + sent = fake_model[0] + system_text = "\n".join(m["content"] for m in sent if m["role"] == "system") + assert "git push origin main" in system_text + + +def test_shell_context_absent_by_default(monkeypatch, capsys, isolated, fake_model, auto_pick): + monkeypatch.delenv("SHELLLM_SHELL_CONTEXT", raising=False) + monkeypatch.setenv("SHELLLM_LAST_CMD", "git push origin main") + _run(["list", "files"], monkeypatch) + sent = fake_model[0] + system_text = "\n".join(m["content"] for m in sent if m["role"] == "system") + assert "git push origin main" not in system_text + + +def test_no_ctx_flag_skips_injection(monkeypatch, capsys, isolated, fake_model, auto_pick): + _enable_shell_ctx(monkeypatch) + _run(["--no-ctx", "list", "files"], monkeypatch) + sent = fake_model[0] + system_text = "\n".join(m["content"] for m in sent if m["role"] == "system") + assert "git push origin main" not in system_text + + +def test_shell_context_never_persisted(monkeypatch, capsys, isolated, fake_model, auto_pick): + _enable_shell_ctx(monkeypatch) + _run(["retry", "that"], monkeypatch) + store, _ = SessionStore.open(cmd="comma") + assert all(m.get("role") != "system" for m in store.messages) + + +def test_fix_without_context_errors_with_hint(monkeypatch, capsys, isolated): + monkeypatch.delenv("SHELLLM_SHELL_CONTEXT", raising=False) + assert _run(["--fix"], monkeypatch) == 2 + err = capsys.readouterr().err + assert "SHELLLM_SHELL_CONTEXT" in err + + +def test_fix_with_no_ctx_is_rejected(monkeypatch, capsys, isolated): + _enable_shell_ctx(monkeypatch) + assert _run(["--fix", "--no-ctx"], monkeypatch) == 2 + assert "--no-ctx" in capsys.readouterr().err + + +def test_fix_builds_repair_prompt(monkeypatch, capsys, isolated, fake_model, auto_pick): + _enable_shell_ctx(monkeypatch) + assert _run(["--fix"], monkeypatch) == 0 + sent = fake_model[0] + user_msg = next(m["content"] for m in sent if m["role"] == "user") + assert "previous command" in user_msg + system_text = "\n".join(m["content"] for m in sent if m["role"] == "system") + assert "git push origin main" in system_text + + +def test_fix_appends_user_hint(monkeypatch, capsys, isolated, fake_model, auto_pick): + _enable_shell_ctx(monkeypatch) + assert _run(["--fix", "I", "meant", "the", "dev", "branch"], monkeypatch) == 0 + user_msg = next(m["content"] for m in fake_model[0] if m["role"] == "user") + assert "I meant the dev branch" in user_msg + + def test_redirect_for_ask_remember(monkeypatch, capsys, isolated): """`,` doesn't share `?`'s deprecation hints — it has its own surface.""" diff --git a/tests/test_embed.py b/tests/test_embed.py index f90c446..a62b43a 100644 --- a/tests/test_embed.py +++ b/tests/test_embed.py @@ -71,8 +71,8 @@ def test_cosine_zero_vector_returns_zero(): def test_embed_posts_to_server_and_normalizes(monkeypatch): captured: list[dict[str, Any]] = [] - def fake_post(url, *, json, timeout): # noqa: A002 - captured.append({"url": url, "json": json, "timeout": timeout}) + def fake_post(url, *, json, headers=None, timeout): # noqa: A002 + captured.append({"url": url, "json": json, "headers": headers or {}, "timeout": timeout}) return FakeResponse(body={"data": [{"embedding": [3.0, 4.0]}]}) monkeypatch.setenv("SHELLLM_EMBED_URL", "http://embed.test") diff --git a/tests/test_render.py b/tests/test_render.py new file mode 100644 index 0000000..bf4c604 --- /dev/null +++ b/tests/test_render.py @@ -0,0 +1,176 @@ +"""Tests for the optional JS-rendering adapter (``render.render_url``). + +httpx is mocked at the module boundary; no real network is touched. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pytest + +from shellllm import render + + +@dataclass +class FakeResponse: + status_code: int = 200 + body: Any | None = None + + def json(self) -> Any: + return self.body if self.body is not None else {} + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + +@pytest.fixture +def configured(monkeypatch): + monkeypatch.setenv(render.RENDER_URL_ENV, "https://render.test") + monkeypatch.setenv(render.RENDER_KEY_ENV, "test-key") + + +@pytest.fixture +def unset_env(monkeypatch): + monkeypatch.delenv(render.RENDER_URL_ENV, raising=False) + monkeypatch.delenv(render.RENDER_KEY_ENV, raising=False) + monkeypatch.delenv(render.RENDER_TIMEOUT_ENV, raising=False) + + +# ── Configuration ----------------------------------------------------------- + + +def test_is_configured_requires_both_vars(monkeypatch, unset_env): + assert render.is_configured() is False + monkeypatch.setenv(render.RENDER_URL_ENV, "https://x") + assert render.is_configured() is False + monkeypatch.setenv(render.RENDER_KEY_ENV, "k") + assert render.is_configured() is True + + +def test_render_url_noop_when_not_configured(monkeypatch, unset_env): + captured: list = [] + monkeypatch.setattr(render.httpx, "post", lambda *a, **kw: captured.append(a)) + assert render.render_url("https://example.com") is None + assert captured == [] + + +def test_render_url_empty_input_returns_none(configured, monkeypatch): + monkeypatch.setattr( + render.httpx, + "post", + lambda *a, **kw: pytest.fail("should not POST for empty url"), + ) + assert render.render_url(" ") is None + + +# ── Happy path -------------------------------------------------------------- + + +def test_render_url_posts_firecrawl_payload(configured, monkeypatch): + captured: list[dict] = [] + + def fake_post(url, *, json, headers, timeout): # noqa: A002 + captured.append({"url": url, "json": json, "headers": headers, "timeout": timeout}) + return FakeResponse(body={"data": {"markdown": "# rendered"}}) + + monkeypatch.setattr(render.httpx, "post", fake_post) + + out = render.render_url("https://example.com/spa") + assert out == "# rendered" + call = captured[0] + assert call["url"] == "https://render.test/v1/scrape" + assert call["json"] == {"url": "https://example.com/spa", "formats": ["markdown"]} + assert call["headers"]["Authorization"] == "Bearer test-key" + assert call["headers"]["Content-Type"] == "application/json" + + +def test_render_url_falls_through_response_fields(configured, monkeypatch): + """We tolerate older Firecrawl-shape responses that only expose html.""" + + monkeypatch.setattr( + render.httpx, + "post", + lambda *a, **kw: FakeResponse(body={"data": {"html": "

hi

"}}), + ) + assert render.render_url("https://example.com") == "

hi

" + + +def test_render_url_supports_top_level_markdown(configured, monkeypatch): + """Some implementations skip the `data` wrapper.""" + + monkeypatch.setattr( + render.httpx, + "post", + lambda *a, **kw: FakeResponse(body={"markdown": "top-level"}), + ) + assert render.render_url("https://example.com") == "top-level" + + +# ── Failure modes ----------------------------------------------------------- + + +def test_render_url_returns_none_on_http_error(configured, monkeypatch): + monkeypatch.setattr( + render.httpx, + "post", + lambda *a, **kw: FakeResponse(status_code=502), + ) + assert render.render_url("https://example.com") is None + + +def test_render_url_returns_none_on_connection_error(configured, monkeypatch): + def broken(*a, **kw): + raise RuntimeError("connect refused") + + monkeypatch.setattr(render.httpx, "post", broken) + assert render.render_url("https://example.com") is None + + +def test_render_url_returns_none_on_malformed_json(configured, monkeypatch): + monkeypatch.setattr( + render.httpx, + "post", + lambda *a, **kw: FakeResponse(body={"unrelated": "shape"}), + ) + assert render.render_url("https://example.com") is None + + +def test_render_url_returns_none_on_non_dict_payload(configured, monkeypatch): + monkeypatch.setattr( + render.httpx, + "post", + lambda *a, **kw: FakeResponse(body=["not", "a", "dict"]), + ) + assert render.render_url("https://example.com") is None + + +# ── Custom timeout --------------------------------------------------------- + + +def test_render_url_honors_timeout_env(configured, monkeypatch): + monkeypatch.setenv(render.RENDER_TIMEOUT_ENV, "5") + captured: list[float] = [] + + def fake_post(url, *, json, headers, timeout): # noqa: A002 + captured.append(timeout) + return FakeResponse(body={"data": {"markdown": "ok"}}) + + monkeypatch.setattr(render.httpx, "post", fake_post) + render.render_url("https://example.com") + assert captured[0] == 5.0 + + +def test_render_url_bad_timeout_env_falls_back_to_default(configured, monkeypatch): + monkeypatch.setenv(render.RENDER_TIMEOUT_ENV, "not-a-number") + captured: list[float] = [] + + def fake_post(url, *, json, headers, timeout): # noqa: A002 + captured.append(timeout) + return FakeResponse(body={"data": {"markdown": "ok"}}) + + monkeypatch.setattr(render.httpx, "post", fake_post) + render.render_url("https://example.com") + assert captured[0] == render.DEFAULT_TIMEOUT diff --git a/tests/test_shell_context.py b/tests/test_shell_context.py new file mode 100644 index 0000000..26e4280 --- /dev/null +++ b/tests/test_shell_context.py @@ -0,0 +1,175 @@ +"""Tests for the opt-in terminal-context block (shellllm.shell_context).""" + +from __future__ import annotations + +from shellllm.shell_context import ( + MAX_PIPED_CHARS, + build_piped_block, + build_shell_context_block, + redact, +) + +ENV_FULL = { + "SHELLLM_SHELL_CONTEXT": "output", + "SHELLLM_LAST_CMD": "git push origin main", + "SHELLLM_LAST_STATUS": "1", + "SHELLLM_RECENT_HISTORY": "git status\ngit add .\ngit commit -m wip", + "SHELLLM_PANE_OUTPUT": "error: failed to push some refs to 'origin'", +} + + +# ─── ladder levels ────────────────────────────────────────────────────── + + +def test_off_by_default(): + assert build_shell_context_block(env={}) == "" + + +def test_off_even_when_vars_are_set(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="off") + assert build_shell_context_block(env=env) == "" + + +def test_unknown_level_fails_safe_to_off(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="everything") + assert build_shell_context_block(env=env) == "" + + +def test_cmd_level_includes_last_command_and_status(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="cmd") + out = build_shell_context_block(env=env) + assert "git push origin main" in out + assert "1" in out + + +def test_cmd_level_excludes_history_and_output(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="cmd") + out = build_shell_context_block(env=env) + assert "git commit -m wip" not in out + assert "failed to push some refs" not in out + + +def test_history_level_includes_recent_commands(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="history") + out = build_shell_context_block(env=env) + assert "git commit -m wip" in out + assert "failed to push some refs" not in out + + +def test_output_level_includes_pane_output(): + out = build_shell_context_block(env=ENV_FULL) + assert "failed to push some refs" in out + + +def test_failed_status_is_annotated(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="cmd") + out = build_shell_context_block(env=env) + assert "failed" in out.lower() + + +def test_zero_status_not_annotated_as_failed(): + env = dict(ENV_FULL, SHELLLM_SHELL_CONTEXT="cmd", SHELLLM_LAST_STATUS="0") + out = build_shell_context_block(env=env) + assert "failed" not in out.lower() + + +def test_empty_when_level_on_but_no_data(): + env = {"SHELLLM_SHELL_CONTEXT": "cmd"} + assert build_shell_context_block(env=env) == "" + + +# ─── caps ─────────────────────────────────────────────────────────────── + + +def test_long_output_is_tail_truncated(): + env = dict(ENV_FULL, SHELLLM_PANE_OUTPUT="x" * 10_000 + "TAIL_MARKER") + out = build_shell_context_block(env=env) + assert "TAIL_MARKER" in out + assert len(out) < 6_000 + + +def test_history_capped_to_recent_lines(): + lines = [f"cmd-{i}" for i in range(50)] + env = dict( + ENV_FULL, + SHELLLM_SHELL_CONTEXT="history", + SHELLLM_RECENT_HISTORY="\n".join(lines), + ) + out = build_shell_context_block(env=env) + assert "cmd-49" in out # most recent kept + assert "cmd-0" not in out # oldest dropped + + +# ─── redaction ────────────────────────────────────────────────────────── + + +def test_redacts_keyed_assignment(): + out = redact("export OPENAI_API_KEY=sk-abc123def456ghi789jkl") + assert "sk-abc123def456ghi789jkl" not in out + assert "[redacted]" in out + + +def test_redacts_password_colon_form(): + out = redact("password: hunter2hunter2") + assert "hunter2hunter2" not in out + + +def test_redacts_bearer_header(): + out = redact("curl -H 'Authorization: Bearer abc.def.ghi'") + assert "abc.def.ghi" not in out + + +def test_redacts_aws_access_key(): + out = redact("AKIAIOSFODNN7EXAMPLE was leaked") + assert "AKIAIOSFODNN7EXAMPLE" not in out + + +def test_redacts_github_token(): + out = redact( + "git remote set-url origin https://ghp_abcdefghijklmnopqrstuvwxyz123456@github.com/x/y" + ) + assert "ghp_abcdefghijklmnopqrstuvwxyz123456" not in out + + +def test_redacts_jwt(): + jwt = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U" + out = redact(f"token is {jwt}") + assert jwt not in out + + +def test_keeps_git_sha(): + sha = "3bb70645a2b9e7c1d8f4a6b3c9d2e1f0a7b8c9d0" + out = redact(f"git checkout {sha}") + assert sha in out + + +def test_block_is_redacted(): + env = dict(ENV_FULL, SHELLLM_LAST_CMD="export STRIPE_SECRET=sk_live_abcdef123456") + out = build_shell_context_block(env=env) + assert "sk_live_abcdef123456" not in out + + +# ─── piped input ──────────────────────────────────────────────────────── + + +def test_piped_block_empty_for_blank_input(): + assert build_piped_block("") == "" + assert build_piped_block(" \n ") == "" + + +def test_piped_block_wraps_content(): + out = build_piped_block("error: linker `cc` not found") + assert "Piped input" in out + assert "linker `cc` not found" in out + + +def test_piped_block_is_redacted(): + out = build_piped_block("Authorization: Bearer super.secret.token1234") + assert "super.secret.token1234" not in out + + +def test_piped_block_keeps_tail_when_oversized(): + text = "x" * (MAX_PIPED_CHARS + 100) + "FINAL_ERROR" + out = build_piped_block(text) + assert "FINAL_ERROR" in out + assert "truncated" in out diff --git a/zsh/shellllm.zsh b/zsh/shellllm.zsh index 2077a3e..36e684d 100644 --- a/zsh/shellllm.zsh +++ b/zsh/shellllm.zsh @@ -15,6 +15,12 @@ : ${SHELLLM_CTX:=32768} : ${SHELLLM_LOG:=$HOME/.cache/shellllm/llama-server.log} : ${SHELLLM_EMBED_LOG:=$HOME/.cache/shellllm/llama-embed.log} +# Terminal-context ladder (see below). Defaults to `cmd` — previous +# command + exit status — which is what makes `,,` and "why did that +# fail" work out of the box. Local-first means this never leaves the +# machine unless you point SHELLLM_BASE_URL at a hosted API; set it to +# `off` to disable capture entirely. +: ${SHELLLM_SHELL_CONTEXT:=cmd} # ─── Tier registry ────────────────────────────────────────────────────── # @@ -42,6 +48,15 @@ _SHELLLM_TIER_DESC[smart]="latest coder-tuned model — best quality, needs down typeset -ga _SHELLLM_TIER_ORDER _SHELLLM_TIER_ORDER=(fast balanced smart) +# Tier → port. Balanced owns the default port so plain `??` + `?` keep +# their historical behavior; the other tiers get dedicated ports so two +# models can serve side by side and `, --fast` / `? --smart` can route +# a single call to a specific one. +typeset -gA _SHELLLM_TIER_PORT +_SHELLLM_TIER_PORT[fast]="${SHELLLM_PORT_FAST:-8091}" +_SHELLLM_TIER_PORT[balanced]="${SHELLLM_PORT_BALANCED:-$SHELLLM_PORT}" +_SHELLLM_TIER_PORT[smart]="${SHELLLM_PORT_SMART:-8093}" + # Embedding tiers — a separate registry because the model lineage and # size targets are different (we want something small and fast, not a # 27B chat model). `?? --start-embed ` resolves against this map. @@ -61,16 +76,123 @@ _SHELLLM_EMBED_DESC[nomic]="nomic-embed-text-v1.5 — strong general-purpose ret typeset -ga _SHELLLM_EMBED_ORDER _SHELLLM_EMBED_ORDER=(tiny bge nomic) +# ─── terminal context (opt-in privacy ladder) ────────────────────────── +# +# SHELLLM_SHELL_CONTEXT=off|cmd|history|output (default: cmd) +# +# off nothing captured +# cmd previous command + exit status (default) +# history + last 10 commands +# output + recent pane output (tmux only) +# +# Captured values are passed as per-invocation environment — nothing is +# exported into the shell, and the Python side re-checks the level and +# redacts secret-shaped strings before anything reaches the model. +function _shellllm_with_ctx() { + local last_status=$1; shift + local level="${SHELLLM_SHELL_CONTEXT:-off}" + if [[ $level != cmd && $level != history && $level != output ]]; then + "$@" + return $? + fi + local last_cmd hist="" out="" + # The command being executed is already in history, so the *previous* + # one is entry -2. + last_cmd="$(builtin fc -ln -2 -2 2>/dev/null)" + last_cmd="${last_cmd#"${last_cmd%%[![:space:]]*}"}" + if [[ $level == history || $level == output ]]; then + hist="$(builtin fc -ln -11 -2 2>/dev/null)" + fi + if [[ $level == output && -n ${TMUX:-} ]]; then + out="$(command tmux capture-pane -p -S -60 2>/dev/null)" + fi + # The level rides along explicitly: it may be a plain (unexported) + # shell variable, and the Python side re-checks it from the env. + SHELLLM_SHELL_CONTEXT="$level" \ + SHELLLM_LAST_STATUS="$last_status" \ + SHELLLM_LAST_CMD="$last_cmd" \ + SHELLLM_RECENT_HISTORY="$hist" \ + SHELLLM_PANE_OUTPUT="$out" \ + "$@" +} + +# ─── lazy start ───────────────────────────────────────────────────────── +# +# SHELLLM_AUTOSTART=1 makes `,` / `?` bring the default tier up on demand +# instead of erroring when the server is down. Off by default: starting +# a multi-GB model is a deliberate act, and the first start can take a +# minute. The health probe only runs when the feature is on. +function _shellllm_autostart() { + [[ "${SHELLLM_AUTOSTART:-0}" == 1 ]] || return 0 + _shellllm_server_up && return 0 + print -u2 -- "shellllm: llama-server is down — autostarting (SHELLLM_AUTOSTART=1)" + _shellllm_start >&2 +} + +# ─── per-call model routing ───────────────────────────────────────────── +# +# `, --fast …` / `? --smart …` send one call to a specific tier's server +# (see _SHELLLM_TIER_PORT). The flag is consumed here; the Python CLI +# never sees it — it just gets SHELLLM_BASE_URL pointed at the right +# port. The tier must already be up: `?? --start fast`. +# +# Only the FIRST argument routes, so prose mentioning --fast later in a +# prompt (`? what does --fast do in pip`) passes through untouched. +# +# Outputs via globals (zsh has no multi-value returns): +# _shellllm_route_url base URL override, or "" for the default +# _shellllm_route_args remaining args with the tier flag removed +typeset -g _shellllm_route_url +typeset -ga _shellllm_route_args + +function _shellllm_route() { + _shellllm_route_url="" + _shellllm_route_args=("$@") + local tier="" + case "${1:-}" in + --fast|--balanced|--smart) tier="${1#--}"; shift; _shellllm_route_args=("$@") ;; + *) return 0 ;; + esac + local port="${_SHELLLM_TIER_PORT[$tier]}" + if ! _shellllm_server_up "$port"; then + print -u2 -- "shellllm: tier '$tier' isn't running on :${port}" + print -u2 -- " what to do: ?? --start $tier" + return 1 + fi + _shellllm_route_url="http://127.0.0.1:${port}" +} + # ─── `,` — propose, never execute. Lands on the prompt line via print -z. -function ,() { +function _shellllm_comma_run() { + local last_status=$1; shift + _shellllm_route "$@" || return $? + set -- "${_shellllm_route_args[@]}" + [[ -n $_shellllm_route_url ]] && local -x SHELLLM_BASE_URL=$_shellllm_route_url + [[ -z $_shellllm_route_url ]] && { _shellllm_autostart || return $?; } local cmd - cmd=$(${=SHELLLM_COMMA} "$@") || return $? + cmd=$(_shellllm_with_ctx $last_status ${=SHELLLM_COMMA} "$@") || return $? [[ -n $cmd ]] && print -z -- "$cmd" } +function ,() { + _shellllm_comma_run $? "$@" +} + +# ─── `,,` — fix the previous command (`, --fix`). Uses the terminal- +# context ladder (on at `cmd` by default; SHELLLM_SHELL_CONTEXT=off +# disables it and `,,` with it). +function ,,() { + _shellllm_comma_run $? --fix "$@" +} + # ─── `?` — answer. `noglob` is required because `?` is a zsh glob char. function _shellllm_ask_fn() { - ${=SHELLLM_ASK} "$@" + local __last_status=$? + _shellllm_route "$@" || return $? + set -- "${_shellllm_route_args[@]}" + [[ -n $_shellllm_route_url ]] && local -x SHELLLM_BASE_URL=$_shellllm_route_url + [[ -z $_shellllm_route_url ]] && { _shellllm_autostart || return $?; } + _shellllm_with_ctx $__last_status ${=SHELLLM_ASK} "$@" } alias '?'='noglob _shellllm_ask_fn' @@ -103,7 +225,8 @@ function _shellllm_find_gguf() { } function _shellllm_server_up() { - curl -fsS -m 1 "http://127.0.0.1:${SHELLLM_PORT}/health" >/dev/null 2>&1 + local port="${1:-$SHELLLM_PORT}" + curl -fsS -m 1 "http://127.0.0.1:${port}/health" >/dev/null 2>&1 } function _shellllm_embed_up() { @@ -138,9 +261,9 @@ function _shellllm_list_tiers() { desc="${_SHELLLM_TIER_DESC[$tier]}" gguf=$(_shellllm_find_gguf "$repo") if [[ -f $gguf ]]; then - print -- " ${_G}✓${_N} ${_C}${tier}${_N} $desc" + print -- " ${_G}✓${_N} ${_C}${tier}${_N} ${_D}:${_SHELLLM_TIER_PORT[$tier]}${_N} $desc" else - print -- " ${_R}✗${_N} ${_C}${tier}${_N} $desc" + print -- " ${_R}✗${_N} ${_C}${tier}${_N} ${_D}:${_SHELLLM_TIER_PORT[$tier]}${_N} $desc" print -- " ${_D}huggingface-cli download $repo${_N}" fi done @@ -149,13 +272,12 @@ function _shellllm_list_tiers() { # ─── `??` — start (or stop / list / status) the local llama-server. # # ?? start the default tier (balanced) -# ?? --start fast start a specific tier -# ?? --start balanced -# ?? --start smart +# ?? --start fast start a tier on its own port (tiers can +# ?? --start smart run side by side; `, --fast` etc. routes) # ?? --model PATH start with an explicit gguf -# ?? --list show tiers and which are downloaded -# ?? --status up/down -# ?? --stop kill the server +# ?? --list show tiers, ports, and which are downloaded +# ?? --status which tiers/servers are up +# ?? --stop [tier] kill the default server, or one tier's # # If a tier isn't downloaded, you get a copy-pasteable # `huggingface-cli download` line. @@ -245,22 +367,39 @@ function _shellllm_start() { return 0 ;; --list-embed) _shellllm_list_embed_tiers; return 0 ;; --stop) - pkill -f "llama-server.*--port ${SHELLLM_PORT}" \ - && echo "stopped" || echo "(nothing to stop)" + local stop_port="$SHELLLM_PORT" stop_what="default" + if [[ -n "${2:-}" && -n "${_SHELLLM_TIER_PORT[${2:-_}]:-}" ]]; then + stop_port="${_SHELLLM_TIER_PORT[$2]}" + stop_what="$2" + fi + pkill -f "llama-server.*--port ${stop_port}" \ + && echo "stopped ${stop_what} (:${stop_port})" \ + || echo "(nothing to stop on :${stop_port})" return 0 ;; --status) - if _shellllm_server_up; then - echo "up → http://127.0.0.1:${SHELLLM_PORT}" - else - echo "down → run ?? to start" + local t p any=0 + for t in "${_SHELLLM_TIER_ORDER[@]}"; do + p="${_SHELLLM_TIER_PORT[$t]}" + if _shellllm_server_up "$p"; then + echo "up ${t} → http://127.0.0.1:${p}" + any=1 + fi + done + if (( ! any )); then + if _shellllm_server_up; then + echo "up → http://127.0.0.1:${SHELLLM_PORT}" + else + echo "down → run ?? to start" + fi fi return 0 ;; --list|-l) _shellllm_list_tiers; return 0 ;; --help|-h) print -- "?? — start the local llama-server" print -- " ?? [--start ] [--model PATH]" - print -- " ?? --list | --status | --stop" + print -- " ?? --list | --status | --stop [tier]" print -- " ?? --start-embed | --status-embed | --stop-embed | --list-embed" + print -- " per-call routing: , --fast … / ? --smart … (tier must be up)" _shellllm_list_tiers print _shellllm_list_embed_tiers @@ -276,16 +415,6 @@ function _shellllm_start() { return 1 fi - if _shellllm_server_up; then - if [[ -n $tier || -n $model ]]; then - print -u2 -- "?? llama-server is already up. To switch:" - print -u2 -- " ?? --stop && ?? --start ${tier:-…}" - return 1 - fi - echo "llama-server already running on :${SHELLLM_PORT}" - return 0 - fi - local extra_args="" if [[ -n $model ]]; then : # explicit model, no tier args @@ -324,27 +453,43 @@ function _shellllm_start() { return 1 fi - mkdir -p "$(dirname "$SHELLLM_LOG")" + # Tiers bind to their own ports so several can serve side by side + # (`, --fast` / `? --smart` route per call). Explicit --model and env + # fallbacks stay on the default port. + local port="$SHELLLM_PORT" + [[ -n $tier ]] && port="${_SHELLLM_TIER_PORT[$tier]}" + + if _shellllm_server_up "$port"; then + echo "llama-server already running on :${port}${tier:+ (tier $tier)}" + echo " to replace it: ?? --stop${tier:+ $tier} && ?? --start ${tier:-…}" + return 0 + fi + + local log="$SHELLLM_LOG" + [[ "$port" != "$SHELLLM_PORT" ]] && log="${SHELLLM_LOG}.${port}" + + mkdir -p "$(dirname "$log")" echo "starting llama-server" [[ -n $tier ]] && echo " tier : $tier" echo " model : $model" + echo " port : $port" [[ -n $extra_args ]] && echo " extra : $extra_args" - echo " log : $SHELLLM_LOG" + echo " log : $log" nohup llama-server \ -m "$model" \ -c "$SHELLLM_CTX" \ -ngl "$SHELLLM_NGL" \ --host 127.0.0.1 \ - --port "$SHELLLM_PORT" \ + --port "$port" \ ${=extra_args} \ - >"$SHELLLM_LOG" 2>&1 & + >"$log" 2>&1 & disown printf " waiting" local i for i in {1..120}; do - if _shellllm_server_up; then + if _shellllm_server_up "$port"; then echo " ready (${i}s)" return 0 fi @@ -354,8 +499,8 @@ function _shellllm_start() { print -u2 -- "" print -u2 -- "?? still not ready after 120s. what to do:" - print -u2 -- " 1. tail -50 $SHELLLM_LOG" - print -u2 -- " 2. lsof -iTCP:${SHELLLM_PORT} -sTCP:LISTEN" + print -u2 -- " 1. tail -50 $log" + print -u2 -- " 2. lsof -iTCP:${port} -sTCP:LISTEN" print -u2 -- " 3. if log says 'unknown option --spec-type': brew upgrade llama.cpp" return 1 }