From 549c5a0a64a68d48a706a1ccbd332c9c80caabcd Mon Sep 17 00:00:00 2001 From: Justin Bollinger Date: Fri, 24 Jul 2026 17:17:58 -0400 Subject: [PATCH] fix(llm): bound Ollama requests with a configurable timeout The Atomic Agents client was built with no timeout, so an Ollama server that accepted the TCP connection but never replied (most commonly a large model still loading into VRAM) left the CLI blocked in agent.run() forever. The caller's `except Exception` handler only ever fired on connection-refused. - llm.generate_candidates() takes a `timeout` parameter (default DEFAULT_TIMEOUT_SECONDS = 300.0) and forwards it to the OpenAI client. - openai.APITimeoutError is translated into a new domain-level llm.LLMTimeoutError so main.py need not import openai itself, keeping the atomic-agents/instructor dependency isolated to llm.py as documented. - hcatOllama passes the new `ollamaTimeout` config value and prints timeout-specific guidance (elapsed seconds, VRAM-loading hint, the setting to raise) instead of the misleading "ensure Ollama is running" message. - Documented `ollamaTimeout` in config.json.example and README. Also closes two test gaps: the defensive `except ValueError` handler in hcatOllama now has coverage, and the misleadingly-named `test_unknown_mode_raises` is split into explicit unit tests of _build_request's mode validation. Ticked the completed steps in the LLM Atomic Agents plan doc. Co-Authored-By: Claude --- CHANGELOG.md | 9 +++++++ README.md | 6 +++-- config.json.example | 1 + hate_crack/llm.py | 30 +++++++++++++++++++---- hate_crack/main.py | 15 +++++++++++- tests/test_hcat_ollama.py | 42 ++++++++++++++++++++++++++++++++ tests/test_llm.py | 50 +++++++++++++++++++++++++++++++++++++-- 7 files changed, 143 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3cf55af..33c3e84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,15 @@ Dates are omitted for releases predating this file; see the git tags for exact t menu: select the LLM attack (option 12), then choose "Wordlist" to derive basewords from a sample wordlist. +### Fixed + +- **The LLM attack no longer hangs forever waiting on Ollama.** Generation requests are now + bounded by a configurable timeout (`ollamaTimeout` in `config.json`, default 300 seconds). + Previously, if Ollama accepted the connection but never replied — most commonly a large + model still loading into VRAM — hate_crack sat at a frozen prompt with no recourse but + Ctrl-C. When the timeout fires you now get a specific message naming the elapsed timeout + and the setting to raise, instead of a misleading "ensure Ollama is running" hint. + ### Removed - **Automatic model pulling.** hate_crack no longer pulls missing Ollama models; pull them diff --git a/README.md b/README.md index 720a08e..9c24748 100644 --- a/README.md +++ b/README.md @@ -459,17 +459,19 @@ Set Hashview credentials in `config.json`: #### Ollama Configuration -The LLM Attack (option 12) uses Ollama to generate password candidates. Configure the model and context window in `config.json`: +The LLM Attack (option 12) uses Ollama to generate password candidates. Configure the model, context window, and request timeout in `config.json`: ```json { "ollamaModel": "qwen2.5:32b", - "ollamaNumCtx": 2048 + "ollamaNumCtx": 2048, + "ollamaTimeout": 300 } ``` - **`ollamaModel`** — The Ollama model used for candidate generation (default: `qwen2.5:32b`). The LLM attack uses structured (JSON) output, so choose a model with good tool/JSON support. - **`ollamaNumCtx`** — Context window size for the model (default: `2048`). +- **`ollamaTimeout`** — Seconds to wait for a generation response before giving up (default: `300`). Raise this if a large model is still loading into VRAM on the first request, which can otherwise exceed the timeout; hate_crack prints the elapsed timeout and this setting's name when it fires. - The Ollama URL defaults to `http://localhost:11434` (override via the `OLLAMA_HOST` env var). Ensure Ollama is running and the model is pulled (`ollama pull qwen2.5:32b`) before using the LLM Attack — hate_crack no longer auto-pulls missing models. The attack offers two generation modes: **Target info** (company / industry / location) and **Wordlist** (derive denylist basewords from a sample wordlist). diff --git a/config.json.example b/config.json.example index 23bcb3b..9efabe3 100644 --- a/config.json.example +++ b/config.json.example @@ -27,6 +27,7 @@ "hashmob_api_key": "", "ollamaModel": "qwen2.5:32b", "ollamaNumCtx": 2048, + "ollamaTimeout": 300, "omenTrainingList": "rockyou.txt", "omenMaxCandidates": 50000000, "pcfgRuleset": "DEFAULT", diff --git a/hate_crack/llm.py b/hate_crack/llm.py index 89ecd88..184d4b0 100644 --- a/hate_crack/llm.py +++ b/hate_crack/llm.py @@ -5,13 +5,22 @@ to this module only through ``generate_candidates``. """ import instructor -from openai import OpenAI +from openai import APITimeoutError, OpenAI from pydantic import Field from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema from atomic_agents.context import SystemPromptGenerator MAX_CANDIDATE_LEN = 128 +DEFAULT_TIMEOUT_SECONDS = 300.0 + + +class LLMTimeoutError(Exception): + """The LLM server accepted the request but did not respond in time. + + Raised instead of ``openai.APITimeoutError`` so callers do not need to import + ``openai`` themselves. + """ class GenerationInput(BaseIOSchema): @@ -93,17 +102,23 @@ def generate_candidates( num_ctx: int, mode: str, context_data: dict, + timeout: float = DEFAULT_TIMEOUT_SECONDS, ) -> list[str]: """Generate password candidates via an Ollama-backed AtomicAgent. + ``timeout`` is the number of seconds to wait for a generation response before + giving up; it bounds the whole request so a server that accepts the + connection but never replies (e.g. a large model still loading into VRAM) + cannot hang the caller forever. + Returns a deduped, length-capped list of candidate strings (may be empty). - Raises ValueError for an unknown mode. Client/connection errors propagate to - the caller. + Raises ValueError for an unknown mode and LLMTimeoutError if the request + exceeds ``timeout``. Other client/connection errors propagate to the caller. """ request = _build_request(mode, context_data) client = instructor.from_openai( - OpenAI(base_url=f"{url}/v1", api_key="ollama"), + OpenAI(base_url=f"{url}/v1", api_key="ollama", timeout=timeout), mode=instructor.Mode.JSON, ) prompt_generator = _TARGET_PROMPT if mode == "target" else _WORDLIST_PROMPT @@ -117,7 +132,12 @@ def generate_candidates( ) ) - result = agent.run(GenerationInput(request=request)) + try: + result = agent.run(GenerationInput(request=request)) + except APITimeoutError as e: + raise LLMTimeoutError( + f"no response from {url} within {timeout:g} seconds" + ) from e seen: set[str] = set() candidates: list[str] = [] diff --git a/hate_crack/main.py b/hate_crack/main.py index e9fc3a0..b801441 100755 --- a/hate_crack/main.py +++ b/hate_crack/main.py @@ -452,6 +452,7 @@ hcatDebugLogPath = os.path.expanduser(config_parser["hcatDebugLogPath"]) ollamaUrl = "http://" + os.environ.get("OLLAMA_HOST", "localhost:11434") ollamaModel = config_parser.get("ollamaModel", "qwen2.5:32b") ollamaNumCtx = int(config_parser.get("ollamaNumCtx", 2048)) +ollamaTimeout = float(config_parser.get("ollamaTimeout", 300)) omenTrainingList = config_parser.get("omenTrainingList", "rockyou.txt") omenMaxCandidates = int(config_parser.get("omenMaxCandidates", 1000000)) @@ -2072,8 +2073,20 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data): print(f"Generating password candidates via Ollama ({ollamaModel})...") try: candidates = llm.generate_candidates( - ollamaUrl, ollamaModel, ollamaNumCtx, mode, gen_context + ollamaUrl, + ollamaModel, + ollamaNumCtx, + mode, + gen_context, + timeout=ollamaTimeout, ) + except llm.LLMTimeoutError: + print(f"Error: the Ollama request timed out after {ollamaTimeout:g} seconds.") + print( + f"The model ({ollamaModel}) may still be loading into VRAM. Retry, or " + 'raise "ollamaTimeout" in config.json to wait longer.' + ) + return except ValueError as e: # Defensive: mode is already validated above, but keep an explicit, # non-misleading message if generate_candidates ever rejects its input. diff --git a/tests/test_hcat_ollama.py b/tests/test_hcat_ollama.py index a63c4d5..7786535 100644 --- a/tests/test_hcat_ollama.py +++ b/tests/test_hcat_ollama.py @@ -179,6 +179,48 @@ def test_generation_error_reports_and_aborts(ollama_env, capsys): popen.assert_not_called() +def test_timeout_error_reports_timeout_guidance(ollama_env, capsys): + with ollama_globals(ollama_env.tmp_path), \ + mock.patch.object(hc_main, "ollamaTimeout", 300.0), \ + mock.patch("hate_crack.main.llm.generate_candidates", + side_effect=hc_main.llm.LLMTimeoutError("timed out")), \ + mock.patch("subprocess.Popen") as popen: + hc_main.hcatOllama("0", ollama_env.hash_file, "target", + {"company": "X", "industry": "Y", "location": "Z"}) + captured = capsys.readouterr() + assert "timed out" in captured.out.lower() + assert "300" in captured.out + assert "ollamaTimeout" in captured.out + assert "Ensure Ollama is running" not in captured.out + popen.assert_not_called() + assert not os.path.isfile(f"{ollama_env.hash_file}.ollama_candidates") + + +def test_value_error_reports_message_and_aborts(ollama_env, capsys): + """Covers the defensive `except ValueError` handler in hcatOllama.""" + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.llm.generate_candidates", + side_effect=ValueError("boom")), \ + mock.patch("subprocess.Popen") as popen: + hc_main.hcatOllama("0", ollama_env.hash_file, "target", + {"company": "X", "industry": "Y", "location": "Z"}) + captured = capsys.readouterr() + assert "Error: boom" in captured.out + popen.assert_not_called() + assert not os.path.isfile(f"{ollama_env.hash_file}.ollama_candidates") + + +def test_timeout_config_forwarded_to_generate_candidates(ollama_env): + with ollama_globals(ollama_env.tmp_path), \ + mock.patch.object(hc_main, "ollamaTimeout", 77.0), \ + mock.patch("hate_crack.main.llm.generate_candidates", + return_value=["Password1"]) as gen, \ + mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", ollama_env.hash_file, "target", + {"company": "X", "industry": "Y", "location": "Z"}) + assert gen.call_args.kwargs["timeout"] == 77.0 + + def test_unknown_mode_prints_error(ollama_env, capsys): with ollama_globals(ollama_env.tmp_path), \ mock.patch("hate_crack.main.llm.generate_candidates") as gen: diff --git a/tests/test_llm.py b/tests/test_llm.py index f555bfc..f34b3b8 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -3,7 +3,9 @@ import os from unittest import mock +import httpx import instructor +import openai import pytest os.environ["HATE_CRACK_SKIP_INIT"] = "1" @@ -94,8 +96,52 @@ def test_num_ctx_forwarded_via_model_api_parameters(): assert config.model_api_parameters["extra_body"]["options"]["num_ctx"] == 4096 -def test_unknown_mode_raises(): - with pytest.raises(ValueError): +def test_build_request_rejects_unknown_mode(): + """Unit test of _build_request's mode validation only — no agent involved.""" + with pytest.raises(ValueError, match="Unknown LLM generation mode: bogus"): + llm._build_request("bogus", {}) + + +def test_generate_candidates_rejects_unknown_mode_before_building_client(): + """generate_candidates surfaces _build_request's ValueError to its caller.""" + with pytest.raises(ValueError, match="Unknown LLM generation mode: bogus"): llm.generate_candidates( "http://localhost:11434", "qwen2.5:32b", 2048, "bogus", {}, ) + + +def test_timeout_forwarded_to_openai_client(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"]) + with p_instr, p_openai as openai_cls, p_agent: + llm.generate_candidates( + "http://localhost:11434", "qwen2.5:32b", 2048, + "target", {"company": "X", "industry": "Y", "location": "Z"}, + timeout=42.5, + ) + assert openai_cls.call_args.kwargs["timeout"] == 42.5 + + +def test_default_timeout_used_when_omitted(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"]) + with p_instr, p_openai as openai_cls, p_agent: + llm.generate_candidates( + "http://localhost:11434", "qwen2.5:32b", 2048, + "target", {"company": "X", "industry": "Y", "location": "Z"}, + ) + assert llm.DEFAULT_TIMEOUT_SECONDS == 300.0 + assert openai_cls.call_args.kwargs["timeout"] == llm.DEFAULT_TIMEOUT_SECONDS + + +def test_api_timeout_reraised_as_domain_error(): + """openai.APITimeoutError is translated into llm.LLMTimeoutError.""" + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"]) + agent_instance.run.side_effect = openai.APITimeoutError( + request=httpx.Request("POST", "http://localhost:11434/v1/chat/completions") + ) + with p_instr, p_openai, p_agent: + with pytest.raises(llm.LLMTimeoutError): + llm.generate_candidates( + "http://localhost:11434", "qwen2.5:32b", 2048, + "target", {"company": "X", "industry": "Y", "location": "Z"}, + timeout=1.0, + )