mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d8ad8e3572 | ||
|
|
788a5393db | ||
|
|
7b82ff8432 | ||
|
|
9de747d135 | ||
|
|
b313d78f60 | ||
|
|
e037d8d727 | ||
|
|
fade37025d | ||
|
|
f968f8e5a7 | ||
|
|
ac0014fe65 | ||
|
|
86282e83a8 |
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "strix-agent"
|
||||
version = "1.4.0"
|
||||
version = "1.4.1"
|
||||
description = "Open-source AI Hackers for your apps"
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
|
||||
+13
-20
@@ -18,12 +18,12 @@ import logging
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
@@ -221,26 +221,19 @@ def _first(query: dict[str, list[str]], key: str) -> str | None:
|
||||
|
||||
|
||||
def _post_form(payload: dict[str, str]) -> dict[str, Any]:
|
||||
body = urllib.parse.urlencode(payload).encode("ascii")
|
||||
request = urllib.request.Request( # noqa: S310 - fixed https OAuth endpoint
|
||||
TOKEN_URL,
|
||||
data=body,
|
||||
headers={
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen( # noqa: S310 # nosec B310 - fixed https endpoint
|
||||
request, timeout=_TOKEN_TIMEOUT
|
||||
) as response:
|
||||
data = json.loads(response.read() or b"{}")
|
||||
except urllib.error.HTTPError as exc:
|
||||
detail = exc.read().decode("utf-8", "replace")[:300]
|
||||
raise CodexAuthError("token_http_error", f"HTTP {exc.code}: {detail}") from exc
|
||||
except (urllib.error.URLError, TimeoutError, OSError) as exc:
|
||||
response = requests.post(
|
||||
TOKEN_URL,
|
||||
data=payload,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=_TOKEN_TIMEOUT,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise CodexAuthError("unavailable", str(exc)) from exc
|
||||
if response.status_code >= 400:
|
||||
detail = response.text[:300]
|
||||
raise CodexAuthError("token_http_error", f"HTTP {response.status_code}: {detail}")
|
||||
data = json.loads(response.content or b"{}")
|
||||
if not isinstance(data, dict):
|
||||
raise CodexAuthError("bad_response", "token endpoint returned non-object")
|
||||
return data
|
||||
|
||||
+160
-3
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import contextlib
|
||||
import inspect
|
||||
import os
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents import (
|
||||
@@ -13,6 +14,8 @@ from agents import (
|
||||
set_tracing_disabled,
|
||||
)
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.fake_id import FAKE_RESPONSES_ID
|
||||
from agents.models.interface import Model
|
||||
from agents.models.multi_provider import MultiProvider
|
||||
from agents.models.openai_responses import OpenAIResponsesModel
|
||||
from agents.retry import (
|
||||
@@ -21,6 +24,12 @@ from agents.retry import (
|
||||
RetryPolicyContext,
|
||||
retry_policies,
|
||||
)
|
||||
from openai.types.responses import (
|
||||
Response,
|
||||
ResponseCompletedEvent,
|
||||
ResponseOutputItemDoneEvent,
|
||||
ResponseUsage,
|
||||
)
|
||||
from openai.types.shared import Reasoning
|
||||
|
||||
from strix.config import codex
|
||||
@@ -30,8 +39,14 @@ from strix.config.loader import load_settings
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from agents.models.interface import Model, ModelProvider
|
||||
from agents.agent_output import AgentOutputSchemaBase
|
||||
from agents.handoffs import Handoff
|
||||
from agents.items import ModelResponse, TResponseInputItem, TResponseStreamEvent
|
||||
from agents.models.interface import ModelProvider, ModelTracing
|
||||
from agents.tool import Tool
|
||||
from agents.usage import Usage
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.responses import ResponsePromptParam
|
||||
|
||||
from strix.config.settings import ReasoningEffort, Settings
|
||||
|
||||
@@ -135,6 +150,99 @@ class _CodexResponsesModel(OpenAIResponsesModel):
|
||||
await result
|
||||
|
||||
|
||||
def _to_response_usage(usage: Usage | None) -> ResponseUsage | None:
|
||||
if usage is None:
|
||||
return None
|
||||
return ResponseUsage(
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
input_tokens_details=usage.input_tokens_details,
|
||||
output_tokens_details=usage.output_tokens_details,
|
||||
)
|
||||
|
||||
|
||||
class _NonStreamingModel(Model):
|
||||
"""Run a model non-streamed but expose the streaming interface the runner uses.
|
||||
|
||||
Some OpenAI-compatible endpoints (notably gateways serving reasoning models)
|
||||
return valid ``tool_calls`` for a non-streamed completion but, when streamed,
|
||||
emit the tool call as plain text or drop it and close the stream — leaving
|
||||
Strix's tool-driven loop with nothing to execute. Selecting
|
||||
``STRIX_STREAM_MODE=never`` routes through this wrapper, which makes the real
|
||||
request non-streamed (where tool calling works) and synthesizes the minimal
|
||||
event sequence the runner consumes from a stream, so the rest of the pipeline
|
||||
is unchanged. The only user-visible difference is no token-by-token output.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: Model) -> None:
|
||||
self._inner = inner
|
||||
|
||||
async def get_response(self, *args: Any, **kwargs: Any) -> ModelResponse:
|
||||
return await self._inner.get_response(*args, **kwargs)
|
||||
|
||||
async def stream_response(
|
||||
self,
|
||||
system_instructions: str | None,
|
||||
input: str | list[TResponseInputItem], # noqa: A002
|
||||
model_settings: ModelSettings,
|
||||
tools: list[Tool],
|
||||
output_schema: AgentOutputSchemaBase | None,
|
||||
handoffs: list[Handoff],
|
||||
tracing: ModelTracing,
|
||||
*,
|
||||
previous_response_id: str | None = None,
|
||||
conversation_id: str | None = None,
|
||||
prompt: ResponsePromptParam | None = None,
|
||||
) -> AsyncIterator[TResponseStreamEvent]:
|
||||
model_response = await self._inner.get_response(
|
||||
system_instructions,
|
||||
input,
|
||||
model_settings,
|
||||
tools,
|
||||
output_schema,
|
||||
handoffs,
|
||||
tracing,
|
||||
previous_response_id=previous_response_id,
|
||||
conversation_id=conversation_id,
|
||||
prompt=prompt,
|
||||
)
|
||||
|
||||
sequence = 0
|
||||
for index, item in enumerate(model_response.output):
|
||||
yield ResponseOutputItemDoneEvent(
|
||||
item=item,
|
||||
output_index=index,
|
||||
type="response.output_item.done",
|
||||
sequence_number=sequence,
|
||||
)
|
||||
sequence += 1
|
||||
|
||||
response = Response(
|
||||
id=model_response.response_id or FAKE_RESPONSES_ID,
|
||||
created_at=time.time(),
|
||||
model=str(getattr(self._inner, "model", "")),
|
||||
object="response",
|
||||
output=model_response.output,
|
||||
tool_choice="auto",
|
||||
tools=[],
|
||||
parallel_tool_calls=False,
|
||||
usage=_to_response_usage(model_response.usage),
|
||||
)
|
||||
yield ResponseCompletedEvent(
|
||||
response=response,
|
||||
type="response.completed",
|
||||
sequence_number=sequence,
|
||||
)
|
||||
|
||||
def get_retry_advice(self, request: Any) -> Any:
|
||||
return self._inner.get_retry_advice(request)
|
||||
|
||||
|
||||
def _should_run_non_streamed(settings: Settings) -> bool:
|
||||
return settings.llm.stream_mode == "never"
|
||||
|
||||
|
||||
class StrixProvider(MultiProvider):
|
||||
"""Route any non-OpenAI prefix through LiteLLM with the prefix preserved,
|
||||
so users type ``deepseek/deepseek-chat`` rather than
|
||||
@@ -159,14 +267,18 @@ class StrixProvider(MultiProvider):
|
||||
return self._get_fallback_provider("litellm"), original_model_name
|
||||
|
||||
def get_model(self, model_name: str | None) -> Model:
|
||||
settings = load_settings()
|
||||
slug = codex.subscription_model(model_name)
|
||||
if slug:
|
||||
return _CodexResponsesModel(
|
||||
slug,
|
||||
codex.get_subscription_client(),
|
||||
reasoning_effort=load_settings().llm.reasoning_effort,
|
||||
reasoning_effort=settings.llm.reasoning_effort,
|
||||
)
|
||||
return super().get_model(model_name)
|
||||
model = super().get_model(model_name)
|
||||
if _should_run_non_streamed(settings):
|
||||
return _NonStreamingModel(model)
|
||||
return model
|
||||
|
||||
|
||||
DEFAULT_MODEL_RETRY = ModelRetrySettings(
|
||||
@@ -277,6 +389,51 @@ def _configure_litellm_compatibility() -> None:
|
||||
litellm.suppress_debug_info = True
|
||||
|
||||
_register_litellm_cost_callback()
|
||||
_install_openrouter_stream_cost_capture()
|
||||
|
||||
|
||||
def _install_openrouter_stream_cost_capture() -> None:
|
||||
"""Preserve OpenRouter's per-stream cost, which LiteLLM drops when streaming.
|
||||
|
||||
OpenRouter reports the real charge in ``usage.cost`` of the final stream
|
||||
chunk, but LiteLLM rebuilds streamed responses from token-only fields and
|
||||
discards it (its non-streamed path stashes the cost in hidden params; the
|
||||
streaming path does not). Every scan streams, so without this the cost is
|
||||
lost and Strix falls back to a cost-map estimate that is missing entirely
|
||||
for new models (e.g. kimi-k3), reporting $0. Subclass the OpenRouter
|
||||
streaming handler to record the cost keyed by response id so the cost
|
||||
callback can recover the exact charge for the matching rebuilt response.
|
||||
"""
|
||||
import litellm
|
||||
from litellm.llms.openrouter.chat.transformation import (
|
||||
OpenRouterChatCompletionStreamingHandler,
|
||||
OpenrouterConfig,
|
||||
)
|
||||
|
||||
from strix.report.state import streamed_openrouter_costs
|
||||
|
||||
class _StrixOpenRouterStreamingHandler(OpenRouterChatCompletionStreamingHandler):
|
||||
def chunk_parser(self, chunk: dict[str, Any]) -> Any:
|
||||
stream = super().chunk_parser(chunk)
|
||||
streamed_openrouter_costs.remember(
|
||||
chunk.get("id") or getattr(stream, "id", None), chunk.get("usage")
|
||||
)
|
||||
return stream
|
||||
|
||||
class _StrixOpenrouterConfig(OpenrouterConfig):
|
||||
def get_model_response_iterator(
|
||||
self, streaming_response: Any, sync_stream: bool, json_mode: bool | None = False
|
||||
) -> Any:
|
||||
return _StrixOpenRouterStreamingHandler(
|
||||
streaming_response=streaming_response,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
|
||||
# LiteLLM's provider-config factory reads litellm.OpenrouterConfig at call
|
||||
# time, so overriding the attribute is enough for the subclass to take
|
||||
# effect. (type: ignore — mypy rejects reassigning a class attribute.)
|
||||
litellm.OpenrouterConfig = _StrixOpenrouterConfig # type: ignore[misc]
|
||||
|
||||
|
||||
_OPENROUTER_ATTRIBUTION_HEADERS = {
|
||||
|
||||
@@ -9,6 +9,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
|
||||
StreamMode = Literal["auto", "always", "never"]
|
||||
|
||||
_BASE_CONFIG = SettingsConfigDict(
|
||||
case_sensitive=False,
|
||||
@@ -40,6 +41,12 @@ class LlmSettings(BaseSettings):
|
||||
default=False,
|
||||
alias="STRIX_FORCE_REQUIRED_TOOL_CHOICE",
|
||||
)
|
||||
# auto/always stream; never runs non-streamed, for endpoints that stream tool
|
||||
# calls incorrectly (some OpenAI-compatible gateways serving reasoning models).
|
||||
stream_mode: StreamMode = Field(
|
||||
default="auto",
|
||||
alias="STRIX_STREAM_MODE",
|
||||
)
|
||||
prompt_cache: bool = Field(
|
||||
default=True,
|
||||
alias="STRIX_PROMPT_CACHE",
|
||||
|
||||
@@ -200,7 +200,9 @@ class AgentCoordinator:
|
||||
logger.info("agent.status %s=%s", agent_id, status)
|
||||
await self._maybe_snapshot()
|
||||
|
||||
async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool:
|
||||
async def send(
|
||||
self, target_agent_id: str, message: dict[str, Any], *, interrupt: bool = True
|
||||
) -> bool:
|
||||
"""Deliver a user/peer message by appending it to the target SDK session."""
|
||||
if message.get("from") == "user" and self._budget_paused:
|
||||
await self.resume_from_budget_pause(exclude=target_agent_id)
|
||||
@@ -211,7 +213,7 @@ class AgentCoordinator:
|
||||
runtime = self.runtimes.setdefault(target_agent_id, AgentRuntime())
|
||||
session = runtime.session
|
||||
stream = runtime.stream
|
||||
interrupt = runtime.interrupt_on_message
|
||||
interrupt_on_message = runtime.interrupt_on_message
|
||||
if session is None:
|
||||
logger.warning(
|
||||
"agent.send dropped target=%s because its SDK session is not attached",
|
||||
@@ -230,7 +232,7 @@ class AgentCoordinator:
|
||||
async with self._lock:
|
||||
self.pending_counts[target_agent_id] = self.pending_counts.get(target_agent_id, 0) + 1
|
||||
self.runtimes.setdefault(target_agent_id, AgentRuntime()).wake.set()
|
||||
if stream is not None and interrupt:
|
||||
if stream is not None and interrupt and interrupt_on_message:
|
||||
stream.cancel(mode="immediate")
|
||||
await self._maybe_snapshot()
|
||||
return True
|
||||
|
||||
+30
-1
@@ -21,6 +21,7 @@ from openai import (
|
||||
RateLimitError,
|
||||
)
|
||||
|
||||
from strix.config import codex
|
||||
from strix.core.hooks import (
|
||||
BudgetExceededError,
|
||||
BudgetPausedError,
|
||||
@@ -88,6 +89,11 @@ async def _compact_session(
|
||||
)
|
||||
|
||||
|
||||
_GUARDRAIL_PARK_ERROR = (
|
||||
"Blocked by the model's content guardrail (flagged as a possible cybersecurity risk). "
|
||||
"Set STRIX_LLM to a model that isn't blocked and resume the scan to continue."
|
||||
)
|
||||
|
||||
_TRANSIENT_MODEL_STATUS_CODES = frozenset({408, 500, 502, 503, 504})
|
||||
_MAX_TRANSIENT_MODEL_RETRIES = 4
|
||||
_TRANSIENT_MODEL_RETRY_BASE_DELAY_S = 2.0
|
||||
@@ -304,6 +310,7 @@ async def respawn_subagents(
|
||||
if coordinator.parent_of.get(aid) is None or aid == root_id:
|
||||
continue
|
||||
md["_restored_status"] = status
|
||||
md["_restored_error"] = coordinator.errors.get(aid)
|
||||
candidates.append(
|
||||
(
|
||||
aid,
|
||||
@@ -316,7 +323,8 @@ async def respawn_subagents(
|
||||
for child_id, name, parent_id, md in candidates:
|
||||
try:
|
||||
restored_status = str(md.get("_restored_status") or "running")
|
||||
start_parked = interactive and restored_status != "running"
|
||||
recoverable_park = restored_status == "waiting" and bool(md.get("_restored_error"))
|
||||
start_parked = interactive and restored_status != "running" and not recoverable_park
|
||||
|
||||
if start_parked:
|
||||
logger.warning(
|
||||
@@ -572,6 +580,10 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
if session is not None:
|
||||
input_data = []
|
||||
continue
|
||||
if codex.is_content_guardrail_error(exc):
|
||||
return await _handle_content_guardrail(
|
||||
coordinator, agent_id, exc, interactive=interactive
|
||||
)
|
||||
if not interactive:
|
||||
raise
|
||||
if isinstance(exc, MaxTurnsExceeded):
|
||||
@@ -589,6 +601,22 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
return stream
|
||||
|
||||
|
||||
async def _handle_content_guardrail(
|
||||
coordinator: AgentCoordinator,
|
||||
agent_id: str,
|
||||
exc: BaseException,
|
||||
*,
|
||||
interactive: bool,
|
||||
) -> RunResultBase | None:
|
||||
logger.warning("agent %s blocked by the model's content guardrail: %s", agent_id, exc)
|
||||
if interactive:
|
||||
await coordinator.set_status(agent_id, "waiting", error=_GUARDRAIL_PARK_ERROR)
|
||||
return None
|
||||
await coordinator.set_status(agent_id, "failed", error=_GUARDRAIL_PARK_ERROR)
|
||||
await _notify_parent_on_terminal(coordinator, agent_id, "failed")
|
||||
return None
|
||||
|
||||
|
||||
async def _settle_run_result(
|
||||
coordinator: AgentCoordinator,
|
||||
agent_id: str,
|
||||
@@ -685,6 +713,7 @@ async def _notify_parent_on_terminal(
|
||||
"priority": "high",
|
||||
"content": template.format(name=name, agent_id=agent_id),
|
||||
},
|
||||
interrupt=False,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -376,6 +376,12 @@ async def run_strix_scan(
|
||||
|
||||
async with coordinator._lock:
|
||||
root_status = coordinator.statuses.get(root_id)
|
||||
root_error = coordinator.errors.get(root_id)
|
||||
|
||||
root_recoverable_park = root_status == "waiting" and bool(root_error)
|
||||
root_start_parked = bool(
|
||||
interactive and is_resume and root_status != "running" and not root_recoverable_park
|
||||
)
|
||||
|
||||
result = await run_agent_loop(
|
||||
agent=root_agent,
|
||||
@@ -387,7 +393,7 @@ async def run_strix_scan(
|
||||
agent_id=root_id,
|
||||
interactive=interactive,
|
||||
session=root_session,
|
||||
start_parked=bool(interactive and is_resume and root_status != "running"),
|
||||
start_parked=root_start_parked,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
|
||||
@@ -884,7 +884,7 @@ def display_completion_message(args: argparse.Namespace, results_path: Path) ->
|
||||
view_text = Text()
|
||||
view_text.append("\n")
|
||||
view_text.append("View", style="dim")
|
||||
view_text.append(" ")
|
||||
view_text.append(" ")
|
||||
view_text.append(f"strix view {args.run_name}", style="#22c55e")
|
||||
panel_parts.extend(["\n", view_text])
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ class TuiLiveView:
|
||||
self.events: list[dict[str, Any]] = []
|
||||
self._next_event_id = 1
|
||||
self._open_assistant_event_by_agent: dict[str, dict[str, Any]] = {}
|
||||
self._tool_event_by_call_id: dict[str, dict[str, Any]] = {}
|
||||
self._tool_event_by_agent_and_call_id: dict[tuple[str, str], dict[str, Any]] = {}
|
||||
|
||||
def hydrate_from_run_dir(self, run_dir: Path) -> None:
|
||||
state_dir = runtime_state_dir(run_dir)
|
||||
@@ -223,7 +223,8 @@ class TuiLiveView:
|
||||
timestamp: str | None = None,
|
||||
) -> None:
|
||||
call_id = call["call_id"]
|
||||
existing = self._tool_event_by_call_id.get(call_id)
|
||||
event_key = (agent_id, call_id)
|
||||
existing = self._tool_event_by_agent_and_call_id.get(event_key)
|
||||
tool_data = {
|
||||
"tool_name": call["tool_name"],
|
||||
"args": call["args"],
|
||||
@@ -233,7 +234,7 @@ class TuiLiveView:
|
||||
}
|
||||
if existing is None:
|
||||
event = self._append_event(agent_id, "tool", tool_data, timestamp=timestamp)
|
||||
self._tool_event_by_call_id[call_id] = event
|
||||
self._tool_event_by_agent_and_call_id[event_key] = event
|
||||
else:
|
||||
existing["data"].update(tool_data)
|
||||
self._bump_event(existing, timestamp=timestamp)
|
||||
@@ -249,7 +250,8 @@ class TuiLiveView:
|
||||
timestamp: str | None = None,
|
||||
) -> None:
|
||||
call_id = output["call_id"]
|
||||
event = self._tool_event_by_call_id.get(call_id)
|
||||
event_key = (agent_id, call_id)
|
||||
event = self._tool_event_by_agent_and_call_id.get(event_key)
|
||||
if event is None:
|
||||
event = self._append_event(
|
||||
agent_id,
|
||||
@@ -263,7 +265,7 @@ class TuiLiveView:
|
||||
},
|
||||
timestamp=timestamp,
|
||||
)
|
||||
self._tool_event_by_call_id[call_id] = event
|
||||
self._tool_event_by_agent_and_call_id[event_key] = event
|
||||
|
||||
result = _parse_json_value(output["output"])
|
||||
event["data"]["result"] = result
|
||||
|
||||
@@ -11,11 +11,10 @@ import tempfile
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.parse import urlparse
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
import docker
|
||||
import requests
|
||||
from docker.errors import DockerException, ImageNotFound
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
@@ -1088,13 +1087,12 @@ def resolve_diff_scope_context(
|
||||
def _is_http_git_repo(url: str) -> bool:
|
||||
check_url = f"{url.rstrip('/')}/info/refs?service=git-upload-pack"
|
||||
try:
|
||||
req = Request(check_url, headers={"User-Agent": "git/strix"}) # noqa: S310
|
||||
with urlopen(req, timeout=10) as resp: # noqa: S310 # nosec B310
|
||||
return "x-git-upload-pack-advertisement" in resp.headers.get("Content-Type", "")
|
||||
except HTTPError as e:
|
||||
return e.code == 401
|
||||
except (URLError, OSError, ValueError):
|
||||
resp = requests.get(check_url, headers={"User-Agent": "git/strix"}, timeout=10)
|
||||
except (requests.RequestException, ValueError):
|
||||
return False
|
||||
if resp.status_code >= 400:
|
||||
return resp.status_code == 401
|
||||
return "x-git-upload-pack-advertisement" in resp.headers.get("Content-Type", "")
|
||||
|
||||
|
||||
def infer_target_type(target: str) -> tuple[str, dict[str, str]]: # noqa: PLR0911
|
||||
|
||||
@@ -15,12 +15,12 @@ import base64
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
|
||||
from strix.config.loader import load_settings
|
||||
|
||||
|
||||
@@ -147,21 +147,17 @@ def _post_json(path: str, payload: dict[str, Any], *, timeout: int) -> tuple[int
|
||||
map, not raised.
|
||||
"""
|
||||
url = f"{_app_url()}{path}"
|
||||
body = json.dumps(payload).encode("utf-8")
|
||||
request = urllib.request.Request( # noqa: S310 - fixed https relay URL
|
||||
url,
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json", "Accept": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response: # noqa: S310 # nosec B310
|
||||
return response.status, _parse_body(response.read())
|
||||
except urllib.error.HTTPError as exc:
|
||||
return exc.code, _parse_body(exc.read())
|
||||
except (urllib.error.URLError, TimeoutError, OSError) as exc:
|
||||
response = requests.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=timeout,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
logger.warning("relay request to %s failed: %s", path, exc)
|
||||
raise RelayError("unavailable") from exc
|
||||
return response.status_code, _parse_body(response.content)
|
||||
|
||||
|
||||
def _parse_body(raw: bytes) -> dict[str, Any]:
|
||||
|
||||
@@ -107,8 +107,11 @@ def resolve_run_dir(base_dir: Path, run_param: str | None, default_run_dir: Path
|
||||
return candidate
|
||||
|
||||
|
||||
# Name of the cookie carrying the per-process session capability.
|
||||
SESSION_COOKIE = "strix_viewer_session"
|
||||
# Prefix of the cookie carrying the per-process session capability. The bound
|
||||
# port is appended (``strix_viewer_session_<port>``) because browsers scope
|
||||
# cookies by host only, never by port: concurrent viewers on 127.0.0.1 would
|
||||
# otherwise share one cookie slot and clobber each other's session.
|
||||
SESSION_COOKIE_PREFIX = "strix_viewer_session"
|
||||
|
||||
|
||||
class _ViewerState:
|
||||
@@ -135,6 +138,9 @@ class _ViewerState:
|
||||
# enough to steer a live scan, trigger a report, or browse history --
|
||||
# the token is never handed to a caller who merely reaches ``/``.
|
||||
self.session_token = secrets.token_urlsafe(32)
|
||||
# Finalized in ``serve()`` once the port is known (the server binds
|
||||
# after this state is constructed); see SESSION_COOKIE_PREFIX.
|
||||
self.cookie_name = SESSION_COOKIE_PREFIX
|
||||
|
||||
|
||||
def _make_handler(state: _ViewerState) -> type[BaseHTTPRequestHandler]:
|
||||
@@ -476,7 +482,7 @@ def _make_handler(state: _ViewerState) -> type[BaseHTTPRequestHandler]:
|
||||
the browser this process handed the page to can pass. A direct
|
||||
caller on an exposed port has no cookie and is rejected.
|
||||
"""
|
||||
supplied = self._cookies().get(SESSION_COOKIE, "")
|
||||
supplied = self._cookies().get(state.cookie_name, "")
|
||||
return bool(supplied) and secrets.compare_digest(supplied, state.session_token)
|
||||
|
||||
def _token_presented(self, query: dict[str, list[str]]) -> bool:
|
||||
@@ -512,7 +518,7 @@ def _make_handler(state: _ViewerState) -> type[BaseHTTPRequestHandler]:
|
||||
# SameSite=Strict (never sent from a cross-site context).
|
||||
self.send_header(
|
||||
"Set-Cookie",
|
||||
f"{SESSION_COOKIE}={state.session_token}; Path=/; HttpOnly; SameSite=Strict",
|
||||
f"{state.cookie_name}={state.session_token}; Path=/; HttpOnly; SameSite=Strict",
|
||||
)
|
||||
self.end_headers()
|
||||
self.wfile.write(content)
|
||||
@@ -586,6 +592,7 @@ def serve(
|
||||
|
||||
httpd.daemon_threads = True
|
||||
bound_port = int(httpd.server_address[1])
|
||||
state.cookie_name = f"{SESSION_COOKIE_PREFIX}_{bound_port}"
|
||||
url = f"http://{host}:{bound_port}"
|
||||
|
||||
thread = threading.Thread(target=httpd.serve_forever, name="strix-viewer", daemon=True)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
from datetime import UTC, datetime
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
@@ -95,6 +96,8 @@ def get_global_report_state() -> Optional["ReportState"]:
|
||||
def set_global_report_state(report_state: "ReportState") -> None:
|
||||
global _global_report_state # noqa: PLW0603
|
||||
_global_report_state = report_state
|
||||
# New run: drop any streamed-cost entries a prior run left unconsumed.
|
||||
streamed_openrouter_costs.clear()
|
||||
|
||||
|
||||
class ReportState:
|
||||
@@ -507,6 +510,72 @@ class ReportState:
|
||||
self._sync_llm_usage_record()
|
||||
|
||||
|
||||
def openrouter_stream_cost(usage: Any) -> float | None:
|
||||
"""Total OpenRouter-reported cost from a raw stream ``usage`` block, or None.
|
||||
|
||||
Non-BYOK responses bill everything to ``usage.cost``. BYOK responses put the
|
||||
OpenRouter fee in ``usage.cost`` (often 0) and the provider charge in
|
||||
``usage.cost_details.upstream_inference_cost``, so BYOK totals sum the two.
|
||||
"""
|
||||
if not isinstance(usage, dict):
|
||||
return None
|
||||
total = 0.0
|
||||
cost = usage.get("cost")
|
||||
if isinstance(cost, int | float) and cost > 0:
|
||||
total += float(cost)
|
||||
if bool(usage.get("is_byok")):
|
||||
details = usage.get("cost_details")
|
||||
upstream = details.get("upstream_inference_cost") if isinstance(details, dict) else None
|
||||
if isinstance(upstream, int | float) and upstream > 0:
|
||||
total += float(upstream)
|
||||
return total if total > 0 else None
|
||||
|
||||
|
||||
def _response_id(completion_response: Any) -> str | None:
|
||||
response_id = getattr(completion_response, "id", None)
|
||||
if response_id is None and isinstance(completion_response, dict):
|
||||
response_id = cast("dict[str, Any]", completion_response).get("id")
|
||||
return response_id if isinstance(response_id, str) and response_id else None
|
||||
|
||||
|
||||
class StreamedOpenRouterCosts:
|
||||
"""Correlates OpenRouter's per-stream cost from the parser to the cost callback.
|
||||
|
||||
LiteLLM rebuilds streamed responses from token-only chunks and drops the
|
||||
``usage.cost`` OpenRouter reports in its final stream chunk (its non-streamed
|
||||
path preserves it; streaming snapshots hidden params at stream start). Every
|
||||
scan streams, so the OpenRouter streaming handler (see strix.config.models)
|
||||
records the cost here keyed by response id, and the callback takes it back out
|
||||
for the matching rebuilt response. Entries are removed on read; ``clear()``
|
||||
runs per scan so nothing accumulates across runs.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._costs: dict[str, float] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def remember(self, response_id: Any, usage: Any) -> None:
|
||||
cost = openrouter_stream_cost(usage)
|
||||
if cost is None or not (isinstance(response_id, str) and response_id):
|
||||
return
|
||||
with self._lock:
|
||||
self._costs[response_id] = cost
|
||||
|
||||
def take(self, completion_response: Any) -> float | None:
|
||||
response_id = _response_id(completion_response)
|
||||
if response_id is None:
|
||||
return None
|
||||
with self._lock:
|
||||
return self._costs.pop(response_id, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._costs.clear()
|
||||
|
||||
|
||||
streamed_openrouter_costs = StreamedOpenRouterCosts()
|
||||
|
||||
|
||||
def litellm_cost_callback(
|
||||
kwargs: Any,
|
||||
completion_response: Any,
|
||||
@@ -541,6 +610,11 @@ def litellm_cost_callback(
|
||||
if cost is None:
|
||||
cost = _usage_reported_cost(completion_response)
|
||||
|
||||
# Recover the exact OpenRouter cost the streaming handler stashed for this
|
||||
# response — LiteLLM drops it from streamed usage, so nothing above sees it.
|
||||
if cost is None:
|
||||
cost = streamed_openrouter_costs.take(completion_response)
|
||||
|
||||
if cost is None:
|
||||
cost = _estimate_response_cost(kwargs, completion_response)
|
||||
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
import json
|
||||
import logging
|
||||
import urllib.request
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import requests
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.telemetry._common import (
|
||||
SESSION_ID,
|
||||
@@ -37,13 +37,7 @@ def _send(event: str, properties: dict[str, Any]) -> bool:
|
||||
"distinct_id": SESSION_ID,
|
||||
"properties": properties,
|
||||
}
|
||||
req = urllib.request.Request( # noqa: S310
|
||||
f"{_POSTHOG_HOST}/capture/",
|
||||
data=json.dumps(payload).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=10): # noqa: S310 # nosec B310
|
||||
pass
|
||||
requests.post(f"{_POSTHOG_HOST}/capture/", json=payload, timeout=10)
|
||||
except Exception: # noqa: BLE001
|
||||
logger.debug("posthog send failed for event %s", event, exc_info=True)
|
||||
return False
|
||||
|
||||
@@ -2,10 +2,11 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import requests
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.telemetry._common import (
|
||||
SESSION_ID,
|
||||
@@ -42,9 +43,7 @@ def _send(event: str, properties: dict[str, Any]) -> bool:
|
||||
url = f"{_SCARF_ENDPOINT}{path}"
|
||||
if query:
|
||||
url = f"{url}?{query}"
|
||||
req = urllib.request.Request(url, method="POST") # noqa: S310
|
||||
with urllib.request.urlopen(req, timeout=10): # noqa: S310 # nosec B310
|
||||
pass
|
||||
requests.post(url, timeout=10)
|
||||
except Exception: # noqa: BLE001
|
||||
logger.debug("scarf send failed for event %s", event, exc_info=True)
|
||||
return False
|
||||
|
||||
@@ -7,8 +7,10 @@ import hashlib
|
||||
import json
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from strix.config import codex
|
||||
|
||||
@@ -52,6 +54,18 @@ def test_authorize_url_carries_pkce_and_client() -> None:
|
||||
assert "state=st8" in url
|
||||
|
||||
|
||||
def test_post_form_returns_parsed_body() -> None:
|
||||
resp = mock.MagicMock()
|
||||
resp.status_code = 200
|
||||
resp.content = b'{"access_token": "tok"}'
|
||||
|
||||
with mock.patch.object(requests, "post", return_value=resp) as post:
|
||||
data = codex._post_form({"grant_type": "refresh_token"})
|
||||
|
||||
assert data == {"access_token": "tok"}
|
||||
assert post.call_args.kwargs["timeout"] == codex._TOKEN_TIMEOUT
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[
|
||||
|
||||
+105
-2
@@ -7,9 +7,25 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
from strix.config.models import _configure_litellm_compatibility
|
||||
from strix.report.state import litellm_cost_callback
|
||||
from strix.config.models import (
|
||||
_configure_litellm_compatibility,
|
||||
_install_openrouter_stream_cost_capture,
|
||||
)
|
||||
from strix.report.state import (
|
||||
ReportState,
|
||||
litellm_cost_callback,
|
||||
openrouter_stream_cost,
|
||||
set_global_report_state,
|
||||
streamed_openrouter_costs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_streamed_costs() -> None:
|
||||
streamed_openrouter_costs.clear()
|
||||
|
||||
|
||||
def test_streaming_logging_stays_enabled_for_cost_callback() -> None:
|
||||
@@ -151,3 +167,90 @@ def test_cost_callback_records_nothing_when_no_cost_available() -> None:
|
||||
litellm_cost_callback({"response_cost": None, "model": "x/y"}, response)
|
||||
|
||||
report_state.record_observed_llm_cost.assert_not_called()
|
||||
|
||||
|
||||
def test_openrouter_stream_cost_extracts_plain_and_byok_totals() -> None:
|
||||
assert openrouter_stream_cost({"cost": 0.003168}) == pytest.approx(0.003168)
|
||||
assert openrouter_stream_cost(
|
||||
{"cost": 0.01, "is_byok": True, "cost_details": {"upstream_inference_cost": 0.2}}
|
||||
) == pytest.approx(0.21)
|
||||
# Upstream cost is only added for BYOK responses.
|
||||
assert openrouter_stream_cost(
|
||||
{"cost": 0.05, "is_byok": False, "cost_details": {"upstream_inference_cost": 0.04}}
|
||||
) == pytest.approx(0.05)
|
||||
assert openrouter_stream_cost({"prompt_tokens": 10}) is None
|
||||
assert openrouter_stream_cost(None) is None
|
||||
|
||||
|
||||
def test_cost_callback_recovers_streamed_openrouter_cost_by_response_id() -> None:
|
||||
report_state = MagicMock()
|
||||
streamed_openrouter_costs.remember("gen-abc", {"cost": 0.42})
|
||||
# LiteLLM strips cost from the rebuilt streamed usage; only the id survives.
|
||||
response = SimpleNamespace(id="gen-abc", usage=SimpleNamespace(cost=None), _hidden_params={})
|
||||
|
||||
with (
|
||||
patch("strix.report.state.get_global_report_state", return_value=report_state),
|
||||
patch("litellm.completion_cost", side_effect=ValueError("unknown model")),
|
||||
):
|
||||
litellm_cost_callback({"response_cost": None, "model": "moonshotai/kimi-k3"}, response)
|
||||
|
||||
report_state.record_observed_llm_cost.assert_called_once_with(0.42)
|
||||
# The entry is consumed so a later response cannot double-count it.
|
||||
assert streamed_openrouter_costs.take(response) is None
|
||||
|
||||
|
||||
def test_streamed_openrouter_cost_prefers_provider_report_over_estimate() -> None:
|
||||
report_state = MagicMock()
|
||||
streamed_openrouter_costs.remember("gen-xyz", {"cost": 0.9})
|
||||
response = SimpleNamespace(
|
||||
id="gen-xyz",
|
||||
usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
_hidden_params={},
|
||||
)
|
||||
|
||||
with (
|
||||
patch("strix.report.state.get_global_report_state", return_value=report_state),
|
||||
patch("litellm.completion_cost", return_value=0.1) as estimate,
|
||||
):
|
||||
litellm_cost_callback({"response_cost": None, "model": "moonshotai/kimi-k3"}, response)
|
||||
|
||||
report_state.record_observed_llm_cost.assert_called_once_with(0.9)
|
||||
estimate.assert_not_called()
|
||||
|
||||
|
||||
def test_streamed_openrouter_costs_ignores_entries_without_cost() -> None:
|
||||
streamed_openrouter_costs.remember("gen-none", {"prompt_tokens": 10})
|
||||
streamed_openrouter_costs.remember("", {"cost": 0.5})
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-none")) is None
|
||||
|
||||
|
||||
def test_streamed_openrouter_costs_cleared_on_new_run() -> None:
|
||||
streamed_openrouter_costs.remember("gen-stale", {"cost": 0.7})
|
||||
set_global_report_state(ReportState.__new__(ReportState))
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-stale")) is None
|
||||
|
||||
|
||||
def test_openrouter_stream_handler_records_cost() -> None:
|
||||
_install_openrouter_stream_cost_capture()
|
||||
# Resolve the config the way LiteLLM does in production so we prove the
|
||||
# override is actually reachable through provider resolution, not just as a
|
||||
# directly-constructed class.
|
||||
config = ProviderConfigManager.get_provider_chat_config(
|
||||
model="moonshotai/kimi-k3", provider=LlmProviders.OPENROUTER
|
||||
)
|
||||
assert config is not None
|
||||
assert type(config).__name__ == "_StrixOpenrouterConfig"
|
||||
handler = config.get_model_response_iterator(streaming_response=iter([]), sync_stream=True)
|
||||
|
||||
chunk = {
|
||||
"id": "gen-stream",
|
||||
"created": 1,
|
||||
"model": "moonshotai/kimi-k3",
|
||||
"choices": [{"index": 0, "delta": {"content": None}}],
|
||||
"usage": {"prompt_tokens": 89, "completion_tokens": 138, "cost": 0.0035055},
|
||||
}
|
||||
handler.chunk_parser(chunk)
|
||||
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-stream")) == pytest.approx(
|
||||
0.0035055
|
||||
)
|
||||
|
||||
+112
-1
@@ -6,13 +6,21 @@ import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from agents.memory import SQLiteSession
|
||||
from agents.tool_context import ToolContext
|
||||
|
||||
from strix.config import codex
|
||||
from strix.core import execution
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.core.execution import _notify_parent_on_terminal, _notify_root_on_budget_reserve
|
||||
from strix.core.execution import (
|
||||
_handle_content_guardrail,
|
||||
_notify_parent_on_terminal,
|
||||
_notify_root_on_budget_reserve,
|
||||
respawn_subagents,
|
||||
)
|
||||
from strix.tools.finish.tool import finish_scan
|
||||
|
||||
|
||||
@@ -465,3 +473,106 @@ async def test_notify_parent_on_terminal_ignores_non_terminal_status(tmp_path: A
|
||||
|
||||
assert coordinator.pending_counts.get("root", 0) == 0
|
||||
session.close()
|
||||
|
||||
|
||||
class _RecordingStream:
|
||||
def __init__(self) -> None:
|
||||
self.cancelled = False
|
||||
self.cancel_mode: str | None = None
|
||||
|
||||
def cancel(self, mode: str = "immediate") -> None:
|
||||
self.cancelled = True
|
||||
self.cancel_mode = mode
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_notice_does_not_cancel_parent_stream(tmp_path: Any) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
session = SQLiteSession("root", tmp_path / "agents.db")
|
||||
stream = _RecordingStream()
|
||||
await coordinator.attach_runtime("root", session=session, interrupt_on_message=True)
|
||||
await coordinator.attach_stream("root", stream)
|
||||
|
||||
await _notify_parent_on_terminal(coordinator, "child", "crashed")
|
||||
|
||||
assert stream.cancelled is False
|
||||
assert coordinator.pending_counts.get("root", 0) > 0
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_interactive_parks_agent_wakeable(tmp_path: Any) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
exc = codex.CodexContentGuardrailError("chatgpt/gpt-5.6-sol")
|
||||
|
||||
result = await _handle_content_guardrail(coordinator, "child", exc, interactive=True)
|
||||
|
||||
assert result is None
|
||||
assert coordinator.statuses["child"] == "waiting"
|
||||
assert "STRIX_LLM" in coordinator.errors["child"]
|
||||
|
||||
waiter = asyncio.create_task(coordinator.wait_for_message("child"))
|
||||
await asyncio.sleep(0)
|
||||
assert not waiter.done()
|
||||
session = SQLiteSession("child", tmp_path / "agents.db")
|
||||
await coordinator.attach_runtime("child", session=session)
|
||||
await coordinator.send("child", {"from": "user", "content": "switched model, resume"})
|
||||
await asyncio.wait_for(waiter, timeout=1.0)
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_noninteractive_fails_only_blocked_agent(tmp_path: Any) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
session = SQLiteSession("root", tmp_path / "agents.db")
|
||||
await coordinator.attach_runtime("root", session=session)
|
||||
exc = codex.CodexContentGuardrailError("chatgpt/gpt-5.6-sol")
|
||||
|
||||
result = await _handle_content_guardrail(coordinator, "child", exc, interactive=False)
|
||||
|
||||
assert result is None
|
||||
assert coordinator.statuses["child"] == "failed"
|
||||
assert "STRIX_LLM" in coordinator.errors["child"]
|
||||
assert coordinator.statuses["root"] == "running"
|
||||
assert coordinator.pending_counts.get("root", 0) > 0
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_revives_guardrail_parked_child_but_not_plain_waiting(
|
||||
tmp_path: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("blocked", "recon", parent_id="root")
|
||||
await coordinator.register("peer_waiter", "recon", parent_id="root")
|
||||
await coordinator.set_status("blocked", "waiting", error="STRIX_LLM guardrail")
|
||||
await coordinator.set_status("peer_waiter", "waiting")
|
||||
|
||||
parked: dict[str, bool] = {}
|
||||
|
||||
async def _fake_start_child_runner(**kwargs: Any) -> None:
|
||||
parked[kwargs["child_id"]] = bool(kwargs["start_parked"])
|
||||
|
||||
monkeypatch.setattr(execution, "_start_child_runner", _fake_start_child_runner)
|
||||
|
||||
await respawn_subagents(
|
||||
coordinator=coordinator,
|
||||
factory=lambda **_kwargs: object(),
|
||||
agents_db_path=tmp_path / "agents.db",
|
||||
sessions_to_close=[],
|
||||
run_config=MagicMock(),
|
||||
max_turns=10,
|
||||
interactive=True,
|
||||
parent_ctx={"agent_id": "root", "parent_id": None},
|
||||
root_id="root",
|
||||
)
|
||||
|
||||
assert parked["blocked"] is False
|
||||
assert parked["peer_waiter"] is True
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
"""Tests for the non-streaming wrapper used on custom OpenAI-compatible endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import TYPE_CHECKING, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from agents.items import ModelResponse
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import Model, ModelTracing
|
||||
from agents.usage import Usage
|
||||
from openai.types.responses import (
|
||||
ResponseCompletedEvent,
|
||||
ResponseFunctionToolCall,
|
||||
ResponseOutputItemDoneEvent,
|
||||
ResponseStreamEvent,
|
||||
)
|
||||
|
||||
from strix.config.models import StrixProvider, _NonStreamingModel, _to_response_usage
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
|
||||
def _tool_call() -> ResponseFunctionToolCall:
|
||||
return ResponseFunctionToolCall(
|
||||
arguments='{"command": "ls"}',
|
||||
call_id="call_1",
|
||||
name="terminal_execute",
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
|
||||
class _FakeModel(Model):
|
||||
def __init__(self, response: ModelResponse) -> None:
|
||||
self.model = "fake-model"
|
||||
self._response = response
|
||||
self.get_response_calls = 0
|
||||
|
||||
async def get_response(self, *_args: object, **_kwargs: object) -> ModelResponse:
|
||||
self.get_response_calls += 1
|
||||
return self._response
|
||||
|
||||
async def stream_response( # pragma: no cover
|
||||
self, *_args: object, **_kwargs: object
|
||||
) -> AsyncIterator[ResponseStreamEvent]:
|
||||
for _ in range(0):
|
||||
yield cast("ResponseStreamEvent", None)
|
||||
raise AssertionError("inner stream_response must never be called")
|
||||
|
||||
|
||||
def _settings(*, api_base: str | None, stream_mode: str = "auto") -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
llm=SimpleNamespace(
|
||||
model="openai/glm",
|
||||
api_base=api_base,
|
||||
stream_mode=stream_mode,
|
||||
reasoning_effort="high",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_to_response_usage_maps_token_details() -> None:
|
||||
usage = Usage(requests=1, input_tokens=10, output_tokens=5, total_tokens=15)
|
||||
usage.input_tokens_details.cached_tokens = 4
|
||||
usage.output_tokens_details.reasoning_tokens = 3
|
||||
mapped = _to_response_usage(usage)
|
||||
assert mapped is not None
|
||||
assert (mapped.input_tokens, mapped.output_tokens, mapped.total_tokens) == (10, 5, 15)
|
||||
assert mapped.input_tokens_details.cached_tokens == 4
|
||||
assert mapped.output_tokens_details.reasoning_tokens == 3
|
||||
|
||||
|
||||
def test_to_response_usage_none() -> None:
|
||||
assert _to_response_usage(None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_response_synthesizes_tool_call_from_non_streamed() -> None:
|
||||
tool_call = _tool_call()
|
||||
inner = _FakeModel(
|
||||
ModelResponse(
|
||||
output=[tool_call],
|
||||
usage=Usage(requests=1, input_tokens=10, output_tokens=5, total_tokens=15),
|
||||
response_id="resp_123",
|
||||
)
|
||||
)
|
||||
wrapper = _NonStreamingModel(inner)
|
||||
|
||||
events = [
|
||||
event
|
||||
async for event in wrapper.stream_response(
|
||||
"sys",
|
||||
"hi",
|
||||
ModelSettings(),
|
||||
[],
|
||||
None,
|
||||
[],
|
||||
ModelTracing.DISABLED,
|
||||
)
|
||||
]
|
||||
|
||||
assert inner.get_response_calls == 1
|
||||
item_done = [e for e in events if isinstance(e, ResponseOutputItemDoneEvent)]
|
||||
completed = [e for e in events if isinstance(e, ResponseCompletedEvent)]
|
||||
assert len(item_done) == 1
|
||||
assert item_done[0].item == tool_call
|
||||
assert len(completed) == 1
|
||||
final = completed[0].response
|
||||
assert final.output == [tool_call]
|
||||
assert final.id == "resp_123"
|
||||
assert final.usage is not None
|
||||
assert final.usage.total_tokens == 15
|
||||
# sequence numbers are strictly increasing
|
||||
assert [e.sequence_number for e in events] == list(range(len(events)))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_response_delegates_get_response() -> None:
|
||||
inner = _FakeModel(
|
||||
ModelResponse(output=[], usage=Usage(), response_id=None),
|
||||
)
|
||||
wrapper = _NonStreamingModel(inner)
|
||||
result = await wrapper.get_response(
|
||||
"sys", "hi", ModelSettings(), [], None, [], ModelTracing.DISABLED
|
||||
)
|
||||
assert result is inner._response
|
||||
assert inner.get_response_calls == 1
|
||||
|
||||
|
||||
def test_get_model_auto_streams_custom_endpoint() -> None:
|
||||
sentinel = _FakeModel(ModelResponse(output=[], usage=Usage(), response_id=None))
|
||||
with (
|
||||
patch("strix.config.models.load_settings", return_value=_settings(api_base="http://x/v1")),
|
||||
patch(
|
||||
"agents.models.multi_provider.MultiProvider.get_model",
|
||||
return_value=sentinel,
|
||||
),
|
||||
):
|
||||
model = StrixProvider().get_model("openai/glm")
|
||||
assert model is sentinel
|
||||
|
||||
|
||||
def test_get_model_auto_streams_hosted() -> None:
|
||||
sentinel = _FakeModel(ModelResponse(output=[], usage=Usage(), response_id=None))
|
||||
with (
|
||||
patch("strix.config.models.load_settings", return_value=_settings(api_base="")),
|
||||
patch(
|
||||
"agents.models.multi_provider.MultiProvider.get_model",
|
||||
return_value=sentinel,
|
||||
),
|
||||
):
|
||||
model = StrixProvider().get_model("openai/gpt-4o")
|
||||
assert model is sentinel
|
||||
|
||||
|
||||
def test_get_model_stream_mode_never_wraps() -> None:
|
||||
sentinel = _FakeModel(ModelResponse(output=[], usage=Usage(), response_id=None))
|
||||
with (
|
||||
patch(
|
||||
"strix.config.models.load_settings",
|
||||
return_value=_settings(api_base="http://x/v1", stream_mode="never"),
|
||||
),
|
||||
patch(
|
||||
"agents.models.multi_provider.MultiProvider.get_model",
|
||||
return_value=sentinel,
|
||||
),
|
||||
):
|
||||
model = StrixProvider().get_model("openai/glm")
|
||||
assert isinstance(model, _NonStreamingModel)
|
||||
+125
-2
@@ -4,9 +4,11 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from strix.core.paths import latest_run_dir, runs_base_dir
|
||||
from strix.interface.viewer.server import serve
|
||||
@@ -86,6 +88,75 @@ def test_build_run_state_from_agents_json(tmp_path: Path) -> None:
|
||||
assert state["events"] == []
|
||||
|
||||
|
||||
def test_build_run_state_keeps_same_call_id_separate_per_agent(tmp_path: Path) -> None:
|
||||
run_dir = _make_run(tmp_path, "tools", status="completed", end_time=None)
|
||||
agents_db = run_dir / ".state" / "agents.db"
|
||||
rows = [
|
||||
(
|
||||
"root",
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "exec_command_0",
|
||||
"name": "exec_command",
|
||||
"arguments": json.dumps({"cmd": "echo root"}),
|
||||
},
|
||||
),
|
||||
(
|
||||
"root",
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "exec_command_0",
|
||||
"output": json.dumps({"success": True, "output": "root"}),
|
||||
},
|
||||
),
|
||||
(
|
||||
"child",
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "exec_command_0",
|
||||
"name": "exec_command",
|
||||
"arguments": json.dumps({"cmd": "echo child"}),
|
||||
},
|
||||
),
|
||||
(
|
||||
"child",
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "exec_command_0",
|
||||
"output": json.dumps({"success": True, "output": "child"}),
|
||||
},
|
||||
),
|
||||
]
|
||||
with sqlite3.connect(agents_db) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
create table agent_messages (
|
||||
id integer primary key,
|
||||
session_id text not null,
|
||||
message_data text not null,
|
||||
created_at text not null
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.executemany(
|
||||
"""
|
||||
insert into agent_messages (session_id, message_data, created_at)
|
||||
values (?, ?, '2026-01-01T00:00:00+00:00')
|
||||
""",
|
||||
[(agent_id, json.dumps(message)) for agent_id, message in rows],
|
||||
)
|
||||
|
||||
state = build_run_state(run_dir)
|
||||
tools = [event for event in state["events"] if event["type"] == "tool"]
|
||||
|
||||
assert len(tools) == 2
|
||||
by_agent = {event["agent_id"]: event for event in tools}
|
||||
assert by_agent["root"]["data"]["args"] == {"cmd": "echo root"}
|
||||
assert by_agent["root"]["data"]["result"]["output"] == "root"
|
||||
assert by_agent["child"]["data"]["args"] == {"cmd": "echo child"}
|
||||
assert by_agent["child"]["data"]["result"]["output"] == "child"
|
||||
|
||||
|
||||
def _get(url: str, *, cookie: str | None = None) -> tuple[int, str, bytes]:
|
||||
headers = {"Cookie": cookie} if cookie else {}
|
||||
req = urllib.request.Request(url, headers=headers) # noqa: S310 - localhost test server
|
||||
@@ -271,6 +342,11 @@ def _session_cookie(url: str, token: str) -> str:
|
||||
return raw.split(";", 1)[0]
|
||||
|
||||
|
||||
def _cookie_name(url: str) -> str:
|
||||
"""The per-server session cookie name, derived from the bound port."""
|
||||
return f"strix_viewer_session_{urlsplit(url).port}"
|
||||
|
||||
|
||||
def _get_status(url: str, *, cookie: str | None = None) -> int:
|
||||
headers = {"Cookie": cookie} if cookie else {}
|
||||
req = urllib.request.Request(url, headers=headers) # noqa: S310 - localhost test server
|
||||
@@ -311,7 +387,7 @@ def test_capability_issued_only_for_tokened_bootstrap(
|
||||
# Only the correct bootstrap token mints the session cookie.
|
||||
with urllib.request.urlopen(f"{url}/?token={token}") as resp: # noqa: S310 # nosec B310
|
||||
cookie = str(resp.headers.get("Set-Cookie", ""))
|
||||
assert "strix_viewer_session=" in cookie
|
||||
assert f"{_cookie_name(url)}=" in cookie
|
||||
assert "HttpOnly" in cookie and "SameSite=Strict" in cookie
|
||||
|
||||
# Static assets never carry it.
|
||||
@@ -344,7 +420,7 @@ def test_unauthorized_client_cannot_acquire_capability(
|
||||
url,
|
||||
"/api/agents/steer",
|
||||
{"agent_id": "root", "message": "pwn"},
|
||||
cookie="strix_viewer_session=",
|
||||
cookie=f"{_cookie_name(url)}=",
|
||||
)
|
||||
assert status == 403
|
||||
assert delivered == []
|
||||
@@ -541,6 +617,53 @@ def test_runs_list_requires_session_and_verification(
|
||||
httpd.server_close()
|
||||
|
||||
|
||||
def test_concurrent_servers_use_distinct_cookies(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Cookies are host-scoped, not port-scoped: two viewers on 127.0.0.1 must
|
||||
not share a cookie slot, and one server's cookie must not pass the other's
|
||||
session gate."""
|
||||
run_a = _make_run(tmp_path / "a", "run-a", status="running", end_time=None)
|
||||
run_b = _make_run(tmp_path / "b", "run-b", status="running", end_time=None)
|
||||
_bundle(tmp_path, monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
"strix.interface.viewer.auth.read_auth", lambda: {"email": "a@b.com", "token": "t"}
|
||||
)
|
||||
monkeypatch.setattr("strix.interface.viewer.auth.is_verified", lambda: True)
|
||||
|
||||
httpd_a, url_a, token_a = serve(run_a, open_browser=False)
|
||||
httpd_b, url_b, token_b = serve(run_b, open_browser=False)
|
||||
try:
|
||||
cookie_a = _session_cookie(url_a, token_a)
|
||||
cookie_b = _session_cookie(url_b, token_b)
|
||||
|
||||
# The two servers mint differently named cookies, so a browser stores both.
|
||||
assert cookie_a.split("=", 1)[0] == _cookie_name(url_a)
|
||||
assert cookie_b.split("=", 1)[0] == _cookie_name(url_b)
|
||||
assert cookie_a.split("=", 1)[0] != cookie_b.split("=", 1)[0]
|
||||
|
||||
def _status(url: str, cookie: str) -> dict[str, object]:
|
||||
_, _, body = _get(f"{url}/api/auth/status", cookie=cookie)
|
||||
return dict(json.loads(body))
|
||||
|
||||
# Each server honors its own cookie...
|
||||
assert _status(url_a, cookie_a)["verified"] is True
|
||||
assert _status(url_b, cookie_b)["verified"] is True
|
||||
# ...but treats the other server's cookie as session-less.
|
||||
assert _status(url_a, cookie_b)["verified"] is False
|
||||
assert _status(url_b, cookie_a)["verified"] is False
|
||||
# Even both cookies together (what a real browser would send) only
|
||||
# match the token minted by the receiving server.
|
||||
both = f"{cookie_a}; {cookie_b}"
|
||||
assert _status(url_a, both)["verified"] is True
|
||||
assert _status(url_b, both)["verified"] is True
|
||||
finally:
|
||||
httpd_a.shutdown()
|
||||
httpd_a.server_close()
|
||||
httpd_b.shutdown()
|
||||
httpd_b.server_close()
|
||||
|
||||
|
||||
def test_server_rejects_path_traversal(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
run_dir = _make_run(tmp_path, "guard", status="completed", end_time="2026-01-01T00:00:00Z")
|
||||
secret = tmp_path / "secret.txt"
|
||||
|
||||
Reference in New Issue
Block a user