mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-30 01:19:34 +00:00
Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
71d5c49b82 | ||
|
|
1d888d42a9 | ||
|
|
786df99dc7 | ||
|
|
9cffb1d6b7 | ||
|
|
6c1d0a4828 | ||
|
|
b0d2d3cff6 | ||
|
|
d72d6f97a6 | ||
|
|
cb059b59e2 | ||
|
|
e8552cd558 | ||
|
|
3ad91338f7 | ||
|
|
348a445a37 | ||
|
|
2c258213f5 | ||
|
|
ae4fee1165 | ||
|
|
cd23380a66 | ||
|
|
ac8083bd85 | ||
|
|
42e6650ab8 |
+1
-18
@@ -4,7 +4,7 @@ DATABASE_POOL_SIZE=10
|
|||||||
|
|
||||||
# LLM Provider
|
# LLM Provider
|
||||||
# LLM_BACKEND=nearai # default
|
# LLM_BACKEND=nearai # default
|
||||||
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, github_copilot, tinfoil, openai_codex, gemini_oauth
|
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, github_copilot, tinfoil, openai_codex
|
||||||
# LLM_REQUEST_TIMEOUT_SECS=120 # Increase for local LLMs (Ollama, vLLM, LM Studio)
|
# LLM_REQUEST_TIMEOUT_SECS=120 # Increase for local LLMs (Ollama, vLLM, LM Studio)
|
||||||
|
|
||||||
# === Anthropic Direct ===
|
# === Anthropic Direct ===
|
||||||
@@ -110,23 +110,6 @@ NEARAI_AUTH_URL=https://private.near.ai
|
|||||||
# OPENAI_CODEX_AUTH_URL=https://auth.openai.com # override (rare)
|
# OPENAI_CODEX_AUTH_URL=https://auth.openai.com # override (rare)
|
||||||
# OPENAI_CODEX_API_URL=https://chatgpt.com/backend-api/codex # override (rare)
|
# OPENAI_CODEX_API_URL=https://chatgpt.com/backend-api/codex # override (rare)
|
||||||
|
|
||||||
# === Google Gemini (OAuth, Gemini CLI compatible) ===
|
|
||||||
# LLM_BACKEND=gemini_oauth
|
|
||||||
# GEMINI_MODEL=gemini-2.5-flash # default
|
|
||||||
# GEMINI_CREDENTIALS_PATH=~/.gemini/oauth_creds.json # default
|
|
||||||
# GEMINI_API_KEY=... # optional: use API key instead of OAuth
|
|
||||||
# GEMINI_API_KEY_AUTH_MECHANISM=query # "query" (default) or "header"
|
|
||||||
# GEMINI_SAFETY_BLOCK_NONE=true # disable safety filters (default: false)
|
|
||||||
# GEMINI_CLI_CUSTOM_HEADERS=Key:Value,Key2:Value2
|
|
||||||
# GEMINI_TOP_P=0.95
|
|
||||||
# GEMINI_TOP_K=40
|
|
||||||
# GEMINI_SEED=42
|
|
||||||
# GEMINI_PRESENCE_PENALTY=0.0
|
|
||||||
# GEMINI_FREQUENCY_PENALTY=0.0
|
|
||||||
# GEMINI_RESPONSE_MIME_TYPE=application/json
|
|
||||||
# GEMINI_RESPONSE_JSON_SCHEMA={"type":"object"}
|
|
||||||
# GEMINI_CACHED_CONTENT=cachedContents/abc123
|
|
||||||
|
|
||||||
# For full provider setup guide see docs/LLM_PROVIDERS.md
|
# For full provider setup guide see docs/LLM_PROVIDERS.md
|
||||||
|
|
||||||
# Channel Configuration
|
# Channel Configuration
|
||||||
|
|||||||
@@ -55,6 +55,9 @@ RUN npm install -g @anthropic-ai/claude-code@latest
|
|||||||
# Copy the binary
|
# Copy the binary
|
||||||
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
|
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
|
||||||
|
|
||||||
|
# Install IronClaw Python SDK for programmatic tool calling (PTC)
|
||||||
|
COPY sdk/python/ironclaw_tools.py /usr/lib/python3/dist-packages/ironclaw_tools.py
|
||||||
|
|
||||||
# Create non-root user (UID 1000 matches the orchestrator's container config)
|
# Create non-root user (UID 1000 matches the orchestrator's container config)
|
||||||
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
||||||
&& mkdir -p /workspace \
|
&& mkdir -p /workspace \
|
||||||
|
|||||||
+5
-14
@@ -3,7 +3,6 @@
|
|||||||
This document tracks feature parity between IronClaw (Rust implementation) and OpenClaw (TypeScript reference implementation). Use this to coordinate work across developers.
|
This document tracks feature parity between IronClaw (Rust implementation) and OpenClaw (TypeScript reference implementation). Use this to coordinate work across developers.
|
||||||
|
|
||||||
**Legend:**
|
**Legend:**
|
||||||
|
|
||||||
- ✅ Implemented
|
- ✅ Implemented
|
||||||
- 🚧 Partial (in progress or incomplete)
|
- 🚧 Partial (in progress or incomplete)
|
||||||
- ❌ Not implemented
|
- ❌ Not implemented
|
||||||
@@ -205,7 +204,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
| Skills (modular capabilities) | ✅ | ✅ | Prompt-based skills with trust gating, attenuation, activation criteria, catalog, selector |
|
| Skills (modular capabilities) | ✅ | ✅ | Prompt-based skills with trust gating, attenuation, activation criteria, catalog, selector |
|
||||||
| Skill routing blocks | ✅ | 🚧 | ActivationCriteria (keywords, patterns, tags) but no "Use when / Don't use when" blocks |
|
| Skill routing blocks | ✅ | 🚧 | ActivationCriteria (keywords, patterns, tags) but no "Use when / Don't use when" blocks |
|
||||||
| Skill path compaction | ✅ | ❌ | ~ prefix to reduce prompt tokens |
|
| Skill path compaction | ✅ | ❌ | ~ prefix to reduce prompt tokens |
|
||||||
| Thinking modes (off/minimal/low/medium/high/xhigh/adaptive) | ✅ | 🚧 | thinkingConfig for Gemini models (thinkingBudget/thinkingLevel); no per-level control yet |
|
| Thinking modes (off/minimal/low/medium/high/xhigh/adaptive) | ✅ | ❌ | Configurable reasoning depth |
|
||||||
| Per-model thinkingDefault override | ✅ | ❌ | Override thinking level per model; Anthropic Claude 4.6 defaults to adaptive |
|
| Per-model thinkingDefault override | ✅ | ❌ | Override thinking level per model; Anthropic Claude 4.6 defaults to adaptive |
|
||||||
| Block-level streaming | ✅ | ❌ | |
|
| Block-level streaming | ✅ | ❌ | |
|
||||||
| Tool-level streaming | ✅ | ❌ | |
|
| Tool-level streaming | ✅ | ❌ | |
|
||||||
@@ -237,13 +236,9 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
| NEAR AI | ✅ | ✅ | - | Primary provider |
|
| NEAR AI | ✅ | ✅ | - | Primary provider |
|
||||||
| Anthropic (Claude) | ✅ | 🚧 | - | Via NEAR AI proxy; Opus 4.5, Sonnet 4, Sonnet 4.6, adaptive thinking default |
|
| Anthropic (Claude) | ✅ | 🚧 | - | Via NEAR AI proxy; Opus 4.5, Sonnet 4, Sonnet 4.6, adaptive thinking default |
|
||||||
| OpenAI | ✅ | 🚧 | - | Via NEAR AI proxy; GPT-5.4 + Codex OAuth |
|
| OpenAI | ✅ | 🚧 | - | Via NEAR AI proxy; GPT-5.4 + Codex OAuth |
|
||||||
| AWS Bedrock | ✅ | ✅ | - | Native Converse API via aws-sdk-bedrockruntime (requires `--features bedrock`) |
|
| AWS Bedrock | ✅ | ❌ | P3 | |
|
||||||
| Google Gemini | ✅ | ✅ | - | OAuth (PKCE + S256), function calling, thinkingConfig, generationConfig |
|
| Google Gemini | ✅ | ❌ | P3 | |
|
||||||
| io.net | ✅ | ✅ | P3 | Via `ionet` adapter |
|
| NVIDIA API | ✅ | ❌ | P3 | New provider |
|
||||||
| Mistral | ✅ | ✅ | P3 | Via `mistral` adapter |
|
|
||||||
| Yandex AI Studio | ✅ | ✅ | P3 | Via `yandex` adapter |
|
|
||||||
| Cloudflare Workers AI | ✅ | ✅ | P3 | Via `cloudflare` adapter |
|
|
||||||
| NVIDIA API | ✅ | ✅ | P3 | Via `nvidia` adapter and `providers.json` |
|
|
||||||
| OpenRouter | ✅ | ✅ | - | Via OpenAI-compatible provider (RigAdapter) |
|
| OpenRouter | ✅ | ✅ | - | Via OpenAI-compatible provider (RigAdapter) |
|
||||||
| Tinfoil | ❌ | ✅ | - | Private inference provider (IronClaw-only) |
|
| Tinfoil | ❌ | ✅ | - | Private inference provider (IronClaw-only) |
|
||||||
| OpenAI-compatible | ❌ | ✅ | - | Generic OpenAI-compatible endpoint (RigAdapter) |
|
| OpenAI-compatible | ❌ | ✅ | - | Generic OpenAI-compatible endpoint (RigAdapter) |
|
||||||
@@ -471,7 +466,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
| Device pairing | ✅ | ❌ | |
|
| Device pairing | ✅ | ❌ | |
|
||||||
| Tailscale identity | ✅ | ❌ | |
|
| Tailscale identity | ✅ | ❌ | |
|
||||||
| Trusted-proxy auth | ✅ | ❌ | Header-based reverse proxy auth |
|
| Trusted-proxy auth | ✅ | ❌ | Header-based reverse proxy auth |
|
||||||
| OAuth flows | ✅ | 🚧 | NEAR AI OAuth + Gemini OAuth (PKCE, S256) + hosted extension/MCP OAuth broker; external auth-proxy rollout still pending |
|
| OAuth flows | ✅ | 🚧 | NEAR AI OAuth plus hosted extension/MCP OAuth broker; external auth-proxy rollout still pending |
|
||||||
| DM pairing verification | ✅ | ✅ | ironclaw pairing approve, host APIs |
|
| DM pairing verification | ✅ | ✅ | ironclaw pairing approve, host APIs |
|
||||||
| Allowlist/blocklist | ✅ | 🚧 | allow_from + pairing store |
|
| Allowlist/blocklist | ✅ | 🚧 | allow_from + pairing store |
|
||||||
| Per-group tool policies | ✅ | ❌ | |
|
| Per-group tool policies | ✅ | ❌ | |
|
||||||
@@ -528,7 +523,6 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
## Implementation Priorities
|
## Implementation Priorities
|
||||||
|
|
||||||
### P0 - Core (Already Done)
|
### P0 - Core (Already Done)
|
||||||
|
|
||||||
- ✅ TUI channel with approval overlays
|
- ✅ TUI channel with approval overlays
|
||||||
- ✅ HTTP webhook channel
|
- ✅ HTTP webhook channel
|
||||||
- ✅ DM pairing (ironclaw pairing list/approve, host APIs)
|
- ✅ DM pairing (ironclaw pairing list/approve, host APIs)
|
||||||
@@ -556,7 +550,6 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
- ✅ OpenAI-compatible / OpenRouter provider support
|
- ✅ OpenAI-compatible / OpenRouter provider support
|
||||||
|
|
||||||
### P1 - High Priority
|
### P1 - High Priority
|
||||||
|
|
||||||
- ❌ Slack channel (real implementation)
|
- ❌ Slack channel (real implementation)
|
||||||
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
|
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
|
||||||
- ❌ WhatsApp channel
|
- ❌ WhatsApp channel
|
||||||
@@ -564,7 +557,6 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
- ✅ Hooks system (core lifecycle hooks + bundled/plugin/workspace hooks + outbound webhooks)
|
- ✅ Hooks system (core lifecycle hooks + bundled/plugin/workspace hooks + outbound webhooks)
|
||||||
|
|
||||||
### P2 - Medium Priority
|
### P2 - Medium Priority
|
||||||
|
|
||||||
- ❌ Media handling (images, PDFs)
|
- ❌ Media handling (images, PDFs)
|
||||||
- ✅ Ollama/local model support (via rig::providers::ollama)
|
- ✅ Ollama/local model support (via rig::providers::ollama)
|
||||||
- ❌ Configuration hot-reload
|
- ❌ Configuration hot-reload
|
||||||
@@ -573,7 +565,6 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
|||||||
- ❌ Partial output preservation on abort
|
- ❌ Partial output preservation on abort
|
||||||
|
|
||||||
### P3 - Lower Priority
|
### P3 - Lower Priority
|
||||||
|
|
||||||
- ❌ Discord channel
|
- ❌ Discord channel
|
||||||
- ❌ Matrix channel
|
- ❌ Matrix channel
|
||||||
- ❌ Other messaging platforms
|
- ❌ Other messaging platforms
|
||||||
|
|||||||
+3
-48
@@ -1,8 +1,8 @@
|
|||||||
# LLM Provider Configuration
|
# LLM Provider Configuration
|
||||||
|
|
||||||
IronClaw defaults to NEAR AI for model access, but supports any OpenAI-compatible
|
IronClaw defaults to NEAR AI for model access, but supports any OpenAI-compatible
|
||||||
endpoint as well as Anthropic, Ollama, and Google Gemini directly. This guide covers
|
endpoint as well as Anthropic and Ollama directly. This guide covers the most common
|
||||||
the most common configurations.
|
configurations.
|
||||||
|
|
||||||
## Provider Overview
|
## Provider Overview
|
||||||
|
|
||||||
@@ -11,7 +11,7 @@ the most common configurations.
|
|||||||
| NEAR AI | `nearai` | OAuth (browser) | Default; multi-model |
|
| NEAR AI | `nearai` | OAuth (browser) | Default; multi-model |
|
||||||
| Anthropic | `anthropic` | `ANTHROPIC_API_KEY` | Claude models |
|
| Anthropic | `anthropic` | `ANTHROPIC_API_KEY` | Claude models |
|
||||||
| OpenAI | `openai` | `OPENAI_API_KEY` | GPT models |
|
| OpenAI | `openai` | `OPENAI_API_KEY` | GPT models |
|
||||||
| Google Gemini | `gemini_oauth` | OAuth (browser) | Gemini models; function calling |
|
| Google Gemini | `gemini` | `GEMINI_API_KEY` | Gemini models |
|
||||||
| io.net | `ionet` | `IONET_API_KEY` | Intelligence API |
|
| io.net | `ionet` | `IONET_API_KEY` | Intelligence API |
|
||||||
| Mistral | `mistral` | `MISTRAL_API_KEY` | Mistral models |
|
| Mistral | `mistral` | `MISTRAL_API_KEY` | Mistral models |
|
||||||
| Yandex AI Studio | `yandex` | `YANDEX_API_KEY` | YandexGPT models |
|
| Yandex AI Studio | `yandex` | `YANDEX_API_KEY` | YandexGPT models |
|
||||||
@@ -62,51 +62,6 @@ Popular models: `gpt-4o`, `gpt-4o-mini`, `o3-mini`
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Google Gemini (OAuth)
|
|
||||||
|
|
||||||
Uses Google OAuth with PKCE (S256) for authentication — no API key required.
|
|
||||||
On first run, a browser opens for Google account login. Credentials (including
|
|
||||||
refresh token) are saved to `~/.gemini/oauth_creds.json` with `0600` permissions.
|
|
||||||
|
|
||||||
```env
|
|
||||||
LLM_BACKEND=gemini_oauth
|
|
||||||
GEMINI_MODEL=gemini-2.5-flash
|
|
||||||
```
|
|
||||||
|
|
||||||
### Supported features
|
|
||||||
|
|
||||||
| Feature | Status | Notes |
|
|
||||||
|---|---|---|
|
|
||||||
| Function calling | ✅ | `functionDeclarations` / `functionCall` / `functionResponse` |
|
|
||||||
| `generationConfig` | ✅ | `temperature`, `maxOutputTokens` passed from request |
|
|
||||||
| `thinkingConfig` | ✅ | `thinkingBudget`/`thinkingLevel` for thinking-capable models (does NOT set `includeThoughts`) |
|
|
||||||
| `toolConfig` | ✅ | `functionCallingConfig.mode`: `AUTO`/`ANY`/`NONE` |
|
|
||||||
| SSE streaming | ✅ | Cloud Code API with `streamGenerateContent?alt=sse` |
|
|
||||||
| Token refresh | ✅ | Automatic via refresh token |
|
|
||||||
|
|
||||||
### Popular models
|
|
||||||
|
|
||||||
| Model | ID | Notes |
|
|
||||||
|---|---|---|
|
|
||||||
| Gemini 3.1 Pro | `gemini-3.1-pro-preview` | Latest, strongest reasoning |
|
|
||||||
| Gemini 3.1 Pro Custom Tools | `gemini-3.1-pro-preview-customtools` | Enhanced tool use |
|
|
||||||
| Gemini 3 Pro | `gemini-3-pro-preview` | Preview |
|
|
||||||
| Gemini 3 Flash | `gemini-3-flash-preview` | Fast preview with thinking |
|
|
||||||
| Gemini 3.1 Flash Lite | `gemini-3.1-flash-lite-preview` | Preview, lightweight |
|
|
||||||
| Gemini 2.5 Pro | `gemini-2.5-pro` | Stable, strong reasoning |
|
|
||||||
| Gemini 2.5 Flash | `gemini-2.5-flash` | Fast, good quality |
|
|
||||||
| Gemini 2.5 Flash Lite | `gemini-2.5-flash-lite` | Fastest, lightweight |
|
|
||||||
|
|
||||||
### Cloud Code API vs standard API
|
|
||||||
|
|
||||||
Models containing `-preview` (with hyphen) or `gemini-3` in the name, as well
|
|
||||||
as any `gemini-` model with major version >= 2, route through the Cloud Code
|
|
||||||
API (`cloudcode-pa.googleapis.com`) which supports SSE streaming
|
|
||||||
and project-scoped access. Other models use the standard Generative Language
|
|
||||||
API (`generativelanguage.googleapis.com`).
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## GitHub Copilot
|
## GitHub Copilot
|
||||||
|
|
||||||
GitHub Copilot exposes chat endpoint at
|
GitHub Copilot exposes chat endpoint at
|
||||||
|
|||||||
@@ -0,0 +1,158 @@
|
|||||||
|
"""IronClaw Programmatic Tool Calling SDK for container scripts.
|
||||||
|
|
||||||
|
Thin wrapper using only Python stdlib. Reads connection details from
|
||||||
|
environment variables injected by the orchestrator:
|
||||||
|
|
||||||
|
IRONCLAW_ORCHESTRATOR_URL - Base URL of the orchestrator API
|
||||||
|
IRONCLAW_JOB_ID - UUID of the current job
|
||||||
|
IRONCLAW_WORKER_TOKEN - Bearer token scoped to this job
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
from ironclaw_tools import call_tool, shell, read_file, write_file, http_get
|
||||||
|
|
||||||
|
# Call any registered tool by name
|
||||||
|
result = call_tool("echo", {"message": "hello"})
|
||||||
|
print(result) # "hello"
|
||||||
|
|
||||||
|
# Convenience wrappers
|
||||||
|
output = shell("ls -la")
|
||||||
|
content = read_file("/workspace/README.md")
|
||||||
|
write_file("/workspace/output.txt", "results here")
|
||||||
|
body = http_get("https://api.example.com/data")
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import urllib.request
|
||||||
|
import urllib.error
|
||||||
|
|
||||||
|
|
||||||
|
def _env(name):
|
||||||
|
"""Get a required environment variable."""
|
||||||
|
value = os.environ.get(name)
|
||||||
|
if not value:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Missing required environment variable: {name}. "
|
||||||
|
"This SDK must be run inside an IronClaw container."
|
||||||
|
)
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
def _base_url():
|
||||||
|
"""Build the base URL for tool call requests."""
|
||||||
|
orchestrator = _env("IRONCLAW_ORCHESTRATOR_URL").rstrip("/")
|
||||||
|
job_id = _env("IRONCLAW_JOB_ID")
|
||||||
|
return f"{orchestrator}/worker/{job_id}"
|
||||||
|
|
||||||
|
|
||||||
|
def _token():
|
||||||
|
"""Get the bearer token."""
|
||||||
|
return _env("IRONCLAW_WORKER_TOKEN")
|
||||||
|
|
||||||
|
|
||||||
|
def call_tool(name, params=None, timeout_secs=60):
|
||||||
|
"""Call a tool on the orchestrator by name.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Tool name (e.g., "echo", "shell", "read_file").
|
||||||
|
params: Dictionary of parameters to pass to the tool.
|
||||||
|
timeout_secs: Timeout in seconds (default 60, max 300).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool output as a string.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If the tool call fails.
|
||||||
|
"""
|
||||||
|
url = f"{_base_url()}/tools/call"
|
||||||
|
server_timeout = min(int(timeout_secs), 300)
|
||||||
|
body = {
|
||||||
|
"tool_name": name,
|
||||||
|
"parameters": params or {},
|
||||||
|
"timeout_secs": server_timeout,
|
||||||
|
}
|
||||||
|
|
||||||
|
data = json.dumps(body).encode("utf-8")
|
||||||
|
req = urllib.request.Request(
|
||||||
|
url,
|
||||||
|
data=data,
|
||||||
|
headers={
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
"Authorization": f"Bearer {_token()}",
|
||||||
|
},
|
||||||
|
method="POST",
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Client-side timeout slightly longer than server-side to account
|
||||||
|
# for network latency, preventing premature client timeouts.
|
||||||
|
client_timeout = server_timeout + 5
|
||||||
|
with urllib.request.urlopen(req, timeout=client_timeout) as resp:
|
||||||
|
result = json.loads(resp.read().decode("utf-8"))
|
||||||
|
except urllib.error.HTTPError as e:
|
||||||
|
body_text = e.read().decode("utf-8", errors="replace") if e.fp else ""
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Tool call failed: HTTP {e.code}: {body_text}"
|
||||||
|
) from None
|
||||||
|
except urllib.error.URLError as e:
|
||||||
|
raise RuntimeError(f"Connection to orchestrator failed: {e.reason}") from None
|
||||||
|
|
||||||
|
if not result.get("success"):
|
||||||
|
raise RuntimeError(f"Tool '{name}' failed: {result.get('error', 'unknown error')}")
|
||||||
|
|
||||||
|
return result.get("output", "")
|
||||||
|
|
||||||
|
|
||||||
|
def shell(command, timeout_secs=60):
|
||||||
|
"""Execute a shell command via the orchestrator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
command: Shell command string to execute.
|
||||||
|
timeout_secs: Timeout in seconds (default 60).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Command output as a string.
|
||||||
|
"""
|
||||||
|
return call_tool("shell", {"command": command}, timeout_secs=timeout_secs)
|
||||||
|
|
||||||
|
|
||||||
|
def read_file(path):
|
||||||
|
"""Read a file via the orchestrator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
path: Absolute path to the file.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
File contents as a string.
|
||||||
|
"""
|
||||||
|
return call_tool("read_file", {"path": path})
|
||||||
|
|
||||||
|
|
||||||
|
def write_file(path, content):
|
||||||
|
"""Write a file via the orchestrator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
path: Absolute path to write to.
|
||||||
|
content: String content to write.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Write confirmation message.
|
||||||
|
"""
|
||||||
|
return call_tool("write_file", {"path": path, "content": content})
|
||||||
|
|
||||||
|
|
||||||
|
def http_get(url, headers=None, timeout_secs=30):
|
||||||
|
"""Make an HTTP GET request via the orchestrator's HTTP tool.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
url: URL to fetch.
|
||||||
|
headers: Optional dictionary of headers.
|
||||||
|
timeout_secs: Timeout in seconds (default 30).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Response body as a string.
|
||||||
|
"""
|
||||||
|
params = {"url": url, "method": "GET"}
|
||||||
|
if headers:
|
||||||
|
params["headers"] = headers
|
||||||
|
return call_tool("http", params, timeout_secs=timeout_secs)
|
||||||
@@ -0,0 +1,148 @@
|
|||||||
|
"""Tests for the IronClaw Programmatic Tool Calling Python SDK."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import patch, MagicMock
|
||||||
|
import urllib.error
|
||||||
|
|
||||||
|
# Ensure ironclaw_tools is importable regardless of working directory.
|
||||||
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||||
|
|
||||||
|
|
||||||
|
class TestEnvMissing(unittest.TestCase):
|
||||||
|
"""Test that missing env vars produce clear errors."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
# Clear all relevant env vars
|
||||||
|
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
|
||||||
|
os.environ.pop(var, None)
|
||||||
|
|
||||||
|
def test_env_missing(self):
|
||||||
|
from ironclaw_tools import call_tool
|
||||||
|
with self.assertRaises(RuntimeError) as ctx:
|
||||||
|
call_tool("echo", {"message": "hello"})
|
||||||
|
# Should mention the missing variable
|
||||||
|
self.assertIn("IRONCLAW_ORCHESTRATOR_URL", str(ctx.exception))
|
||||||
|
|
||||||
|
|
||||||
|
class TestCallToolRequestFormat(unittest.TestCase):
|
||||||
|
"""Test that call_tool sends correctly formatted requests."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
os.environ["IRONCLAW_ORCHESTRATOR_URL"] = "http://localhost:50051"
|
||||||
|
os.environ["IRONCLAW_JOB_ID"] = "550e8400-e29b-41d4-a716-446655440000"
|
||||||
|
os.environ["IRONCLAW_WORKER_TOKEN"] = "test-token-123"
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
|
||||||
|
os.environ.pop(var, None)
|
||||||
|
|
||||||
|
@patch("ironclaw_tools.urllib.request.urlopen")
|
||||||
|
def test_call_tool_request_format(self, mock_urlopen):
|
||||||
|
from ironclaw_tools import call_tool
|
||||||
|
|
||||||
|
# Mock successful response
|
||||||
|
mock_response = MagicMock()
|
||||||
|
mock_response.read.return_value = json.dumps({
|
||||||
|
"success": True,
|
||||||
|
"output": "hello",
|
||||||
|
"duration_ms": 5,
|
||||||
|
"was_sanitized": False,
|
||||||
|
}).encode("utf-8")
|
||||||
|
mock_response.__enter__ = lambda s: s
|
||||||
|
mock_response.__exit__ = MagicMock(return_value=False)
|
||||||
|
mock_urlopen.return_value = mock_response
|
||||||
|
|
||||||
|
result = call_tool("echo", {"message": "hello"}, timeout_secs=30)
|
||||||
|
|
||||||
|
# Verify the request was made
|
||||||
|
mock_urlopen.assert_called_once()
|
||||||
|
call_args = mock_urlopen.call_args
|
||||||
|
req = call_args[0][0] # First positional arg is the Request object
|
||||||
|
|
||||||
|
# Check URL
|
||||||
|
self.assertIn("/worker/550e8400-e29b-41d4-a716-446655440000/tools/call", req.full_url)
|
||||||
|
|
||||||
|
# Check headers
|
||||||
|
self.assertEqual(req.get_header("Content-type"), "application/json")
|
||||||
|
self.assertEqual(req.get_header("Authorization"), "Bearer test-token-123")
|
||||||
|
|
||||||
|
# Check body
|
||||||
|
body = json.loads(req.data.decode("utf-8"))
|
||||||
|
self.assertEqual(body["tool_name"], "echo")
|
||||||
|
self.assertEqual(body["parameters"], {"message": "hello"})
|
||||||
|
self.assertEqual(body["timeout_secs"], 30)
|
||||||
|
|
||||||
|
# Check return value
|
||||||
|
self.assertEqual(result, "hello")
|
||||||
|
|
||||||
|
|
||||||
|
class TestCallToolHttpError(unittest.TestCase):
|
||||||
|
"""Test HTTP error handling."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
os.environ["IRONCLAW_ORCHESTRATOR_URL"] = "http://localhost:50051"
|
||||||
|
os.environ["IRONCLAW_JOB_ID"] = "550e8400-e29b-41d4-a716-446655440000"
|
||||||
|
os.environ["IRONCLAW_WORKER_TOKEN"] = "test-token-123"
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
|
||||||
|
os.environ.pop(var, None)
|
||||||
|
|
||||||
|
@patch("ironclaw_tools.urllib.request.urlopen")
|
||||||
|
def test_call_tool_http_error(self, mock_urlopen):
|
||||||
|
from ironclaw_tools import call_tool
|
||||||
|
|
||||||
|
mock_urlopen.side_effect = urllib.error.HTTPError(
|
||||||
|
url="http://localhost:50051/worker/test/tools/call",
|
||||||
|
code=500,
|
||||||
|
msg="Internal Server Error",
|
||||||
|
hdrs=None,
|
||||||
|
fp=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.assertRaises(RuntimeError) as ctx:
|
||||||
|
call_tool("echo", {"message": "hello"})
|
||||||
|
|
||||||
|
self.assertIn("500", str(ctx.exception))
|
||||||
|
|
||||||
|
|
||||||
|
class TestConvenienceWrappers(unittest.TestCase):
|
||||||
|
"""Test that convenience wrappers call call_tool correctly."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
os.environ["IRONCLAW_ORCHESTRATOR_URL"] = "http://localhost:50051"
|
||||||
|
os.environ["IRONCLAW_JOB_ID"] = "550e8400-e29b-41d4-a716-446655440000"
|
||||||
|
os.environ["IRONCLAW_WORKER_TOKEN"] = "test-token-123"
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
|
||||||
|
os.environ.pop(var, None)
|
||||||
|
|
||||||
|
@patch("ironclaw_tools.call_tool")
|
||||||
|
def test_convenience_wrappers(self, mock_call_tool):
|
||||||
|
from ironclaw_tools import shell, read_file, write_file, http_get
|
||||||
|
|
||||||
|
mock_call_tool.return_value = "output"
|
||||||
|
|
||||||
|
# Test shell
|
||||||
|
shell("ls -la")
|
||||||
|
mock_call_tool.assert_called_with("shell", {"command": "ls -la"}, timeout_secs=60)
|
||||||
|
|
||||||
|
# Test read_file
|
||||||
|
read_file("/workspace/README.md")
|
||||||
|
mock_call_tool.assert_called_with("read_file", {"path": "/workspace/README.md"})
|
||||||
|
|
||||||
|
# Test write_file
|
||||||
|
write_file("/workspace/out.txt", "content")
|
||||||
|
mock_call_tool.assert_called_with("write_file", {"path": "/workspace/out.txt", "content": "content"})
|
||||||
|
|
||||||
|
# Test http_get
|
||||||
|
http_get("https://api.example.com/data")
|
||||||
|
mock_call_tool.assert_called_with("http", {"url": "https://api.example.com/data", "method": "GET"}, timeout_secs=30)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
+7
-7
@@ -729,13 +729,13 @@ impl AppBuilder {
|
|||||||
self.init_database().await?;
|
self.init_database().await?;
|
||||||
self.init_secrets().await?;
|
self.init_secrets().await?;
|
||||||
|
|
||||||
// Post-init validation: backends with dedicated config (nearai, gemini_oauth,
|
// Post-init validation: if a non-nearai backend was selected but
|
||||||
// bedrock, openai_codex) handle their own credential resolution. For registry-based
|
// credentials were never resolved (deferred resolution found no keys),
|
||||||
// backends, fail early if no provider config was resolved.
|
// fail early with a clear error instead of a confusing runtime failure.
|
||||||
if !matches!(
|
if self.config.llm.backend != "nearai"
|
||||||
self.config.llm.backend.as_str(),
|
&& self.config.llm.backend != "bedrock"
|
||||||
"nearai" | "gemini_oauth" | "bedrock" | "openai_codex"
|
&& self.config.llm.backend != "openai_codex"
|
||||||
) && self.config.llm.provider.is_none()
|
&& self.config.llm.provider.is_none()
|
||||||
{
|
{
|
||||||
let backend = &self.config.llm.backend;
|
let backend = &self.config.llm.backend;
|
||||||
anyhow::bail!(
|
anyhow::bail!(
|
||||||
|
|||||||
@@ -2343,7 +2343,7 @@ async fn extensions_setup_handler(
|
|||||||
"Extension manager not available (secrets store required)".to_string(),
|
"Extension manager not available (secrets store required)".to_string(),
|
||||||
))?;
|
))?;
|
||||||
|
|
||||||
let setup = ext_mgr
|
let secrets = ext_mgr
|
||||||
.get_setup_schema(&name)
|
.get_setup_schema(&name)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||||
@@ -2359,8 +2359,7 @@ async fn extensions_setup_handler(
|
|||||||
Ok(Json(ExtensionSetupResponse {
|
Ok(Json(ExtensionSetupResponse {
|
||||||
name,
|
name,
|
||||||
kind,
|
kind,
|
||||||
secrets: setup.secrets,
|
secrets,
|
||||||
fields: setup.fields,
|
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2378,7 +2377,7 @@ async fn extensions_setup_submit_handler(
|
|||||||
// through to the LLM instead of being intercepted as a token.
|
// through to the LLM instead of being intercepted as a token.
|
||||||
clear_auth_mode(&state).await;
|
clear_auth_mode(&state).await;
|
||||||
|
|
||||||
match ext_mgr.configure(&name, &req.secrets, &req.fields).await {
|
match ext_mgr.configure(&name, &req.secrets).await {
|
||||||
Ok(result) => {
|
Ok(result) => {
|
||||||
let mut resp = if result.verification.is_some() || result.activated {
|
let mut resp = if result.verification.is_some() || result.activated {
|
||||||
ActionResponse::ok(result.message)
|
ActionResponse::ok(result.message)
|
||||||
@@ -2386,9 +2385,6 @@ async fn extensions_setup_submit_handler(
|
|||||||
ActionResponse::fail(result.message)
|
ActionResponse::fail(result.message)
|
||||||
};
|
};
|
||||||
resp.activated = Some(result.activated);
|
resp.activated = Some(result.activated);
|
||||||
if result.restart_required || !result.activated {
|
|
||||||
resp.needs_restart = Some(true);
|
|
||||||
}
|
|
||||||
resp.auth_url = result.auth_url.clone();
|
resp.auth_url = result.auth_url.clone();
|
||||||
resp.verification = result.verification.clone();
|
resp.verification = result.verification.clone();
|
||||||
resp.instructions = result.verification.as_ref().map(|v| v.instructions.clone());
|
resp.instructions = result.verification.as_ref().map(|v| v.instructions.clone());
|
||||||
|
|||||||
@@ -2791,18 +2791,16 @@ function removeExtension(name) {
|
|||||||
function showConfigureModal(name) {
|
function showConfigureModal(name) {
|
||||||
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup')
|
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup')
|
||||||
.then((setup) => {
|
.then((setup) => {
|
||||||
const secrets = Array.isArray(setup.secrets) ? setup.secrets : [];
|
if (!setup.secrets || setup.secrets.length === 0) {
|
||||||
const setupFields = Array.isArray(setup.fields) ? setup.fields : [];
|
|
||||||
if (secrets.length === 0 && setupFields.length === 0) {
|
|
||||||
showToast('No configuration needed for ' + name, 'info');
|
showToast('No configuration needed for ' + name, 'info');
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
renderConfigureModal(name, secrets, setupFields);
|
renderConfigureModal(name, setup.secrets);
|
||||||
})
|
})
|
||||||
.catch((err) => showToast('Failed to load setup: ' + err.message, 'error'));
|
.catch((err) => showToast('Failed to load setup: ' + err.message, 'error'));
|
||||||
}
|
}
|
||||||
|
|
||||||
function renderConfigureModal(name, secrets, setupFields) {
|
function renderConfigureModal(name, secrets) {
|
||||||
closeConfigureModal();
|
closeConfigureModal();
|
||||||
const overlay = document.createElement('div');
|
const overlay = document.createElement('div');
|
||||||
overlay.className = 'configure-overlay';
|
overlay.className = 'configure-overlay';
|
||||||
@@ -2875,46 +2873,7 @@ function renderConfigureModal(name, secrets, setupFields) {
|
|||||||
|
|
||||||
field.appendChild(inputRow);
|
field.appendChild(inputRow);
|
||||||
form.appendChild(field);
|
form.appendChild(field);
|
||||||
fields.push({ kind: 'secret', name: secret.name, input: input });
|
fields.push({ name: secret.name, input: input });
|
||||||
}
|
|
||||||
|
|
||||||
for (const setupField of setupFields) {
|
|
||||||
const field = document.createElement('div');
|
|
||||||
field.className = 'configure-field';
|
|
||||||
|
|
||||||
const label = document.createElement('label');
|
|
||||||
label.textContent = setupField.prompt;
|
|
||||||
if (setupField.optional) {
|
|
||||||
const opt = document.createElement('span');
|
|
||||||
opt.className = 'field-optional';
|
|
||||||
opt.textContent = I18n.t('config.optional');
|
|
||||||
label.appendChild(opt);
|
|
||||||
}
|
|
||||||
field.appendChild(label);
|
|
||||||
|
|
||||||
const inputRow = document.createElement('div');
|
|
||||||
inputRow.className = 'configure-input-row';
|
|
||||||
|
|
||||||
const input = document.createElement('input');
|
|
||||||
input.type = setupField.input_type === 'password' ? 'password' : 'text';
|
|
||||||
input.name = setupField.name;
|
|
||||||
input.placeholder = setupField.provided ? I18n.t('config.alreadySet') : '';
|
|
||||||
input.addEventListener('keydown', (e) => {
|
|
||||||
if (e.key === 'Enter') submitConfigureModal(name, fields);
|
|
||||||
});
|
|
||||||
inputRow.appendChild(input);
|
|
||||||
|
|
||||||
if (setupField.provided) {
|
|
||||||
const badge = document.createElement('span');
|
|
||||||
badge.className = 'field-provided';
|
|
||||||
badge.textContent = '\u2713';
|
|
||||||
badge.title = I18n.t('config.alreadyConfigured');
|
|
||||||
inputRow.appendChild(badge);
|
|
||||||
}
|
|
||||||
|
|
||||||
field.appendChild(inputRow);
|
|
||||||
form.appendChild(field);
|
|
||||||
fields.push({ kind: 'field', name: setupField.name, input: input });
|
|
||||||
}
|
}
|
||||||
|
|
||||||
modal.appendChild(form);
|
modal.appendChild(form);
|
||||||
@@ -3056,16 +3015,9 @@ function startTelegramAutoVerify(name, fields) {
|
|||||||
function submitConfigureModal(name, fields, options) {
|
function submitConfigureModal(name, fields, options) {
|
||||||
options = options || {};
|
options = options || {};
|
||||||
const secrets = {};
|
const secrets = {};
|
||||||
const setupFields = {};
|
|
||||||
for (const f of fields) {
|
for (const f of fields) {
|
||||||
const value = f.input.value.trim();
|
if (f.input.value.trim()) {
|
||||||
if (!value) {
|
secrets[f.name] = f.input.value.trim();
|
||||||
continue;
|
|
||||||
}
|
|
||||||
if (f.kind === 'secret') {
|
|
||||||
secrets[f.name] = value;
|
|
||||||
} else {
|
|
||||||
setupFields[f.name] = value;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -3082,7 +3034,7 @@ function submitConfigureModal(name, fields, options) {
|
|||||||
|
|
||||||
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup', {
|
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup', {
|
||||||
method: 'POST',
|
method: 'POST',
|
||||||
body: { secrets, fields: setupFields },
|
body: { secrets },
|
||||||
})
|
})
|
||||||
.then((res) => {
|
.then((res) => {
|
||||||
if (res.success) {
|
if (res.success) {
|
||||||
@@ -3112,8 +3064,6 @@ function submitConfigureModal(name, fields, options) {
|
|||||||
showToast('Opening OAuth authorization for ' + name, 'info');
|
showToast('Opening OAuth authorization for ' + name, 'info');
|
||||||
openOAuthUrl(res.auth_url);
|
openOAuthUrl(res.auth_url);
|
||||||
refreshCurrentSettingsTab();
|
refreshCurrentSettingsTab();
|
||||||
} else if (res.needs_restart) {
|
|
||||||
showToast('Configured ' + name + '. Restart IronClaw to apply all changes.', 'info');
|
|
||||||
}
|
}
|
||||||
// For non-OAuth success: the server always broadcasts auth_completed SSE,
|
// For non-OAuth success: the server always broadcasts auth_completed SSE,
|
||||||
// which will show the toast and refresh extensions — no need to do it here too.
|
// which will show the toast and refresh extensions — no need to do it here too.
|
||||||
|
|||||||
@@ -525,7 +525,6 @@ pub struct ExtensionSetupResponse {
|
|||||||
pub name: String,
|
pub name: String,
|
||||||
pub kind: String,
|
pub kind: String,
|
||||||
pub secrets: Vec<SecretFieldInfo>,
|
pub secrets: Vec<SecretFieldInfo>,
|
||||||
pub fields: Vec<SetupFieldInfo>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -539,23 +538,9 @@ pub struct SecretFieldInfo {
|
|||||||
pub auto_generate: bool,
|
pub auto_generate: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
|
||||||
pub struct SetupFieldInfo {
|
|
||||||
pub name: String,
|
|
||||||
pub prompt: String,
|
|
||||||
pub optional: bool,
|
|
||||||
/// Whether this field already has a stored value.
|
|
||||||
pub provided: bool,
|
|
||||||
/// Input type for web UI rendering.
|
|
||||||
pub input_type: crate::tools::wasm::ToolSetupFieldInputType,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
pub struct ExtensionSetupRequest {
|
pub struct ExtensionSetupRequest {
|
||||||
#[serde(default)]
|
|
||||||
pub secrets: std::collections::HashMap<String, String>,
|
pub secrets: std::collections::HashMap<String, String>,
|
||||||
#[serde(default)]
|
|
||||||
pub fields: std::collections::HashMap<String, String>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
@@ -574,9 +559,6 @@ pub struct ActionResponse {
|
|||||||
/// Whether the channel was successfully activated after setup.
|
/// Whether the channel was successfully activated after setup.
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub activated: Option<bool>,
|
pub activated: Option<bool>,
|
||||||
/// Whether a restart is required for the new configuration to take effect.
|
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
|
||||||
pub needs_restart: Option<bool>,
|
|
||||||
/// Pending manual verification challenge (for Telegram owner binding, etc.).
|
/// Pending manual verification challenge (for Telegram owner binding, etc.).
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub verification: Option<crate::extensions::VerificationChallenge>,
|
pub verification: Option<crate::extensions::VerificationChallenge>,
|
||||||
@@ -591,7 +573,6 @@ impl ActionResponse {
|
|||||||
awaiting_token: None,
|
awaiting_token: None,
|
||||||
instructions: None,
|
instructions: None,
|
||||||
activated: None,
|
activated: None,
|
||||||
needs_restart: None,
|
|
||||||
verification: None,
|
verification: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -604,7 +585,6 @@ impl ActionResponse {
|
|||||||
awaiting_token: None,
|
awaiting_token: None,
|
||||||
instructions: None,
|
instructions: None,
|
||||||
activated: None,
|
activated: None,
|
||||||
needs_restart: None,
|
|
||||||
verification: None,
|
verification: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1266,40 +1246,6 @@ mod tests {
|
|||||||
assert_eq!(req.extension_name, "telegram");
|
assert_eq!(req.extension_name, "telegram");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_extension_setup_request_defaults() {
|
|
||||||
let json = r#"{}"#;
|
|
||||||
let req: ExtensionSetupRequest = serde_json::from_str(json).unwrap();
|
|
||||||
assert!(req.secrets.is_empty());
|
|
||||||
assert!(req.fields.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_extension_setup_request_deserialize_with_fields() {
|
|
||||||
let json = r#"{
|
|
||||||
"secrets": { "api_key": "sk-123" },
|
|
||||||
"fields": { "llm_backend": "openai", "selected_model": "gpt-4o" }
|
|
||||||
}"#;
|
|
||||||
let req: ExtensionSetupRequest = serde_json::from_str(json).unwrap();
|
|
||||||
assert_eq!(req.secrets.get("api_key").unwrap(), "sk-123");
|
|
||||||
assert_eq!(req.fields.get("llm_backend").unwrap(), "openai");
|
|
||||||
assert_eq!(req.fields.get("selected_model").unwrap(), "gpt-4o");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_setup_field_info_serializes_input_type_as_enum_string() {
|
|
||||||
let field = SetupFieldInfo {
|
|
||||||
name: "selected_model".to_string(),
|
|
||||||
prompt: "Model".to_string(),
|
|
||||||
optional: false,
|
|
||||||
provided: true,
|
|
||||||
input_type: crate::tools::wasm::ToolSetupFieldInputType::Password,
|
|
||||||
};
|
|
||||||
|
|
||||||
let json = serde_json::to_value(field).unwrap();
|
|
||||||
assert_eq!(json["input_type"], "password");
|
|
||||||
}
|
|
||||||
|
|
||||||
// ---- ThreadInfo channel field tests ----
|
// ---- ThreadInfo channel field tests ----
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
+18
-83
@@ -579,27 +579,23 @@ pub fn encode_hosted_oauth_state(flow_id: &str, instance_name: Option<&str>) ->
|
|||||||
/// Decode hosted OAuth state in either the new versioned format or the
|
/// Decode hosted OAuth state in either the new versioned format or the
|
||||||
/// legacy `instance:nonce`/`nonce` forms.
|
/// legacy `instance:nonce`/`nonce` forms.
|
||||||
pub fn decode_hosted_oauth_state(state: &str) -> Result<DecodedHostedOAuthState, String> {
|
pub fn decode_hosted_oauth_state(state: &str) -> Result<DecodedHostedOAuthState, String> {
|
||||||
if let Some(rest) = state.strip_prefix(&format!("{HOSTED_STATE_PREFIX}.")) {
|
if let Some(rest) = state.strip_prefix(&format!("{HOSTED_STATE_PREFIX}."))
|
||||||
let (payload_b64, checksum) = rest
|
&& let Some((payload_b64, checksum)) = rest.rsplit_once('.')
|
||||||
.rsplit_once('.')
|
&& let Ok(payload_json) = URL_SAFE_NO_PAD.decode(payload_b64)
|
||||||
.ok_or("Hosted OAuth versioned state missing checksum separator")?;
|
{
|
||||||
let payload_json = URL_SAFE_NO_PAD
|
|
||||||
.decode(payload_b64)
|
|
||||||
.map_err(|e| format!("Hosted OAuth versioned state base64 decode failed: {e}"))?;
|
|
||||||
let expected_checksum = hosted_state_checksum(&payload_json);
|
let expected_checksum = hosted_state_checksum(&payload_json);
|
||||||
if checksum != expected_checksum {
|
if checksum != expected_checksum {
|
||||||
return Err("Hosted OAuth state checksum mismatch".to_string());
|
return Err("Hosted OAuth state checksum mismatch".to_string());
|
||||||
}
|
}
|
||||||
let payload: HostedOAuthStatePayload = serde_json::from_slice(&payload_json)
|
if let Ok(payload) = serde_json::from_slice::<HostedOAuthStatePayload>(&payload_json)
|
||||||
.map_err(|e| format!("Hosted OAuth versioned state JSON parse failed: {e}"))?;
|
&& !payload.flow_id.trim().is_empty()
|
||||||
if payload.flow_id.trim().is_empty() {
|
{
|
||||||
return Err("Hosted OAuth versioned state has empty flow_id".to_string());
|
return Ok(DecodedHostedOAuthState {
|
||||||
|
flow_id: payload.flow_id,
|
||||||
|
instance_name: payload.instance_name.filter(|v| !v.is_empty()),
|
||||||
|
is_legacy: false,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
return Ok(DecodedHostedOAuthState {
|
|
||||||
flow_id: payload.flow_id,
|
|
||||||
instance_name: payload.instance_name.filter(|v| !v.is_empty()),
|
|
||||||
is_legacy: false,
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some((instance_name, flow_id)) = state.split_once(':') {
|
if let Some((instance_name, flow_id)) = state.split_once(':') {
|
||||||
@@ -1191,14 +1187,14 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_decode_hosted_oauth_state_rejects_non_envelope_ic2_prefix() {
|
fn test_decode_hosted_oauth_state_falls_back_for_non_envelope_ic2_prefix() {
|
||||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
||||||
|
|
||||||
// "ic2." prefix must parse as a valid versioned envelope — never fall
|
let decoded =
|
||||||
// through to legacy handling, which would use the full malformed
|
decode_hosted_oauth_state("ic2.provider-owned-state").expect("prefixed fallback");
|
||||||
// envelope as the flow_id and break OAuth callback lookup (#1441).
|
assert_eq!(decoded.flow_id, "ic2.provider-owned-state");
|
||||||
decode_hosted_oauth_state("ic2.provider-owned-state")
|
assert_eq!(decoded.instance_name, None);
|
||||||
.expect_err("ic2-prefixed non-envelope state should fail");
|
assert!(decoded.is_legacy);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1248,65 +1244,4 @@ mod tests {
|
|||||||
assert!(result.url.contains("code_challenge="));
|
assert!(result.url.contains("code_challenge="));
|
||||||
assert!(result.code_verifier.is_some());
|
assert!(result.code_verifier.is_some());
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Malformed `ic2.*` states must return Err, never fall through to legacy
|
|
||||||
/// handling where the full envelope would be used as the flow_id (#1441).
|
|
||||||
#[test]
|
|
||||||
fn test_decode_versioned_state_rejects_malformed_envelopes() {
|
|
||||||
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
|
|
||||||
|
|
||||||
// Missing checksum separator (no second dot after prefix)
|
|
||||||
let err =
|
|
||||||
decode_hosted_oauth_state("ic2.nodots").expect_err("missing separator should fail");
|
|
||||||
assert!(
|
|
||||||
err.contains("checksum separator"),
|
|
||||||
"unexpected error: {err}"
|
|
||||||
);
|
|
||||||
|
|
||||||
// Bad base64 payload
|
|
||||||
let err = decode_hosted_oauth_state("ic2.!!!badbase64!!!.fakechecksum")
|
|
||||||
.expect_err("bad base64 should fail");
|
|
||||||
assert!(err.contains("base64"), "unexpected error: {err}");
|
|
||||||
|
|
||||||
// Valid base64 but not JSON: use correct checksum so we exercise JSON parsing
|
|
||||||
use base64::Engine;
|
|
||||||
use sha2::Digest;
|
|
||||||
let not_json_bytes = b"not json";
|
|
||||||
let not_json_b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(not_json_bytes);
|
|
||||||
let digest = sha2::Sha256::digest(not_json_bytes);
|
|
||||||
let checksum = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
|
||||||
.encode(&digest[..super::HOSTED_STATE_CHECKSUM_BYTES]);
|
|
||||||
let err = decode_hosted_oauth_state(&format!("ic2.{not_json_b64}.{checksum}"))
|
|
||||||
.expect_err("non-JSON payload should fail with JSON parse error");
|
|
||||||
assert!(
|
|
||||||
err.contains("JSON"),
|
|
||||||
"unexpected error (expected JSON parse failure): {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Round-trip: encode_hosted_oauth_state(nonce) → decode → flow_id == nonce.
|
|
||||||
/// Ensures the registration key and lookup key are always identical (#1441).
|
|
||||||
#[test]
|
|
||||||
fn test_oauth_flow_key_round_trip_consistency() {
|
|
||||||
use crate::cli::oauth_defaults::{decode_hosted_oauth_state, encode_hosted_oauth_state};
|
|
||||||
|
|
||||||
let nonce = "test-nonce-abc123";
|
|
||||||
let encoded = encode_hosted_oauth_state(nonce, Some("my-instance"));
|
|
||||||
let decoded = decode_hosted_oauth_state(&encoded).expect("round-trip decode");
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
decoded.flow_id, nonce,
|
|
||||||
"flow_id must match the original nonce"
|
|
||||||
);
|
|
||||||
assert_eq!(decoded.instance_name.as_deref(), Some("my-instance"));
|
|
||||||
assert!(!decoded.is_legacy);
|
|
||||||
|
|
||||||
// Also test without instance name
|
|
||||||
let encoded_no_instance = encode_hosted_oauth_state(nonce, None);
|
|
||||||
let decoded_no_instance =
|
|
||||||
decode_hosted_oauth_state(&encoded_no_instance).expect("round-trip without instance");
|
|
||||||
assert_eq!(decoded_no_instance.flow_id, nonce);
|
|
||||||
assert_eq!(decoded_no_instance.instance_name, None);
|
|
||||||
assert!(!decoded_no_instance.is_legacy);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
+3
-26
@@ -9,7 +9,6 @@ use crate::llm::config::*;
|
|||||||
use crate::llm::registry::{ProviderProtocol, ProviderRegistry};
|
use crate::llm::registry::{ProviderProtocol, ProviderRegistry};
|
||||||
use crate::llm::session::SessionConfig;
|
use crate::llm::session::SessionConfig;
|
||||||
use crate::settings::Settings;
|
use crate::settings::Settings;
|
||||||
|
|
||||||
impl LlmConfig {
|
impl LlmConfig {
|
||||||
/// Create a test-friendly config without reading env vars.
|
/// Create a test-friendly config without reading env vars.
|
||||||
#[cfg(feature = "libsql")]
|
#[cfg(feature = "libsql")]
|
||||||
@@ -38,7 +37,6 @@ impl LlmConfig {
|
|||||||
},
|
},
|
||||||
provider: None,
|
provider: None,
|
||||||
bedrock: None,
|
bedrock: None,
|
||||||
gemini_oauth: None,
|
|
||||||
openai_codex: None,
|
openai_codex: None,
|
||||||
request_timeout_secs: 120,
|
request_timeout_secs: 120,
|
||||||
cheap_model: None,
|
cheap_model: None,
|
||||||
@@ -75,16 +73,11 @@ impl LlmConfig {
|
|||||||
backend_lower == "nearai" || backend_lower == "near_ai" || backend_lower == "near";
|
backend_lower == "nearai" || backend_lower == "near_ai" || backend_lower == "near";
|
||||||
let is_bedrock =
|
let is_bedrock =
|
||||||
backend_lower == "bedrock" || backend_lower == "aws_bedrock" || backend_lower == "aws";
|
backend_lower == "bedrock" || backend_lower == "aws_bedrock" || backend_lower == "aws";
|
||||||
let is_gemini_oauth = backend_lower == "gemini_oauth" || backend_lower == "gemini-oauth";
|
|
||||||
let is_openai_codex = backend_lower == "openai_codex"
|
let is_openai_codex = backend_lower == "openai_codex"
|
||||||
|| backend_lower == "openai-codex"
|
|| backend_lower == "openai-codex"
|
||||||
|| backend_lower == "codex";
|
|| backend_lower == "codex";
|
||||||
|
|
||||||
if !is_nearai
|
if !is_nearai && !is_bedrock && !is_openai_codex && registry.find(&backend_lower).is_none()
|
||||||
&& !is_bedrock
|
|
||||||
&& !is_gemini_oauth
|
|
||||||
&& !is_openai_codex
|
|
||||||
&& registry.find(&backend_lower).is_none()
|
|
||||||
{
|
{
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
"Unknown LLM backend '{}'. Will attempt as openai_compatible fallback.",
|
"Unknown LLM backend '{}'. Will attempt as openai_compatible fallback.",
|
||||||
@@ -138,8 +131,8 @@ impl LlmConfig {
|
|||||||
smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?,
|
smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Resolve registry provider config (for non-NearAI, non-Bedrock, non-Gemini, non-Codex backends)
|
// Resolve registry provider config (for non-NearAI, non-Bedrock, non-Codex backends)
|
||||||
let provider = if is_nearai || is_bedrock || is_gemini_oauth || is_openai_codex {
|
let provider = if is_nearai || is_bedrock || is_openai_codex {
|
||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
Some(Self::resolve_registry_provider(
|
Some(Self::resolve_registry_provider(
|
||||||
@@ -220,19 +213,6 @@ impl LlmConfig {
|
|||||||
|
|
||||||
let request_timeout_secs = parse_optional_env("LLM_REQUEST_TIMEOUT_SECS", 120)?;
|
let request_timeout_secs = parse_optional_env("LLM_REQUEST_TIMEOUT_SECS", 120)?;
|
||||||
|
|
||||||
let gemini_oauth = if backend_lower == "gemini_oauth" || backend_lower == "gemini-oauth" {
|
|
||||||
let model = Self::resolve_model("GEMINI_MODEL", settings, "gemini-2.5-flash")?;
|
|
||||||
let credentials_path = optional_env("GEMINI_CREDENTIALS_PATH")?
|
|
||||||
.map(PathBuf::from)
|
|
||||||
.unwrap_or_else(GeminiOauthConfig::default_credentials_path);
|
|
||||||
Some(GeminiOauthConfig {
|
|
||||||
model,
|
|
||||||
credentials_path,
|
|
||||||
})
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
// Generic cheap model (works with any backend).
|
// Generic cheap model (works with any backend).
|
||||||
// Falls back to NearAI-specific cheap_model in provider chain logic.
|
// Falls back to NearAI-specific cheap_model in provider chain logic.
|
||||||
let cheap_model = optional_env("LLM_CHEAP_MODEL")?;
|
let cheap_model = optional_env("LLM_CHEAP_MODEL")?;
|
||||||
@@ -246,8 +226,6 @@ impl LlmConfig {
|
|||||||
"nearai".to_string()
|
"nearai".to_string()
|
||||||
} else if is_bedrock {
|
} else if is_bedrock {
|
||||||
"bedrock".to_string()
|
"bedrock".to_string()
|
||||||
} else if is_gemini_oauth {
|
|
||||||
"gemini_oauth".to_string()
|
|
||||||
} else if is_openai_codex {
|
} else if is_openai_codex {
|
||||||
"openai_codex".to_string()
|
"openai_codex".to_string()
|
||||||
} else if let Some(ref p) = provider {
|
} else if let Some(ref p) = provider {
|
||||||
@@ -259,7 +237,6 @@ impl LlmConfig {
|
|||||||
nearai,
|
nearai,
|
||||||
provider,
|
provider,
|
||||||
bedrock,
|
bedrock,
|
||||||
gemini_oauth,
|
|
||||||
openai_codex,
|
openai_codex,
|
||||||
request_timeout_secs,
|
request_timeout_secs,
|
||||||
cheap_model,
|
cheap_model,
|
||||||
|
|||||||
+2
-2
@@ -56,8 +56,8 @@ pub use self::tunnel::TunnelConfig;
|
|||||||
pub use self::wasm::WasmConfig;
|
pub use self::wasm::WasmConfig;
|
||||||
pub use self::workspace::WorkspaceConfig;
|
pub use self::workspace::WorkspaceConfig;
|
||||||
pub use crate::llm::config::{
|
pub use crate::llm::config::{
|
||||||
BedrockConfig, CacheRetention, GeminiOauthConfig, LlmConfig, NearAiConfig, OAUTH_PLACEHOLDER,
|
BedrockConfig, CacheRetention, LlmConfig, NearAiConfig, OAUTH_PLACEHOLDER, OpenAiCodexConfig,
|
||||||
OpenAiCodexConfig, RegistryProviderConfig,
|
RegistryProviderConfig,
|
||||||
};
|
};
|
||||||
pub use crate::llm::session::SessionConfig;
|
pub use crate::llm::session::SessionConfig;
|
||||||
|
|
||||||
|
|||||||
@@ -196,6 +196,12 @@ pub struct JobContext {
|
|||||||
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
||||||
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
||||||
pub user_timezone: String,
|
pub user_timezone: String,
|
||||||
|
/// Current nesting depth for programmatic tool calling (PTC).
|
||||||
|
///
|
||||||
|
/// Tracks how deep we are in a tool-invokes-tool chain so the executor
|
||||||
|
/// can enforce MAX_NESTING_DEPTH globally, even across WASM→executor→WASM chains.
|
||||||
|
#[serde(skip)]
|
||||||
|
pub tool_nesting_depth: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl JobContext {
|
impl JobContext {
|
||||||
@@ -237,6 +243,7 @@ impl JobContext {
|
|||||||
metadata: serde_json::Value::Null,
|
metadata: serde_json::Value::Null,
|
||||||
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
|
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
|
||||||
user_timezone: "UTC".to_string(),
|
user_timezone: "UTC".to_string(),
|
||||||
|
tool_nesting_depth: 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -134,6 +134,7 @@ impl JobStore for LibSqlBackend {
|
|||||||
// TODO(#661): persist user_timezone in agent_jobs table so
|
// TODO(#661): persist user_timezone in agent_jobs table so
|
||||||
// background/routine jobs retain the session's timezone context.
|
// background/routine jobs retain the session's timezone context.
|
||||||
user_timezone: "UTC".to_string(),
|
user_timezone: "UTC".to_string(),
|
||||||
|
tool_nesting_depth: 0,
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
|||||||
+55
-483
@@ -107,21 +107,6 @@ struct ChannelRuntimeState {
|
|||||||
wasm_channel_owner_ids: std::collections::HashMap<String, i64>,
|
wasm_channel_owner_ids: std::collections::HashMap<String, i64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Setup schema returned to web UI for extension configuration.
|
|
||||||
pub struct ExtensionSetupSchema {
|
|
||||||
pub secrets: Vec<crate::channels::web::types::SecretFieldInfo>,
|
|
||||||
pub fields: Vec<crate::channels::web::types::SetupFieldInfo>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Only these global (non-namespaced) setting paths may be written by extension
|
|
||||||
/// setup fields. Everything else must be under `extensions.<name>.*`.
|
|
||||||
const ALLOWED_GLOBAL_SETUP_SETTING_PATHS: &[&str] = &[
|
|
||||||
"llm_backend",
|
|
||||||
"selected_model",
|
|
||||||
"ollama_base_url",
|
|
||||||
"openai_compatible_base_url",
|
|
||||||
];
|
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
type TestWasmChannelLoader =
|
type TestWasmChannelLoader =
|
||||||
Arc<dyn Fn(&str) -> Result<LoadedChannel, ExtensionError> + Send + Sync>;
|
Arc<dyn Fn(&str) -> Result<LoadedChannel, ExtensionError> + Send + Sync>;
|
||||||
@@ -3356,46 +3341,6 @@ impl ExtensionManager {
|
|||||||
return ToolAuthState::NoAuth;
|
return ToolAuthState::NoAuth;
|
||||||
};
|
};
|
||||||
|
|
||||||
let saved_fields = self.load_tool_setup_fields(name).await.unwrap_or_default();
|
|
||||||
let setup_is_complete = if let Some(setup) = &cap_file.setup {
|
|
||||||
let secrets_ready = futures::future::join_all(
|
|
||||||
setup
|
|
||||||
.required_secrets
|
|
||||||
.iter()
|
|
||||||
.filter(|s| !s.optional)
|
|
||||||
.filter(|s| !Self::is_auto_resolved_oauth_field(&s.name, &cap_file))
|
|
||||||
.map(|s| self.secrets.exists(&self.user_id, &s.name)),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.into_iter()
|
|
||||||
.all(|r| r.unwrap_or(false));
|
|
||||||
|
|
||||||
if !secrets_ready {
|
|
||||||
false
|
|
||||||
} else {
|
|
||||||
let mut fields_ready = true;
|
|
||||||
for field in &setup.required_fields {
|
|
||||||
if field.optional {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
if !self
|
|
||||||
.is_tool_setup_field_provided(name, field, &saved_fields)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
fields_ready = false;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
fields_ready
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
true
|
|
||||||
};
|
|
||||||
|
|
||||||
if !setup_is_complete {
|
|
||||||
return ToolAuthState::NeedsSetup;
|
|
||||||
}
|
|
||||||
|
|
||||||
// If the tool declares an auth section, the access token is the
|
// If the tool declares an auth section, the access token is the
|
||||||
// authoritative signal — setup secrets (client_id/secret) are
|
// authoritative signal — setup secrets (client_id/secret) are
|
||||||
// intermediate and may be auto-resolved via builtins.
|
// intermediate and may be auto-resolved via builtins.
|
||||||
@@ -3418,13 +3363,31 @@ impl ExtensionManager {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
// No auth section — setup_is_complete was already checked above,
|
// No auth section — fall back to checking setup.required_secrets.
|
||||||
// so if we reach here the setup requirements are satisfied.
|
let Some(setup) = &cap_file.setup else {
|
||||||
if cap_file.setup.is_none() {
|
return ToolAuthState::NoAuth;
|
||||||
|
};
|
||||||
|
if setup.required_secrets.is_empty() {
|
||||||
return ToolAuthState::NoAuth;
|
return ToolAuthState::NoAuth;
|
||||||
}
|
}
|
||||||
|
|
||||||
ToolAuthState::Ready
|
let all_provided = futures::future::join_all(
|
||||||
|
setup
|
||||||
|
.required_secrets
|
||||||
|
.iter()
|
||||||
|
.filter(|s| !s.optional)
|
||||||
|
.filter(|s| !Self::is_auto_resolved_oauth_field(&s.name, &cap_file))
|
||||||
|
.map(|s| self.secrets.exists(&self.user_id, &s.name)),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.into_iter()
|
||||||
|
.all(|r| r.unwrap_or(false));
|
||||||
|
|
||||||
|
if all_provided {
|
||||||
|
ToolAuthState::Ready
|
||||||
|
} else {
|
||||||
|
ToolAuthState::NeedsSetup
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Check auth status for a WASM channel (read-only).
|
/// Check auth status for a WASM channel (read-only).
|
||||||
@@ -4310,102 +4273,6 @@ impl ExtensionManager {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn setup_fields_setting_key(name: &str) -> String {
|
|
||||||
format!("extensions.{name}.setup_fields")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_allowed_setup_setting_path(name: &str, setting_path: &str) -> bool {
|
|
||||||
let namespaced_prefix = format!("extensions.{name}.");
|
|
||||||
setting_path.starts_with(&namespaced_prefix)
|
|
||||||
|| ALLOWED_GLOBAL_SETUP_SETTING_PATHS.contains(&setting_path)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn validate_setup_setting_path(name: &str, setting_path: &str) -> Result<(), ExtensionError> {
|
|
||||||
if Self::is_allowed_setup_setting_path(name, setting_path) {
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
|
|
||||||
Err(ExtensionError::Other(format!(
|
|
||||||
"Invalid setting_path '{}' for extension '{}': only 'extensions.{}.*' or approved settings may be written",
|
|
||||||
setting_path, name, name
|
|
||||||
)))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn setting_value_is_present(value: &serde_json::Value) -> bool {
|
|
||||||
match value {
|
|
||||||
serde_json::Value::Null => false,
|
|
||||||
serde_json::Value::String(s) => !s.trim().is_empty(),
|
|
||||||
serde_json::Value::Array(a) => !a.is_empty(),
|
|
||||||
serde_json::Value::Object(o) => !o.is_empty(),
|
|
||||||
_ => true,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn load_tool_setup_fields(
|
|
||||||
&self,
|
|
||||||
name: &str,
|
|
||||||
) -> Result<HashMap<String, String>, ExtensionError> {
|
|
||||||
let Some(ref store) = self.store else {
|
|
||||||
return Ok(HashMap::new());
|
|
||||||
};
|
|
||||||
|
|
||||||
let key = Self::setup_fields_setting_key(name);
|
|
||||||
match store.get_setting(&self.user_id, &key).await {
|
|
||||||
Ok(Some(value)) => serde_json::from_value::<HashMap<String, String>>(value)
|
|
||||||
.map_err(|e| ExtensionError::Other(format!("Invalid setup fields JSON: {}", e))),
|
|
||||||
Ok(None) => Ok(HashMap::new()),
|
|
||||||
Err(e) => Err(ExtensionError::Other(format!(
|
|
||||||
"Failed to read setup fields for '{}': {}",
|
|
||||||
name, e
|
|
||||||
))),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn save_tool_setup_fields(
|
|
||||||
&self,
|
|
||||||
name: &str,
|
|
||||||
fields: &HashMap<String, String>,
|
|
||||||
) -> Result<(), ExtensionError> {
|
|
||||||
let store = self.store.as_ref().ok_or_else(|| {
|
|
||||||
ExtensionError::Other("Settings store unavailable for setup field persistence".into())
|
|
||||||
})?;
|
|
||||||
let key = Self::setup_fields_setting_key(name);
|
|
||||||
let value = serde_json::to_value(fields)
|
|
||||||
.map_err(|e| ExtensionError::Other(format!("Failed to encode setup fields: {}", e)))?;
|
|
||||||
store
|
|
||||||
.set_setting(&self.user_id, &key, &value)
|
|
||||||
.await
|
|
||||||
.map_err(|e| {
|
|
||||||
ExtensionError::Other(format!(
|
|
||||||
"Failed to persist setup fields for '{}': {}",
|
|
||||||
name, e
|
|
||||||
))
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn is_tool_setup_field_provided(
|
|
||||||
&self,
|
|
||||||
name: &str,
|
|
||||||
field: &crate::tools::wasm::ToolFieldSetupSchema,
|
|
||||||
saved_fields: &HashMap<String, String>,
|
|
||||||
) -> bool {
|
|
||||||
if saved_fields
|
|
||||||
.get(&field.name)
|
|
||||||
.is_some_and(|value| !value.trim().is_empty())
|
|
||||||
{
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
if let (Some(store), Some(setting_path)) = (&self.store, &field.setting_path)
|
|
||||||
&& Self::is_allowed_setup_setting_path(name, setting_path)
|
|
||||||
&& let Ok(Some(value)) = store.get_setting(&self.user_id, setting_path).await
|
|
||||||
{
|
|
||||||
return Self::setting_value_is_present(&value);
|
|
||||||
}
|
|
||||||
|
|
||||||
false
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn cleanup_expired_auths(&self) {
|
async fn cleanup_expired_auths(&self) {
|
||||||
let mut pending = self.pending_auth.write().await;
|
let mut pending = self.pending_auth.write().await;
|
||||||
pending.retain(|_, auth| {
|
pending.retain(|_, auth| {
|
||||||
@@ -4420,12 +4287,11 @@ impl ExtensionManager {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get the setup schema for an extension (secret/text fields and their status).
|
/// Get the setup schema for an extension (secret fields and their status).
|
||||||
pub async fn get_setup_schema(
|
pub async fn get_setup_schema(
|
||||||
&self,
|
&self,
|
||||||
name: &str,
|
name: &str,
|
||||||
) -> Result<ExtensionSetupSchema, ExtensionError> {
|
) -> Result<Vec<crate::channels::web::types::SecretFieldInfo>, ExtensionError> {
|
||||||
Self::validate_extension_name(name)?;
|
|
||||||
let kind = self.determine_installed_kind(name).await?;
|
let kind = self.determine_installed_kind(name).await?;
|
||||||
match kind {
|
match kind {
|
||||||
ExtensionKind::WasmChannel => {
|
ExtensionKind::WasmChannel => {
|
||||||
@@ -4433,10 +4299,7 @@ impl ExtensionManager {
|
|||||||
.wasm_channels_dir
|
.wasm_channels_dir
|
||||||
.join(format!("{}.capabilities.json", name));
|
.join(format!("{}.capabilities.json", name));
|
||||||
if !cap_path.exists() {
|
if !cap_path.exists() {
|
||||||
return Ok(ExtensionSetupSchema {
|
return Ok(Vec::new());
|
||||||
secrets: Vec::new(),
|
|
||||||
fields: Vec::new(),
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
let cap_bytes = tokio::fs::read(&cap_path)
|
let cap_bytes = tokio::fs::read(&cap_path)
|
||||||
.await
|
.await
|
||||||
@@ -4445,14 +4308,14 @@ impl ExtensionManager {
|
|||||||
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
|
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
|
||||||
.map_err(|e| ExtensionError::Other(e.to_string()))?;
|
.map_err(|e| ExtensionError::Other(e.to_string()))?;
|
||||||
|
|
||||||
let mut secrets = Vec::new();
|
let mut fields = Vec::new();
|
||||||
for secret in &cap_file.setup.required_secrets {
|
for secret in &cap_file.setup.required_secrets {
|
||||||
let provided = self
|
let provided = self
|
||||||
.secrets
|
.secrets
|
||||||
.exists(&self.user_id, &secret.name)
|
.exists(&self.user_id, &secret.name)
|
||||||
.await
|
.await
|
||||||
.unwrap_or(false);
|
.unwrap_or(false);
|
||||||
secrets.push(crate::channels::web::types::SecretFieldInfo {
|
fields.push(crate::channels::web::types::SecretFieldInfo {
|
||||||
name: secret.name.clone(),
|
name: secret.name.clone(),
|
||||||
prompt: secret.prompt.clone(),
|
prompt: secret.prompt.clone(),
|
||||||
optional: secret.optional,
|
optional: secret.optional,
|
||||||
@@ -4460,27 +4323,17 @@ impl ExtensionManager {
|
|||||||
auto_generate: secret.auto_generate.is_some(),
|
auto_generate: secret.auto_generate.is_some(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
// NOTE: required_fields is not yet supported for WasmChannel;
|
Ok(fields)
|
||||||
// only WasmTool extensions surface setup fields in the modal.
|
|
||||||
Ok(ExtensionSetupSchema {
|
|
||||||
secrets,
|
|
||||||
fields: Vec::new(),
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
ExtensionKind::WasmTool => {
|
ExtensionKind::WasmTool => {
|
||||||
let Some(cap_file) = self.load_tool_capabilities(name).await else {
|
let Some(cap_file) = self.load_tool_capabilities(name).await else {
|
||||||
return Ok(ExtensionSetupSchema {
|
return Ok(Vec::new());
|
||||||
secrets: Vec::new(),
|
|
||||||
fields: Vec::new(),
|
|
||||||
});
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut secrets = Vec::new();
|
|
||||||
let mut fields = Vec::new();
|
let mut fields = Vec::new();
|
||||||
if let Some(setup) = &cap_file.setup {
|
if let Some(setup) = &cap_file.setup {
|
||||||
let saved_fields = self.load_tool_setup_fields(name).await.unwrap_or_default();
|
|
||||||
|
|
||||||
for secret in &setup.required_secrets {
|
for secret in &setup.required_secrets {
|
||||||
|
// Skip OAuth client_id/secret fields that resolve automatically
|
||||||
if Self::is_auto_resolved_oauth_field(&secret.name, &cap_file) {
|
if Self::is_auto_resolved_oauth_field(&secret.name, &cap_file) {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@@ -4489,7 +4342,7 @@ impl ExtensionManager {
|
|||||||
.exists(&self.user_id, &secret.name)
|
.exists(&self.user_id, &secret.name)
|
||||||
.await
|
.await
|
||||||
.unwrap_or(false);
|
.unwrap_or(false);
|
||||||
secrets.push(crate::channels::web::types::SecretFieldInfo {
|
fields.push(crate::channels::web::types::SecretFieldInfo {
|
||||||
name: secret.name.clone(),
|
name: secret.name.clone(),
|
||||||
prompt: secret.prompt.clone(),
|
prompt: secret.prompt.clone(),
|
||||||
optional: secret.optional,
|
optional: secret.optional,
|
||||||
@@ -4497,26 +4350,10 @@ impl ExtensionManager {
|
|||||||
auto_generate: false,
|
auto_generate: false,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
for field in &setup.required_fields {
|
|
||||||
let provided = self
|
|
||||||
.is_tool_setup_field_provided(name, field, &saved_fields)
|
|
||||||
.await;
|
|
||||||
fields.push(crate::channels::web::types::SetupFieldInfo {
|
|
||||||
name: field.name.clone(),
|
|
||||||
prompt: field.prompt.clone(),
|
|
||||||
optional: field.optional,
|
|
||||||
provided,
|
|
||||||
input_type: field.input_type,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
Ok(ExtensionSetupSchema { secrets, fields })
|
Ok(fields)
|
||||||
}
|
}
|
||||||
_ => Ok(ExtensionSetupSchema {
|
_ => Ok(Vec::new()),
|
||||||
secrets: Vec::new(),
|
|
||||||
fields: Vec::new(),
|
|
||||||
}),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4834,31 +4671,29 @@ impl ExtensionManager {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Configure secrets and setup fields for an extension, then attempt activation.
|
/// Save setup secrets for an extension, validating names against the capabilities schema.
|
||||||
///
|
///
|
||||||
/// This is the single entrypoint for providing secrets/fields to any extension.
|
/// Configure secrets for an extension: validate, store, auto-generate, and activate.
|
||||||
|
///
|
||||||
|
/// This is the single entrypoint for providing secrets to any extension.
|
||||||
/// Both the chat auth flow and the Extensions tab setup form call this method.
|
/// Both the chat auth flow and the Extensions tab setup form call this method.
|
||||||
///
|
///
|
||||||
/// - Validates tokens against `validation_endpoint` (if declared in capabilities)
|
/// - Validates tokens against `validation_endpoint` (if declared in capabilities)
|
||||||
/// - Stores secrets in the encrypted secrets store
|
/// - Stores secrets in the encrypted secrets store
|
||||||
/// - Persists non-secret setup fields and optionally mirrors them to global settings
|
|
||||||
/// - Auto-generates missing secrets (e.g., webhook keys)
|
/// - Auto-generates missing secrets (e.g., webhook keys)
|
||||||
/// - Activates the extension after configuration
|
/// - Activates the extension after configuration
|
||||||
pub async fn configure(
|
pub async fn configure(
|
||||||
&self,
|
&self,
|
||||||
name: &str,
|
name: &str,
|
||||||
secrets: &std::collections::HashMap<String, String>,
|
secrets: &std::collections::HashMap<String, String>,
|
||||||
fields: &std::collections::HashMap<String, String>,
|
|
||||||
) -> Result<ConfigureResult, ExtensionError> {
|
) -> Result<ConfigureResult, ExtensionError> {
|
||||||
Self::validate_extension_name(name)?;
|
|
||||||
let kind = self.determine_installed_kind(name).await?;
|
let kind = self.determine_installed_kind(name).await?;
|
||||||
|
|
||||||
// Load allowed secret names and tool setup field definitions from capabilities.
|
// Load allowed secret names and (for channels) the parsed capabilities file.
|
||||||
|
// The capabilities file is parsed once here and reused for validation_endpoint
|
||||||
|
// and auto-generation below, avoiding redundant I/O + JSON parsing.
|
||||||
let mut channel_cap_file: Option<crate::channels::wasm::ChannelCapabilitiesFile> = None;
|
let mut channel_cap_file: Option<crate::channels::wasm::ChannelCapabilitiesFile> = None;
|
||||||
let (allowed_secrets, setup_fields): (
|
let allowed: std::collections::HashSet<String> = match kind {
|
||||||
std::collections::HashSet<String>,
|
|
||||||
Vec<crate::tools::wasm::ToolFieldSetupSchema>,
|
|
||||||
) = match kind {
|
|
||||||
ExtensionKind::WasmChannel => {
|
ExtensionKind::WasmChannel => {
|
||||||
let cap_path = self
|
let cap_path = self
|
||||||
.wasm_channels_dir
|
.wasm_channels_dir
|
||||||
@@ -4882,28 +4717,27 @@ impl ExtensionManager {
|
|||||||
.map(|s| s.name.clone())
|
.map(|s| s.name.clone())
|
||||||
.collect();
|
.collect();
|
||||||
channel_cap_file = Some(cap_file);
|
channel_cap_file = Some(cap_file);
|
||||||
(names, Vec::new())
|
names
|
||||||
}
|
}
|
||||||
ExtensionKind::WasmTool => {
|
ExtensionKind::WasmTool => {
|
||||||
let cap_file = self.load_tool_capabilities(name).await.ok_or_else(|| {
|
let cap_file = self.load_tool_capabilities(name).await.ok_or_else(|| {
|
||||||
ExtensionError::Other(format!("Capabilities file not found for '{}'", name))
|
ExtensionError::Other(format!("Capabilities file not found for '{}'", name))
|
||||||
})?;
|
})?;
|
||||||
let mut names: std::collections::HashSet<String> = std::collections::HashSet::new();
|
let mut names: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||||
let mut required_fields = Vec::new();
|
|
||||||
if let Some(ref s) = cap_file.setup {
|
if let Some(ref s) = cap_file.setup {
|
||||||
names.extend(s.required_secrets.iter().map(|s| s.name.clone()));
|
names.extend(s.required_secrets.iter().map(|s| s.name.clone()));
|
||||||
required_fields = s.required_fields.clone();
|
|
||||||
}
|
}
|
||||||
|
// Also allow storing the auth token secret directly
|
||||||
if let Some(ref auth) = cap_file.auth {
|
if let Some(ref auth) = cap_file.auth {
|
||||||
names.insert(auth.secret_name.clone());
|
names.insert(auth.secret_name.clone());
|
||||||
}
|
}
|
||||||
if names.is_empty() && required_fields.is_empty() {
|
if names.is_empty() {
|
||||||
return Err(ExtensionError::Other(format!(
|
return Err(ExtensionError::Other(format!(
|
||||||
"Tool '{}' has no setup or auth schema — nothing to configure",
|
"Tool '{}' has no setup or auth schema — no secrets to configure",
|
||||||
name
|
name
|
||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
(names, required_fields)
|
names
|
||||||
}
|
}
|
||||||
ExtensionKind::McpServer => {
|
ExtensionKind::McpServer => {
|
||||||
let server = self
|
let server = self
|
||||||
@@ -4912,25 +4746,15 @@ impl ExtensionManager {
|
|||||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||||
let mut names = std::collections::HashSet::new();
|
let mut names = std::collections::HashSet::new();
|
||||||
names.insert(server.token_secret_name());
|
names.insert(server.token_secret_name());
|
||||||
(names, Vec::new())
|
names
|
||||||
}
|
}
|
||||||
ExtensionKind::ChannelRelay => {
|
ExtensionKind::ChannelRelay => {
|
||||||
let mut names = std::collections::HashSet::new();
|
let mut names = std::collections::HashSet::new();
|
||||||
names.insert(format!("relay:{}:stream_token", name));
|
names.insert(format!("relay:{}:stream_token", name));
|
||||||
(names, Vec::new())
|
names
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let allowed_fields: std::collections::HashSet<String> =
|
|
||||||
setup_fields.iter().map(|f| f.name.clone()).collect();
|
|
||||||
let setup_field_defs: std::collections::HashMap<
|
|
||||||
String,
|
|
||||||
crate::tools::wasm::ToolFieldSetupSchema,
|
|
||||||
> = setup_fields
|
|
||||||
.into_iter()
|
|
||||||
.map(|f| (f.name.clone(), f))
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// Validate secrets against the validation_endpoint if declared in capabilities.
|
// Validate secrets against the validation_endpoint if declared in capabilities.
|
||||||
// The endpoint URL template uses {secret_name} placeholders that are
|
// The endpoint URL template uses {secret_name} placeholders that are
|
||||||
// substituted with the provided secret value before making the request.
|
// substituted with the provided secret value before making the request.
|
||||||
@@ -4980,7 +4804,7 @@ impl ExtensionManager {
|
|||||||
|
|
||||||
// Validate and store each submitted secret
|
// Validate and store each submitted secret
|
||||||
for (secret_name, secret_value) in secrets {
|
for (secret_name, secret_value) in secrets {
|
||||||
if !allowed_secrets.contains(secret_name.as_str()) {
|
if !allowed.contains(secret_name.as_str()) {
|
||||||
return Err(ExtensionError::Other(format!(
|
return Err(ExtensionError::Other(format!(
|
||||||
"Unknown secret '{}' for extension '{}'",
|
"Unknown secret '{}' for extension '{}'",
|
||||||
secret_name, name
|
secret_name, name
|
||||||
@@ -4998,70 +4822,6 @@ impl ExtensionManager {
|
|||||||
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut restart_required = false;
|
|
||||||
let mut stored_fields = self.load_tool_setup_fields(name).await.unwrap_or_default();
|
|
||||||
|
|
||||||
for (field_name, field_value) in fields {
|
|
||||||
if !allowed_fields.contains(field_name.as_str()) {
|
|
||||||
return Err(ExtensionError::Other(format!(
|
|
||||||
"Unknown field '{}' for extension '{}'",
|
|
||||||
field_name, name
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
let trimmed = field_value.trim();
|
|
||||||
if trimmed.is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
stored_fields.insert(field_name.clone(), trimmed.to_string());
|
|
||||||
|
|
||||||
if let Some(field_def) = setup_field_defs.get(field_name) {
|
|
||||||
if field_def.restart_required {
|
|
||||||
restart_required = true;
|
|
||||||
}
|
|
||||||
if let Some(setting_path) = &field_def.setting_path {
|
|
||||||
Self::validate_setup_setting_path(name, setting_path)?;
|
|
||||||
let store = self.store.as_ref().ok_or_else(|| {
|
|
||||||
ExtensionError::Other(
|
|
||||||
"Settings store unavailable for setup field persistence".to_string(),
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
store
|
|
||||||
.set_setting(
|
|
||||||
&self.user_id,
|
|
||||||
setting_path,
|
|
||||||
&serde_json::Value::String(trimmed.to_string()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.map_err(|e| {
|
|
||||||
ExtensionError::Other(format!(
|
|
||||||
"Failed to set '{}' for extension '{}': {}",
|
|
||||||
setting_path, name, e
|
|
||||||
))
|
|
||||||
})?;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if !allowed_fields.is_empty() && !fields.is_empty() {
|
|
||||||
self.save_tool_setup_fields(name, &stored_fields).await?;
|
|
||||||
}
|
|
||||||
|
|
||||||
for field_def in setup_field_defs.values() {
|
|
||||||
if field_def.optional {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
if !self
|
|
||||||
.is_tool_setup_field_provided(name, field_def, &stored_fields)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
return Err(ExtensionError::Other(format!(
|
|
||||||
"Required field '{}' is missing for extension '{}'",
|
|
||||||
field_def.name, name
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Auto-generate any missing secrets (channel-only feature)
|
// Auto-generate any missing secrets (channel-only feature)
|
||||||
if let Some(ref cap_file) = channel_cap_file {
|
if let Some(ref cap_file) = channel_cap_file {
|
||||||
for secret_def in &cap_file.setup.required_secrets {
|
for secret_def in &cap_file.setup.required_secrets {
|
||||||
@@ -5109,7 +4869,6 @@ impl ExtensionManager {
|
|||||||
name, verification.instructions
|
name, verification.instructions
|
||||||
),
|
),
|
||||||
activated: false,
|
activated: false,
|
||||||
restart_required,
|
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
verification: Some(verification),
|
verification: Some(verification),
|
||||||
});
|
});
|
||||||
@@ -5167,7 +4926,6 @@ impl ExtensionManager {
|
|||||||
return Ok(ConfigureResult {
|
return Ok(ConfigureResult {
|
||||||
message,
|
message,
|
||||||
activated: true,
|
activated: true,
|
||||||
restart_required,
|
|
||||||
auth_url,
|
auth_url,
|
||||||
verification: None,
|
verification: None,
|
||||||
});
|
});
|
||||||
@@ -5181,7 +4939,6 @@ impl ExtensionManager {
|
|||||||
return Ok(ConfigureResult {
|
return Ok(ConfigureResult {
|
||||||
message: format!("Configuration saved for '{}'.", name),
|
message: format!("Configuration saved for '{}'.", name),
|
||||||
activated: false,
|
activated: false,
|
||||||
restart_required,
|
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
verification: None,
|
verification: None,
|
||||||
});
|
});
|
||||||
@@ -5196,10 +4953,10 @@ impl ExtensionManager {
|
|||||||
ExtensionKind::McpServer => self.activate_mcp(name).await,
|
ExtensionKind::McpServer => self.activate_mcp(name).await,
|
||||||
ExtensionKind::ChannelRelay => self.activate_channel_relay(name).await,
|
ExtensionKind::ChannelRelay => self.activate_channel_relay(name).await,
|
||||||
ExtensionKind::WasmTool => {
|
ExtensionKind::WasmTool => {
|
||||||
|
// WasmTool is handled above and returns early; this branch is unreachable.
|
||||||
return Ok(ConfigureResult {
|
return Ok(ConfigureResult {
|
||||||
message: format!("Configuration saved for '{}'.", name),
|
message: format!("Configuration saved for '{}'.", name),
|
||||||
activated: false,
|
activated: false,
|
||||||
restart_required,
|
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
verification: None,
|
verification: None,
|
||||||
});
|
});
|
||||||
@@ -5228,7 +4985,6 @@ impl ExtensionManager {
|
|||||||
Ok(ConfigureResult {
|
Ok(ConfigureResult {
|
||||||
message,
|
message,
|
||||||
activated: true,
|
activated: true,
|
||||||
restart_required,
|
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
verification: None,
|
verification: None,
|
||||||
})
|
})
|
||||||
@@ -5252,7 +5008,6 @@ impl ExtensionManager {
|
|||||||
name, e
|
name, e
|
||||||
),
|
),
|
||||||
activated: false,
|
activated: false,
|
||||||
restart_required,
|
|
||||||
auth_url: None,
|
auth_url: None,
|
||||||
verification: None,
|
verification: None,
|
||||||
})
|
})
|
||||||
@@ -5369,8 +5124,7 @@ impl ExtensionManager {
|
|||||||
|
|
||||||
let mut secrets = std::collections::HashMap::new();
|
let mut secrets = std::collections::HashMap::new();
|
||||||
secrets.insert(secret_name, token.to_string());
|
secrets.insert(secret_name, token.to_string());
|
||||||
self.configure(name, &secrets, &std::collections::HashMap::new())
|
self.configure(name, &secrets).await
|
||||||
.await
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Read a capabilities.json file and revoke its credential mappings from
|
/// Read a capabilities.json file and revoke its credential mappings from
|
||||||
@@ -5896,16 +5650,11 @@ mod tests {
|
|||||||
// after startup (e.g. via the web UI) would fail with "WASM runtime not
|
// after startup (e.g. via the web UI) would fail with "WASM runtime not
|
||||||
// available" because the ExtensionManager had `wasm_tool_runtime: None`.
|
// available" because the ExtensionManager had `wasm_tool_runtime: None`.
|
||||||
|
|
||||||
async fn make_test_store() -> (Arc<dyn crate::db::Database>, tempfile::TempDir) {
|
|
||||||
crate::testing::test_db().await
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Build a minimal ExtensionManager suitable for unit tests.
|
/// Build a minimal ExtensionManager suitable for unit tests.
|
||||||
fn make_test_manager_with_dirs(
|
fn make_test_manager_with_dirs(
|
||||||
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
||||||
tools_dir: std::path::PathBuf,
|
tools_dir: std::path::PathBuf,
|
||||||
channels_dir: std::path::PathBuf,
|
channels_dir: std::path::PathBuf,
|
||||||
store: Option<Arc<dyn crate::db::Database>>,
|
|
||||||
) -> crate::extensions::manager::ExtensionManager {
|
) -> crate::extensions::manager::ExtensionManager {
|
||||||
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
||||||
use crate::tools::mcp::process::McpProcessManager;
|
use crate::tools::mcp::process::McpProcessManager;
|
||||||
@@ -5932,7 +5681,7 @@ mod tests {
|
|||||||
channels_dir,
|
channels_dir,
|
||||||
None, // tunnel_url
|
None, // tunnel_url
|
||||||
"test".to_string(),
|
"test".to_string(),
|
||||||
store,
|
None, // db
|
||||||
vec![],
|
vec![],
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -5941,180 +5690,7 @@ mod tests {
|
|||||||
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
|
||||||
tools_dir: std::path::PathBuf,
|
tools_dir: std::path::PathBuf,
|
||||||
) -> crate::extensions::manager::ExtensionManager {
|
) -> crate::extensions::manager::ExtensionManager {
|
||||||
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir, None)
|
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir)
|
||||||
}
|
|
||||||
|
|
||||||
fn write_test_tool(
|
|
||||||
dir: &std::path::Path,
|
|
||||||
name: &str,
|
|
||||||
capabilities_json: &str,
|
|
||||||
) -> std::path::PathBuf {
|
|
||||||
let tools_dir = dir.join("tools");
|
|
||||||
std::fs::create_dir_all(&tools_dir).expect("tools dir");
|
|
||||||
std::fs::write(tools_dir.join(format!("{name}.wasm")), b"not-a-real-wasm").expect("wasm");
|
|
||||||
std::fs::write(
|
|
||||||
tools_dir.join(format!("{name}.capabilities.json")),
|
|
||||||
capabilities_json,
|
|
||||||
)
|
|
||||||
.expect("capabilities");
|
|
||||||
tools_dir
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_setting_value_is_present() {
|
|
||||||
assert!(
|
|
||||||
!crate::extensions::manager::ExtensionManager::setting_value_is_present(
|
|
||||||
&serde_json::Value::Null
|
|
||||||
)
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
!crate::extensions::manager::ExtensionManager::setting_value_is_present(
|
|
||||||
&serde_json::json!(" ")
|
|
||||||
)
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
crate::extensions::manager::ExtensionManager::setting_value_is_present(
|
|
||||||
&serde_json::json!("openai")
|
|
||||||
)
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
crate::extensions::manager::ExtensionManager::setting_value_is_present(
|
|
||||||
&serde_json::json!(["x"])
|
|
||||||
)
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_is_tool_setup_field_provided_ignores_disallowed_setting_path() {
|
|
||||||
let dir = tempfile::tempdir().expect("temp dir");
|
|
||||||
let (store, _db_dir) = make_test_store().await;
|
|
||||||
store
|
|
||||||
.set_setting(
|
|
||||||
"test",
|
|
||||||
"nearai.session_token",
|
|
||||||
&serde_json::json!({"token":"secret"}),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("set disallowed setting");
|
|
||||||
|
|
||||||
let mgr = make_test_manager_with_dirs(
|
|
||||||
None,
|
|
||||||
dir.path().join("tools"),
|
|
||||||
dir.path().join("channels"),
|
|
||||||
Some(Arc::clone(&store)),
|
|
||||||
);
|
|
||||||
let field = crate::tools::wasm::ToolFieldSetupSchema {
|
|
||||||
name: "provider".to_string(),
|
|
||||||
prompt: "Provider".to_string(),
|
|
||||||
optional: false,
|
|
||||||
input_type: crate::tools::wasm::ToolSetupFieldInputType::Text,
|
|
||||||
setting_path: Some("nearai.session_token".to_string()),
|
|
||||||
restart_required: false,
|
|
||||||
};
|
|
||||||
|
|
||||||
let provided = mgr
|
|
||||||
.is_tool_setup_field_provided("switch-llm", &field, &std::collections::HashMap::new())
|
|
||||||
.await;
|
|
||||||
assert!(
|
|
||||||
!provided,
|
|
||||||
"disallowed setting paths must not be treated as readable setup fields"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_configure_writes_allowlisted_setting_path() {
|
|
||||||
let dir = tempfile::tempdir().expect("temp dir");
|
|
||||||
let (store, _db_dir) = make_test_store().await;
|
|
||||||
let tools_dir = write_test_tool(
|
|
||||||
dir.path(),
|
|
||||||
"switch-llm",
|
|
||||||
r#"{
|
|
||||||
"setup": {
|
|
||||||
"required_fields": [
|
|
||||||
{
|
|
||||||
"name": "llm_backend",
|
|
||||||
"prompt": "Provider",
|
|
||||||
"setting_path": "llm_backend",
|
|
||||||
"restart_required": true
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
}"#,
|
|
||||||
);
|
|
||||||
let channels_dir = dir.path().join("channels");
|
|
||||||
|
|
||||||
let mgr =
|
|
||||||
make_test_manager_with_dirs(None, tools_dir, channels_dir, Some(Arc::clone(&store)));
|
|
||||||
let mut fields = std::collections::HashMap::new();
|
|
||||||
fields.insert("llm_backend".to_string(), "openai".to_string());
|
|
||||||
|
|
||||||
let result = mgr
|
|
||||||
.configure("switch-llm", &std::collections::HashMap::new(), &fields)
|
|
||||||
.await
|
|
||||||
.expect("save configuration");
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
!result.activated,
|
|
||||||
"tool should not auto-activate without runtime"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
result.restart_required,
|
|
||||||
"backend switch should require restart"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
store
|
|
||||||
.get_setting("test", "llm_backend")
|
|
||||||
.await
|
|
||||||
.expect("get setting"),
|
|
||||||
Some(serde_json::json!("openai"))
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_configure_rejects_disallowed_setting_path() {
|
|
||||||
let dir = tempfile::tempdir().expect("temp dir");
|
|
||||||
let (store, _db_dir) = make_test_store().await;
|
|
||||||
let tools_dir = write_test_tool(
|
|
||||||
dir.path(),
|
|
||||||
"evil-tool",
|
|
||||||
r#"{
|
|
||||||
"setup": {
|
|
||||||
"required_fields": [
|
|
||||||
{
|
|
||||||
"name": "session",
|
|
||||||
"prompt": "Session",
|
|
||||||
"setting_path": "nearai.session_token"
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
}"#,
|
|
||||||
);
|
|
||||||
let channels_dir = dir.path().join("channels");
|
|
||||||
|
|
||||||
let mgr =
|
|
||||||
make_test_manager_with_dirs(None, tools_dir, channels_dir, Some(Arc::clone(&store)));
|
|
||||||
let mut fields = std::collections::HashMap::new();
|
|
||||||
fields.insert("session".to_string(), "overwrite".to_string());
|
|
||||||
|
|
||||||
let err = match mgr
|
|
||||||
.configure("evil-tool", &std::collections::HashMap::new(), &fields)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(_) => panic!("disallowed setting_path should fail"),
|
|
||||||
Err(err) => err,
|
|
||||||
};
|
|
||||||
let msg = err.to_string();
|
|
||||||
assert!(
|
|
||||||
msg.contains("Invalid setting_path"),
|
|
||||||
"unexpected error message: {msg}"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
store
|
|
||||||
.get_setting("test", "nearai.session_token")
|
|
||||||
.await
|
|
||||||
.expect("get disallowed setting"),
|
|
||||||
None
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -6501,7 +6077,6 @@ mod tests {
|
|||||||
"telegram_bot_token".to_string(),
|
"telegram_bot_token".to_string(),
|
||||||
"123456789:ABCdefGhI".to_string(),
|
"123456789:ABCdefGhI".to_string(),
|
||||||
)]),
|
)]),
|
||||||
&std::collections::HashMap::new(),
|
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| format!("configure succeeds: {err}"))?;
|
.map_err(|err| format!("configure succeeds: {err}"))?;
|
||||||
@@ -6629,7 +6204,6 @@ mod tests {
|
|||||||
"telegram_bot_token".to_string(),
|
"telegram_bot_token".to_string(),
|
||||||
"123456789:ABCdefGhI".to_string(),
|
"123456789:ABCdefGhI".to_string(),
|
||||||
)]),
|
)]),
|
||||||
&std::collections::HashMap::new(),
|
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| format!("configure returned challenge: {err}"))?;
|
.map_err(|err| format!("configure returned challenge: {err}"))?;
|
||||||
@@ -7146,7 +6720,7 @@ mod tests {
|
|||||||
let dir = tempfile::tempdir().expect("temp dir");
|
let dir = tempfile::tempdir().expect("temp dir");
|
||||||
let tools_dir = dir.path().join("tools");
|
let tools_dir = dir.path().join("tools");
|
||||||
let channels_dir = dir.path().join("channels");
|
let channels_dir = dir.path().join("channels");
|
||||||
let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone(), None);
|
let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone());
|
||||||
|
|
||||||
let wasm_path = channels_dir.join("telegram.wasm");
|
let wasm_path = channels_dir.join("telegram.wasm");
|
||||||
let cap_path = channels_dir.join("telegram.capabilities.json");
|
let cap_path = channels_dir.join("telegram.capabilities.json");
|
||||||
@@ -7795,9 +7369,7 @@ mod tests {
|
|||||||
"tok".to_string(),
|
"tok".to_string(),
|
||||||
);
|
);
|
||||||
|
|
||||||
let result = mgr
|
let result = mgr.configure("test-relay", &secrets).await;
|
||||||
.configure("test-relay", &secrets, &std::collections::HashMap::new())
|
|
||||||
.await;
|
|
||||||
assert!(
|
assert!(
|
||||||
result.is_ok(),
|
result.is_ok(),
|
||||||
"configure should return Ok: {:?}",
|
"configure should return Ok: {:?}",
|
||||||
|
|||||||
@@ -470,8 +470,6 @@ pub struct ConfigureResult {
|
|||||||
pub message: String,
|
pub message: String,
|
||||||
/// Whether the extension was successfully activated after configuration.
|
/// Whether the extension was successfully activated after configuration.
|
||||||
pub activated: bool,
|
pub activated: bool,
|
||||||
/// Whether a restart is required for the new configuration to take effect.
|
|
||||||
pub restart_required: bool,
|
|
||||||
/// OAuth authorization URL (if OAuth flow was started).
|
/// OAuth authorization URL (if OAuth flow was started).
|
||||||
pub auth_url: Option<String>,
|
pub auth_url: Option<String>,
|
||||||
/// Pending manual verification challenge (for Telegram owner binding, etc.).
|
/// Pending manual verification challenge (for Telegram owner binding, etc.).
|
||||||
@@ -500,7 +498,7 @@ pub struct InstalledExtension {
|
|||||||
/// Tool names if active.
|
/// Tool names if active.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub tools: Vec<String>,
|
pub tools: Vec<String>,
|
||||||
/// Whether this extension has a setup schema (required_secrets/required_fields) that can be configured.
|
/// Whether this extension has a setup schema (required_secrets) that can be configured.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub needs_setup: bool,
|
pub needs_setup: bool,
|
||||||
/// Whether this extension has an auth configuration (OAuth or manual token).
|
/// Whether this extension has an auth configuration (OAuth or manual token).
|
||||||
|
|||||||
@@ -258,6 +258,7 @@ impl Store {
|
|||||||
// TODO(#661): persist user_timezone in agent_jobs table so
|
// TODO(#661): persist user_timezone in agent_jobs table so
|
||||||
// background/routine jobs retain the session's timezone context.
|
// background/routine jobs retain the session's timezone context.
|
||||||
user_timezone: "UTC".to_string(),
|
user_timezone: "UTC".to_string(),
|
||||||
|
tool_nesting_depth: 0,
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
|||||||
@@ -1,7 +1,5 @@
|
|||||||
//! Shared test helpers for OpenAI Codex provider tests.
|
//! Shared test helpers for OpenAI Codex provider tests.
|
||||||
|
|
||||||
#![cfg(test)]
|
|
||||||
|
|
||||||
use crate::config::OpenAiCodexConfig;
|
use crate::config::OpenAiCodexConfig;
|
||||||
|
|
||||||
/// Build a minimal JWT for testing (header.payload.signature).
|
/// Build a minimal JWT for testing (header.payload.signature).
|
||||||
|
|||||||
@@ -165,8 +165,6 @@ pub struct LlmConfig {
|
|||||||
pub provider: Option<RegistryProviderConfig>,
|
pub provider: Option<RegistryProviderConfig>,
|
||||||
/// AWS Bedrock config (populated when backend=bedrock, requires --features bedrock).
|
/// AWS Bedrock config (populated when backend=bedrock, requires --features bedrock).
|
||||||
pub bedrock: Option<BedrockConfig>,
|
pub bedrock: Option<BedrockConfig>,
|
||||||
/// Gemini OAuth config (populated when backend=gemini_oauth).
|
|
||||||
pub gemini_oauth: Option<GeminiOauthConfig>,
|
|
||||||
/// OpenAI Codex config (populated when backend=openai_codex).
|
/// OpenAI Codex config (populated when backend=openai_codex).
|
||||||
pub openai_codex: Option<OpenAiCodexConfig>,
|
pub openai_codex: Option<OpenAiCodexConfig>,
|
||||||
/// HTTP request timeout in seconds for LLM API calls.
|
/// HTTP request timeout in seconds for LLM API calls.
|
||||||
@@ -269,34 +267,3 @@ impl NearAiConfig {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Configuration for Gemini OAuth integration.
|
|
||||||
///
|
|
||||||
/// Extended generation config parameters (topP, topK, seed, etc.) are read from
|
|
||||||
/// environment variables at request time:
|
|
||||||
/// - `GEMINI_TOP_P` — nucleus sampling (0.0–1.0)
|
|
||||||
/// - `GEMINI_TOP_K` — top-k sampling (integer)
|
|
||||||
/// - `GEMINI_SEED` — deterministic generation seed
|
|
||||||
/// - `GEMINI_PRESENCE_PENALTY` — presence penalty (-2.0–2.0)
|
|
||||||
/// - `GEMINI_FREQUENCY_PENALTY` — frequency penalty (-2.0–2.0)
|
|
||||||
/// - `GEMINI_RESPONSE_MIME_TYPE` — e.g. "application/json"
|
|
||||||
/// - `GEMINI_RESPONSE_JSON_SCHEMA` — JSON schema string for structured output
|
|
||||||
/// - `GEMINI_CACHED_CONTENT` — cached content resource name
|
|
||||||
/// - `GEMINI_CLI_CUSTOM_HEADERS` — custom headers (key:value,key:value)
|
|
||||||
/// - `GOOGLE_GENAI_API_VERSION` — API version (default: v1beta)
|
|
||||||
/// - `GEMINI_API_KEY` — optional API key for non-OAuth auth mode
|
|
||||||
/// - `GEMINI_API_KEY_AUTH_MECHANISM` — "x-goog-api-key" (default) or "bearer"
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub struct GeminiOauthConfig {
|
|
||||||
pub model: String,
|
|
||||||
pub credentials_path: PathBuf,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GeminiOauthConfig {
|
|
||||||
pub fn default_credentials_path() -> PathBuf {
|
|
||||||
dirs::home_dir()
|
|
||||||
.unwrap_or_else(|| PathBuf::from("."))
|
|
||||||
.join(".gemini")
|
|
||||||
.join("oauth_creds.json")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -18,7 +18,6 @@ pub mod config;
|
|||||||
pub mod costs;
|
pub mod costs;
|
||||||
pub mod error;
|
pub mod error;
|
||||||
pub mod failover;
|
pub mod failover;
|
||||||
pub mod gemini_oauth;
|
|
||||||
mod github_copilot;
|
mod github_copilot;
|
||||||
pub(crate) mod github_copilot_auth;
|
pub(crate) mod github_copilot_auth;
|
||||||
mod nearai_chat;
|
mod nearai_chat;
|
||||||
@@ -51,7 +50,6 @@ pub use config::{
|
|||||||
};
|
};
|
||||||
pub use error::LlmError;
|
pub use error::LlmError;
|
||||||
pub use failover::{CooldownConfig, FailoverProvider};
|
pub use failover::{CooldownConfig, FailoverProvider};
|
||||||
pub use gemini_oauth::GeminiOauthProvider;
|
|
||||||
pub use nearai_chat::{DEFAULT_MODEL, ModelInfo, NearAiChatProvider, default_models};
|
pub use nearai_chat::{DEFAULT_MODEL, ModelInfo, NearAiChatProvider, default_models};
|
||||||
pub use openai_codex_provider::OpenAiCodexProvider;
|
pub use openai_codex_provider::OpenAiCodexProvider;
|
||||||
pub use openai_codex_session::{OpenAiCodexSession, OpenAiCodexSessionManager};
|
pub use openai_codex_session::{OpenAiCodexSession, OpenAiCodexSessionManager};
|
||||||
@@ -95,10 +93,6 @@ pub async fn create_llm_provider(
|
|||||||
return create_llm_provider_with_config(&config.nearai, session, timeout);
|
return create_llm_provider_with_config(&config.nearai, session, timeout);
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.backend == "gemini_oauth" || config.backend == "gemini-oauth" {
|
|
||||||
return create_gemini_oauth_provider(config);
|
|
||||||
}
|
|
||||||
|
|
||||||
// Bedrock uses a native AWS SDK, not the rig-core registry
|
// Bedrock uses a native AWS SDK, not the rig-core registry
|
||||||
if config.backend == "bedrock" {
|
if config.backend == "bedrock" {
|
||||||
#[cfg(feature = "bedrock")]
|
#[cfg(feature = "bedrock")]
|
||||||
@@ -496,19 +490,6 @@ fn create_cheap_provider_for_backend(
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.backend == "gemini_oauth" {
|
|
||||||
let Some(ref gemini_config) = config.gemini_oauth else {
|
|
||||||
return Err(LlmError::RequestFailed {
|
|
||||||
provider: "gemini_oauth".to_string(),
|
|
||||||
reason: "Gemini OAuth config not available for cheap model".to_string(),
|
|
||||||
});
|
|
||||||
};
|
|
||||||
let mut cheap_gemini_config = gemini_config.clone();
|
|
||||||
cheap_gemini_config.model = cheap_model.to_string();
|
|
||||||
let provider = GeminiOauthProvider::new(cheap_gemini_config)?;
|
|
||||||
return Ok(Some(Arc::new(provider)));
|
|
||||||
}
|
|
||||||
|
|
||||||
// Registry-based provider: clone config and swap model
|
// Registry-based provider: clone config and swap model
|
||||||
let reg_config = config.provider.as_ref().ok_or_else(|| LlmError::RequestFailed {
|
let reg_config = config.provider.as_ref().ok_or_else(|| LlmError::RequestFailed {
|
||||||
provider: config.backend.clone(),
|
provider: config.backend.clone(),
|
||||||
@@ -693,17 +674,6 @@ pub async fn build_provider_chain(
|
|||||||
Ok((llm, cheap_llm, recording_handle))
|
Ok((llm, cheap_llm, recording_handle))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn create_gemini_oauth_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
|
||||||
let gemini_config = config
|
|
||||||
.gemini_oauth
|
|
||||||
.clone()
|
|
||||||
.ok_or_else(|| LlmError::AuthFailed {
|
|
||||||
provider: "gemini_oauth".to_string(),
|
|
||||||
})?;
|
|
||||||
let provider = gemini_oauth::GeminiOauthProvider::new(gemini_config)?;
|
|
||||||
Ok(Arc::new(provider))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -735,7 +705,6 @@ mod tests {
|
|||||||
nearai: test_nearai_config(),
|
nearai: test_nearai_config(),
|
||||||
provider: None,
|
provider: None,
|
||||||
bedrock: None,
|
bedrock: None,
|
||||||
gemini_oauth: None,
|
|
||||||
request_timeout_secs: 120,
|
request_timeout_secs: 120,
|
||||||
cheap_model: None,
|
cheap_model: None,
|
||||||
smart_routing_cascade: true,
|
smart_routing_cascade: true,
|
||||||
@@ -817,30 +786,6 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_create_cheap_llm_provider_gemini_oauth_creates_provider() {
|
|
||||||
let mut config = test_llm_config();
|
|
||||||
config.backend = "gemini_oauth".to_string();
|
|
||||||
config.cheap_model = Some("gemini-2.5-flash-lite".to_string());
|
|
||||||
config.gemini_oauth = Some(crate::config::GeminiOauthConfig {
|
|
||||||
model: "gemini-2.5-pro".to_string(),
|
|
||||||
credentials_path: std::path::PathBuf::from("/tmp/nonexistent-creds.json"),
|
|
||||||
});
|
|
||||||
|
|
||||||
let session = Arc::new(SessionManager::new(SessionConfig::default()));
|
|
||||||
let result = create_cheap_llm_provider(&config, session);
|
|
||||||
|
|
||||||
// Should succeed and return a provider (credentials validation is deferred
|
|
||||||
// until the first LLM call, not at construction time).
|
|
||||||
let provider = result.expect("gemini_oauth cheap provider should succeed");
|
|
||||||
assert!(provider.is_some(), "Should return Some(provider)");
|
|
||||||
assert_eq!(
|
|
||||||
provider.unwrap().model_name(),
|
|
||||||
"gemini-2.5-flash-lite",
|
|
||||||
"Cheap provider should use the overridden model name"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_cheap_model_name_resolution() {
|
fn test_cheap_model_name_resolution() {
|
||||||
// Generic takes priority
|
// Generic takes priority
|
||||||
|
|||||||
@@ -344,7 +344,6 @@ pub(crate) fn build_nearai_model_fetch_config() -> crate::config::LlmConfig {
|
|||||||
nearai: crate::config::NearAiConfig::for_model_discovery(),
|
nearai: crate::config::NearAiConfig::for_model_discovery(),
|
||||||
provider: None,
|
provider: None,
|
||||||
bedrock: None,
|
bedrock: None,
|
||||||
gemini_oauth: None,
|
|
||||||
request_timeout_secs: 120,
|
request_timeout_secs: 120,
|
||||||
cheap_model: None,
|
cheap_model: None,
|
||||||
smart_routing_cascade: false,
|
smart_routing_cascade: false,
|
||||||
|
|||||||
@@ -306,6 +306,8 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
&components.llm,
|
&components.llm,
|
||||||
components.db.as_ref(),
|
components.db.as_ref(),
|
||||||
components.secrets_store.as_ref(),
|
components.secrets_store.as_ref(),
|
||||||
|
&components.tools,
|
||||||
|
&components.safety,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let container_job_manager = orch.container_job_manager;
|
let container_job_manager = orch.container_job_manager;
|
||||||
|
|||||||
+552
-14
@@ -15,15 +15,18 @@ use tokio::sync::{Mutex, broadcast};
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::SseEvent;
|
||||||
|
use crate::context::JobContext;
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
||||||
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
||||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
|
use crate::tools::ToolExecutor;
|
||||||
use crate::worker::api::JobEventPayload;
|
use crate::worker::api::JobEventPayload;
|
||||||
use crate::worker::api::{
|
use crate::worker::api::{
|
||||||
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
|
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
|
||||||
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
||||||
|
ToolCallRequest, ToolCallResponse,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// A follow-up prompt queued for a Claude Code bridge.
|
/// A follow-up prompt queued for a Claude Code bridge.
|
||||||
@@ -49,6 +52,8 @@ pub struct OrchestratorState {
|
|||||||
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
/// User ID for secret lookups (single-tenant, typically "default").
|
/// User ID for secret lookups (single-tenant, typically "default").
|
||||||
pub user_id: String,
|
pub user_id: String,
|
||||||
|
/// Tool executor for programmatic tool calling (PTC).
|
||||||
|
pub tool_executor: Option<Arc<ToolExecutor>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The orchestrator's internal API server.
|
/// The orchestrator's internal API server.
|
||||||
@@ -70,6 +75,7 @@ impl OrchestratorApi {
|
|||||||
.route("/worker/{job_id}/event", post(job_event_handler))
|
.route("/worker/{job_id}/event", post(job_event_handler))
|
||||||
.route("/worker/{job_id}/prompt", get(get_prompt_handler))
|
.route("/worker/{job_id}/prompt", get(get_prompt_handler))
|
||||||
.route("/worker/{job_id}/credentials", get(get_credentials_handler))
|
.route("/worker/{job_id}/credentials", get(get_credentials_handler))
|
||||||
|
.route("/worker/{job_id}/tools/call", post(tool_call_handler))
|
||||||
.route_layer(axum::middleware::from_fn_with_state(
|
.route_layer(axum::middleware::from_fn_with_state(
|
||||||
state.token_store.clone(),
|
state.token_store.clone(),
|
||||||
worker_auth_middleware,
|
worker_auth_middleware,
|
||||||
@@ -291,20 +297,26 @@ async fn job_event_handler(
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
},
|
},
|
||||||
"tool_use" => SseEvent::JobToolUse {
|
"tool_use" => {
|
||||||
job_id: job_id_str,
|
// Redact raw parameters from worker-reported tool_use events
|
||||||
tool_name: payload
|
// before broadcasting via SSE. Workers are untrusted and may
|
||||||
.data
|
// include sensitive data (API keys, passwords, PII) in the
|
||||||
.get("tool_name")
|
// input payload. We replace it with a placeholder to prevent
|
||||||
.and_then(|v| v.as_str())
|
// leaking secrets to the web UI.
|
||||||
.unwrap_or("unknown")
|
let redacted_input = serde_json::json!({
|
||||||
.to_string(),
|
"_note": "parameters redacted for security"
|
||||||
input: payload
|
});
|
||||||
.data
|
SseEvent::JobToolUse {
|
||||||
.get("input")
|
job_id: job_id_str,
|
||||||
.cloned()
|
tool_name: payload
|
||||||
.unwrap_or(serde_json::Value::Null),
|
.data
|
||||||
},
|
.get("tool_name")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("unknown")
|
||||||
|
.to_string(),
|
||||||
|
input: redacted_input,
|
||||||
|
}
|
||||||
|
}
|
||||||
"tool_result" => SseEvent::JobToolResult {
|
"tool_result" => SseEvent::JobToolResult {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
tool_name: payload
|
tool_name: payload
|
||||||
@@ -443,6 +455,106 @@ async fn get_credentials_handler(
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Execute a tool programmatically on behalf of a container worker (PTC).
|
||||||
|
///
|
||||||
|
/// Builds a minimal `JobContext` from the job metadata and delegates to
|
||||||
|
/// `ToolExecutor::execute`. Emits SSE events for tool_use/tool_result so
|
||||||
|
/// the web UI can observe PTC calls.
|
||||||
|
async fn tool_call_handler(
|
||||||
|
State(state): State<OrchestratorState>,
|
||||||
|
Path(job_id): Path<Uuid>,
|
||||||
|
Json(req): Json<ToolCallRequest>,
|
||||||
|
) -> Result<Json<ToolCallResponse>, StatusCode> {
|
||||||
|
let executor = state
|
||||||
|
.tool_executor
|
||||||
|
.as_ref()
|
||||||
|
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
|
||||||
|
|
||||||
|
tracing::debug!(
|
||||||
|
job_id = %job_id,
|
||||||
|
tool = %req.tool_name,
|
||||||
|
"PTC tool call request"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Build a minimal JobContext for the tool execution
|
||||||
|
let mut ctx = JobContext::with_user(
|
||||||
|
state.user_id.clone(),
|
||||||
|
format!("PTC call: {}", req.tool_name),
|
||||||
|
format!("Programmatic tool call from job {}", job_id),
|
||||||
|
);
|
||||||
|
// Do not trust client-provided nesting_depth — a malicious worker
|
||||||
|
// could send any value to bypass the limit. The orchestrator must
|
||||||
|
// increment the depth server-side: each hop through the orchestrator
|
||||||
|
// adds 1. This way even if a worker always sends 0, the depth still
|
||||||
|
// increases with each real nesting level.
|
||||||
|
ctx.tool_nesting_depth = req.nesting_depth.saturating_add(1);
|
||||||
|
|
||||||
|
// Emit tool_use SSE event with redacted parameters to avoid leaking
|
||||||
|
// sensitive data (API keys, passwords, PII) to the web UI.
|
||||||
|
if let Some(ref tx) = state.job_event_tx {
|
||||||
|
let redacted_params = serde_json::json!({
|
||||||
|
"_note": "parameters redacted for security"
|
||||||
|
});
|
||||||
|
let _ = tx.send((
|
||||||
|
job_id,
|
||||||
|
SseEvent::JobToolUse {
|
||||||
|
job_id: job_id.to_string(),
|
||||||
|
tool_name: req.tool_name.clone(),
|
||||||
|
input: redacted_params,
|
||||||
|
},
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Determine timeout override
|
||||||
|
let timeout_override = req
|
||||||
|
.timeout_secs
|
||||||
|
.map(|s| std::time::Duration::from_secs(s.min(300)));
|
||||||
|
|
||||||
|
// Execute the tool
|
||||||
|
match executor
|
||||||
|
.execute(&req.tool_name, req.parameters, &ctx, timeout_override)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(result) => {
|
||||||
|
// Emit tool_result SSE event
|
||||||
|
if let Some(ref tx) = state.job_event_tx {
|
||||||
|
let _ = tx.send((
|
||||||
|
job_id,
|
||||||
|
SseEvent::JobToolResult {
|
||||||
|
job_id: job_id.to_string(),
|
||||||
|
tool_name: req.tool_name.clone(),
|
||||||
|
output: result.output.clone(),
|
||||||
|
},
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Json(ToolCallResponse {
|
||||||
|
success: true,
|
||||||
|
output: Some(result.output),
|
||||||
|
error: None,
|
||||||
|
duration_ms: result.duration.as_millis() as u64,
|
||||||
|
was_sanitized: result.was_sanitized,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(
|
||||||
|
job_id = %job_id,
|
||||||
|
tool = %req.tool_name,
|
||||||
|
error = %e,
|
||||||
|
"PTC tool call failed"
|
||||||
|
);
|
||||||
|
|
||||||
|
Ok(Json(ToolCallResponse {
|
||||||
|
success: false,
|
||||||
|
output: None,
|
||||||
|
error: Some(e.to_string()),
|
||||||
|
duration_ms: 0,
|
||||||
|
was_sanitized: false,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn format_finish_reason(reason: crate::llm::FinishReason) -> String {
|
fn format_finish_reason(reason: crate::llm::FinishReason) -> String {
|
||||||
match reason {
|
match reason {
|
||||||
crate::llm::FinishReason::Stop => "stop".to_string(),
|
crate::llm::FinishReason::Stop => "stop".to_string(),
|
||||||
@@ -480,6 +592,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -709,6 +822,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: Some(secrets_store),
|
secrets_store: Some(secrets_store),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let router = OrchestratorApi::router(state);
|
let router = OrchestratorApi::router(state);
|
||||||
@@ -744,6 +858,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -799,6 +914,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -847,6 +963,7 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -926,4 +1043,425 @@ mod tests {
|
|||||||
assert_eq!(handle.worker_iteration, 5);
|
assert_eq!(handle.worker_iteration, 5);
|
||||||
assert_eq!(handle.last_worker_status.as_deref(), Some("Iteration 5"));
|
assert_eq!(handle.last_worker_status.as_deref(), Some("Iteration 5"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// -- Programmatic tool calling (PTC) tests --
|
||||||
|
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use crate::config::SafetyConfig;
|
||||||
|
use crate::context::JobContext;
|
||||||
|
use crate::safety::SafetyLayer;
|
||||||
|
use crate::tools::{Tool, ToolError, ToolExecutor, ToolOutput, ToolRegistry};
|
||||||
|
|
||||||
|
/// A tool that sleeps for 10 seconds (used to test timeout enforcement).
|
||||||
|
struct SlowTool;
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl Tool for SlowTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"slow_tool"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"A tool that sleeps"
|
||||||
|
}
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
serde_json::json!({"type": "object"})
|
||||||
|
}
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
_params: serde_json::Value,
|
||||||
|
_ctx: &JobContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
tokio::time::sleep(Duration::from_secs(10)).await;
|
||||||
|
Ok(ToolOutput::text("done", Duration::from_secs(10)))
|
||||||
|
}
|
||||||
|
fn requires_sanitization(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build an `OrchestratorState` with a real `ToolExecutor` wired in.
|
||||||
|
///
|
||||||
|
/// Also returns the broadcast receiver when `with_broadcast` is true,
|
||||||
|
/// so SSE-related tests can observe emitted events.
|
||||||
|
fn test_state_with_executor(
|
||||||
|
with_broadcast: bool,
|
||||||
|
) -> (
|
||||||
|
OrchestratorState,
|
||||||
|
Option<broadcast::Receiver<(Uuid, SseEvent)>>,
|
||||||
|
) {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
}));
|
||||||
|
|
||||||
|
let executor = ToolExecutor::new(Arc::clone(&tools), safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let token_store = TokenStore::new();
|
||||||
|
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
|
||||||
|
|
||||||
|
let (tx, rx) = if with_broadcast {
|
||||||
|
let (tx, rx) = broadcast::channel(16);
|
||||||
|
(Some(tx), Some(rx))
|
||||||
|
} else {
|
||||||
|
(None, None)
|
||||||
|
};
|
||||||
|
|
||||||
|
let state = OrchestratorState {
|
||||||
|
llm: Arc::new(StubLlm::default()),
|
||||||
|
job_manager: Arc::new(jm),
|
||||||
|
token_store,
|
||||||
|
job_event_tx: tx,
|
||||||
|
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
|
||||||
|
store: None,
|
||||||
|
secrets_store: None,
|
||||||
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: Some(Arc::new(executor)),
|
||||||
|
};
|
||||||
|
|
||||||
|
(state, rx)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_echo_success() {
|
||||||
|
let (state, _) = test_state_with_executor(false);
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = state.token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "echo",
|
||||||
|
"parameters": {"message": "hello"},
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
|
||||||
|
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(json["success"], true);
|
||||||
|
assert!(
|
||||||
|
json["output"]
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.contains("hello"))
|
||||||
|
.unwrap_or(false),
|
||||||
|
"output should contain 'hello', got: {:?}",
|
||||||
|
json["output"]
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
json["duration_ms"].is_u64(),
|
||||||
|
"duration_ms should be present as a number"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_not_found() {
|
||||||
|
let (state, _) = test_state_with_executor(false);
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = state.token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "nonexistent_tool",
|
||||||
|
"parameters": {},
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
// Handler returns Ok(Json(...)) even on tool failure
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
|
||||||
|
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(json["success"], false);
|
||||||
|
assert!(
|
||||||
|
json["error"]
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.to_lowercase().contains("not found"))
|
||||||
|
.unwrap_or(false),
|
||||||
|
"error should mention 'not found', got: {:?}",
|
||||||
|
json["error"]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_no_executor() {
|
||||||
|
// Use regular test_state() which has tool_executor: None
|
||||||
|
let state = test_state();
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = state.token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "echo",
|
||||||
|
"parameters": {"message": "hello"},
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_with_sse_events() {
|
||||||
|
let (state, rx) = test_state_with_executor(true);
|
||||||
|
let mut rx = rx.expect("broadcast receiver should be present");
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = state.token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "echo",
|
||||||
|
"parameters": {"message": "hello"},
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
// Collect events from broadcast channel
|
||||||
|
let mut saw_tool_use = false;
|
||||||
|
let mut saw_tool_result = false;
|
||||||
|
while let Ok((recv_id, event)) = rx.try_recv() {
|
||||||
|
assert_eq!(recv_id, job_id);
|
||||||
|
match event {
|
||||||
|
SseEvent::JobToolUse { tool_name, .. } => {
|
||||||
|
assert_eq!(tool_name, "echo");
|
||||||
|
saw_tool_use = true;
|
||||||
|
}
|
||||||
|
SseEvent::JobToolResult { tool_name, .. } => {
|
||||||
|
assert_eq!(tool_name, "echo");
|
||||||
|
saw_tool_result = true;
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert!(saw_tool_use, "should have emitted JobToolUse event");
|
||||||
|
assert!(saw_tool_result, "should have emitted JobToolResult event");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_with_timeout() {
|
||||||
|
// Build a registry that includes our SlowTool
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register(Arc::new(SlowTool)).await;
|
||||||
|
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
}));
|
||||||
|
|
||||||
|
let executor = ToolExecutor::new(Arc::clone(&tools), safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let token_store = TokenStore::new();
|
||||||
|
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
|
||||||
|
|
||||||
|
let state = OrchestratorState {
|
||||||
|
llm: Arc::new(StubLlm::default()),
|
||||||
|
job_manager: Arc::new(jm),
|
||||||
|
token_store: token_store.clone(),
|
||||||
|
job_event_tx: None,
|
||||||
|
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
|
||||||
|
store: None,
|
||||||
|
secrets_store: None,
|
||||||
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: Some(Arc::new(executor)),
|
||||||
|
};
|
||||||
|
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "slow_tool",
|
||||||
|
"parameters": {},
|
||||||
|
"timeout_secs": 1,
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
|
||||||
|
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(json["success"], false);
|
||||||
|
assert!(
|
||||||
|
json["error"]
|
||||||
|
.as_str()
|
||||||
|
.map(|s| {
|
||||||
|
let lower = s.to_lowercase();
|
||||||
|
lower.contains("timed out") || lower.contains("timeout")
|
||||||
|
})
|
||||||
|
.unwrap_or(false),
|
||||||
|
"error should mention timeout, got: {:?}",
|
||||||
|
json["error"]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_auth_required() {
|
||||||
|
let (state, _) = test_state_with_executor(false);
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
// Do NOT create a token -- request should be rejected
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "echo",
|
||||||
|
"parameters": {"message": "hello"},
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
// No Authorization header
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn tool_call_nesting_depth_incremented_server_side() {
|
||||||
|
// A worker sending nesting_depth=4 should get depth=5 after the
|
||||||
|
// orchestrator increments it. With MAX_NESTING_DEPTH=5, this
|
||||||
|
// should be rejected (depth >= max).
|
||||||
|
let (state, _) = test_state_with_executor(false);
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = state.token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"tool_name": "echo",
|
||||||
|
"parameters": {"message": "hello"},
|
||||||
|
"nesting_depth": 4,
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/tools/call", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
|
||||||
|
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||||
|
assert_eq!(json["success"], false);
|
||||||
|
assert!(
|
||||||
|
json["error"]
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.to_lowercase().contains("nesting"))
|
||||||
|
.unwrap_or(false),
|
||||||
|
"error should mention nesting depth, got: {:?}",
|
||||||
|
json["error"]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn job_event_tool_use_redacts_input() {
|
||||||
|
// Worker-reported tool_use events must have their input redacted
|
||||||
|
// before SSE broadcast to prevent leaking sensitive parameters.
|
||||||
|
let (tx, mut rx) = broadcast::channel(16);
|
||||||
|
let token_store = TokenStore::new();
|
||||||
|
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
|
||||||
|
let state = OrchestratorState {
|
||||||
|
llm: Arc::new(StubLlm::default()),
|
||||||
|
job_manager: Arc::new(jm),
|
||||||
|
token_store: token_store.clone(),
|
||||||
|
job_event_tx: Some(tx),
|
||||||
|
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
|
||||||
|
store: None,
|
||||||
|
secrets_store: None,
|
||||||
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let job_id = Uuid::new_v4();
|
||||||
|
let token = token_store.create_token(job_id).await;
|
||||||
|
let router = OrchestratorApi::router(state);
|
||||||
|
|
||||||
|
// Worker sends a tool_use event with sensitive data in input
|
||||||
|
let payload = serde_json::json!({
|
||||||
|
"event_type": "tool_use",
|
||||||
|
"data": {
|
||||||
|
"tool_name": "shell",
|
||||||
|
"input": {"command": "curl -H 'Authorization: Bearer sk-secret-key' https://api.example.com"}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let req = Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri(format!("/worker/{}/event", job_id))
|
||||||
|
.header("Authorization", format!("Bearer {}", token))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let resp = router.oneshot(req).await.unwrap();
|
||||||
|
assert_eq!(resp.status(), StatusCode::OK);
|
||||||
|
|
||||||
|
let (_recv_id, event) = rx.recv().await.unwrap();
|
||||||
|
match event {
|
||||||
|
SseEvent::JobToolUse {
|
||||||
|
tool_name, input, ..
|
||||||
|
} => {
|
||||||
|
assert_eq!(tool_name, "shell");
|
||||||
|
// The input must be redacted, not the raw worker payload
|
||||||
|
assert!(
|
||||||
|
input.get("_note").is_some(),
|
||||||
|
"input should be redacted placeholder, got: {}",
|
||||||
|
input
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!input.to_string().contains("sk-secret-key"),
|
||||||
|
"input must not contain sensitive data"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
other => panic!("Expected JobToolUse, got {:?}", other),
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -49,7 +49,9 @@ use uuid::Uuid;
|
|||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::SseEvent;
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::LlmProvider;
|
use crate::llm::LlmProvider;
|
||||||
|
use crate::safety::SafetyLayer;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
|
use crate::tools::{ToolExecutor, ToolRegistry};
|
||||||
|
|
||||||
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
|
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
|
||||||
/// variable, falling back to 50051.
|
/// variable, falling back to 50051.
|
||||||
@@ -75,6 +77,8 @@ pub async fn setup_orchestrator(
|
|||||||
llm: &Arc<dyn LlmProvider>,
|
llm: &Arc<dyn LlmProvider>,
|
||||||
db: Option<&Arc<dyn Database>>,
|
db: Option<&Arc<dyn Database>>,
|
||||||
secrets_store: Option<&Arc<dyn SecretsStore + Send + Sync>>,
|
secrets_store: Option<&Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
|
tools: &Arc<ToolRegistry>,
|
||||||
|
safety: &Arc<SafetyLayer>,
|
||||||
) -> OrchestratorSetup {
|
) -> OrchestratorSetup {
|
||||||
let prompt_queue = Arc::new(Mutex::new(
|
let prompt_queue = Arc::new(Mutex::new(
|
||||||
HashMap::<Uuid, VecDeque<api::PendingPrompt>>::new(),
|
HashMap::<Uuid, VecDeque<api::PendingPrompt>>::new(),
|
||||||
@@ -125,6 +129,17 @@ pub async fn setup_orchestrator(
|
|||||||
};
|
};
|
||||||
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
|
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
|
||||||
|
|
||||||
|
// Build ToolExecutor for programmatic tool calling (PTC)
|
||||||
|
let tool_executor = Arc::new(ToolExecutor::new(
|
||||||
|
Arc::clone(tools),
|
||||||
|
Arc::clone(safety),
|
||||||
|
std::time::Duration::from_secs(60),
|
||||||
|
));
|
||||||
|
|
||||||
|
// Wire the executor into the shared slot so WASM tools registered
|
||||||
|
// during build_all() can resolve it lazily at execution time.
|
||||||
|
tools.set_tool_executor(Arc::clone(&tool_executor));
|
||||||
|
|
||||||
let orchestrator_state = api::OrchestratorState {
|
let orchestrator_state = api::OrchestratorState {
|
||||||
llm: Arc::clone(llm),
|
llm: Arc::clone(llm),
|
||||||
job_manager: Arc::clone(&jm),
|
job_manager: Arc::clone(&jm),
|
||||||
@@ -134,6 +149,7 @@ pub async fn setup_orchestrator(
|
|||||||
store: db.cloned(),
|
store: db.cloned(),
|
||||||
secrets_store: secrets_store.cloned(),
|
secrets_store: secrets_store.cloned(),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
|
tool_executor: Some(tool_executor),
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
|
|||||||
+102
-205
@@ -1078,40 +1078,23 @@ impl SetupWizard {
|
|||||||
.map(|s| s.display_name().to_string())
|
.map(|s| s.display_name().to_string())
|
||||||
.unwrap_or_else(|| def.id.clone())
|
.unwrap_or_else(|| def.id.clone())
|
||||||
} else {
|
} else {
|
||||||
match current.as_str() {
|
current.clone()
|
||||||
"nearai" => "NEAR AI".to_string(),
|
|
||||||
"gemini_oauth" | "gemini-oauth" => "Gemini API (OAuth)".to_string(),
|
|
||||||
_ => {
|
|
||||||
if let Some(def) = registry.find(¤t) {
|
|
||||||
def.setup
|
|
||||||
.as_ref()
|
|
||||||
.map(|s| s.display_name().to_string())
|
|
||||||
.unwrap_or_else(|| def.id.clone())
|
|
||||||
} else {
|
|
||||||
current.clone()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
print_info(&format!("Current provider: {}", display));
|
print_info(&format!("Current provider: {}", display));
|
||||||
println!();
|
println!();
|
||||||
|
|
||||||
let is_known = current == "nearai"
|
let is_known = current == "nearai"
|
||||||
|| current == "bedrock"
|
|| current == "bedrock"
|
||||||
|| current == "gemini_oauth"
|
|
||||||
|| current == "gemini-oauth"
|
|
||||||
|| current == "openai_codex"
|
|| current == "openai_codex"
|
||||||
|| registry.is_known(¤t);
|
|| registry.is_known(¤t);
|
||||||
|
|
||||||
if is_known && confirm("Keep current provider?", true).map_err(SetupError::Io)? {
|
if is_known && confirm("Keep current provider?", true).map_err(SetupError::Io)? {
|
||||||
if current == "bedrock" {
|
if current == "bedrock" {
|
||||||
|
// Keeping the existing Bedrock config — no need to re-run
|
||||||
|
// the full setup flow (region, auth, cross-region).
|
||||||
print_info("Keeping existing AWS Bedrock configuration.");
|
print_info("Keeping existing AWS Bedrock configuration.");
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
if current == "gemini_oauth" || current == "gemini-oauth" {
|
|
||||||
print_info("Keeping existing Gemini CLI OAuth configuration.");
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
if current == "openai_codex" {
|
if current == "openai_codex" {
|
||||||
print_info("Keeping existing OpenAI Codex configuration.");
|
print_info("Keeping existing OpenAI Codex configuration.");
|
||||||
return Ok(());
|
return Ok(());
|
||||||
@@ -1130,15 +1113,13 @@ impl SetupWizard {
|
|||||||
print_info("Select your inference provider:");
|
print_info("Select your inference provider:");
|
||||||
println!();
|
println!();
|
||||||
|
|
||||||
// Build menu: NearAI first, then Gemini OAuth, then OpenAI Codex, then registry providers, then Bedrock
|
// Build menu: NearAI first, then OpenAI Codex, then registry providers, then Bedrock
|
||||||
let selectable = registry.selectable();
|
let selectable = registry.selectable();
|
||||||
let mut options: Vec<String> = Vec::with_capacity(3 + selectable.len());
|
let mut options: Vec<String> = Vec::with_capacity(2 + selectable.len());
|
||||||
let mut provider_ids: Vec<String> = Vec::with_capacity(3 + selectable.len());
|
let mut provider_ids: Vec<String> = Vec::with_capacity(2 + selectable.len());
|
||||||
|
|
||||||
options.push("NEAR AI - multi-model access via NEAR account".to_string());
|
options.push("NEAR AI - multi-model access via NEAR account".to_string());
|
||||||
provider_ids.push("nearai".to_string());
|
provider_ids.push("nearai".to_string());
|
||||||
options.push("Gemini CLI - Official Gemini API via Gemini CLI OAuth".to_string());
|
|
||||||
provider_ids.push("gemini_oauth".to_string());
|
|
||||||
|
|
||||||
options.push("OpenAI Codex - ChatGPT subscription (Plus/Pro/Max)".to_string());
|
options.push("OpenAI Codex - ChatGPT subscription (Plus/Pro/Max)".to_string());
|
||||||
provider_ids.push("openai_codex".to_string());
|
provider_ids.push("openai_codex".to_string());
|
||||||
@@ -1166,8 +1147,6 @@ impl SetupWizard {
|
|||||||
|
|
||||||
if selected_id == "bedrock" {
|
if selected_id == "bedrock" {
|
||||||
self.setup_bedrock().await?;
|
self.setup_bedrock().await?;
|
||||||
} else if selected_id == "gemini_oauth" {
|
|
||||||
self.setup_gemini_oauth().await?;
|
|
||||||
} else {
|
} else {
|
||||||
self.run_provider_setup(selected_id, ®istry).await?;
|
self.run_provider_setup(selected_id, ®istry).await?;
|
||||||
}
|
}
|
||||||
@@ -1816,40 +1795,6 @@ impl SetupWizard {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn setup_gemini_oauth(&mut self) -> Result<(), SetupError> {
|
|
||||||
self.settings.llm_backend = Some("gemini_oauth".to_string());
|
|
||||||
print_info("Starting Gemini CLI OAuth authentication...");
|
|
||||||
println!();
|
|
||||||
|
|
||||||
let creds_path = crate::config::GeminiOauthConfig::default_credentials_path();
|
|
||||||
let cred_manager =
|
|
||||||
crate::llm::gemini_oauth::CredentialManager::new(&creds_path).map_err(|e| {
|
|
||||||
SetupError::Config(format!(
|
|
||||||
"Failed to initialize Gemini credential manager: {}",
|
|
||||||
e
|
|
||||||
))
|
|
||||||
})?;
|
|
||||||
|
|
||||||
match cred_manager.get_valid_credential().await {
|
|
||||||
Ok(cred) => {
|
|
||||||
print_success("Gemini CLI authentication successful!");
|
|
||||||
if let Some(ref pid) = cred.project_id {
|
|
||||||
print_info(&format!("Cloud Code project: {}", pid));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
return Err(SetupError::Config(format!(
|
|
||||||
"Gemini CLI authentication failed: {}. Please try again.",
|
|
||||||
e
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println!();
|
|
||||||
print_success("Gemini API configured via Gemini CLI");
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Step 4: Model selection.
|
/// Step 4: Model selection.
|
||||||
///
|
///
|
||||||
/// Branches on the selected LLM backend and fetches models from the
|
/// Branches on the selected LLM backend and fetches models from the
|
||||||
@@ -1873,157 +1818,109 @@ impl SetupWizard {
|
|||||||
let backend = self.settings.llm_backend.as_deref().unwrap_or("nearai");
|
let backend = self.settings.llm_backend.as_deref().unwrap_or("nearai");
|
||||||
let registry = crate::llm::ProviderRegistry::load();
|
let registry = crate::llm::ProviderRegistry::load();
|
||||||
|
|
||||||
match backend {
|
if backend == "nearai" {
|
||||||
"nearai" => {
|
// NEAR AI: use existing provider list_models()
|
||||||
// NEAR AI: use existing provider list_models()
|
let fetched = self.fetch_nearai_models().await;
|
||||||
let fetched = self.fetch_nearai_models().await;
|
let models = if fetched.is_empty() {
|
||||||
let models = if fetched.is_empty() {
|
crate::llm::default_models()
|
||||||
crate::llm::default_models()
|
} else {
|
||||||
} else {
|
fetched.iter().map(|m| (m.clone(), m.clone())).collect()
|
||||||
fetched.iter().map(|m| (m.clone(), m.clone())).collect()
|
};
|
||||||
|
self.select_from_model_list(&models)?;
|
||||||
|
} else if let Some(def) = registry.find(backend) {
|
||||||
|
let can_list = def
|
||||||
|
.setup
|
||||||
|
.as_ref()
|
||||||
|
.map(|s| s.can_list_models())
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
if can_list {
|
||||||
|
// Try to fetch models from the provider's /v1/models endpoint
|
||||||
|
let cached_key = self
|
||||||
|
.llm_api_key
|
||||||
|
.as_ref()
|
||||||
|
.map(|k| k.expose_secret().to_string());
|
||||||
|
|
||||||
|
let models = match backend {
|
||||||
|
"anthropic" => fetch_anthropic_models(cached_key.as_deref()).await,
|
||||||
|
"openai" => fetch_openai_models(cached_key.as_deref()).await,
|
||||||
|
"ollama" => {
|
||||||
|
let base_url = self
|
||||||
|
.settings
|
||||||
|
.ollama_base_url
|
||||||
|
.as_deref()
|
||||||
|
.or(def.default_base_url.as_deref())
|
||||||
|
.unwrap_or("http://localhost:11434");
|
||||||
|
let models = fetch_ollama_models(base_url).await;
|
||||||
|
if models.is_empty() {
|
||||||
|
print_info("No models found. Pull one first: ollama pull llama3");
|
||||||
|
}
|
||||||
|
models
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
// Generic OpenAI-compatible model listing
|
||||||
|
let base_url = def.default_base_url.as_deref().unwrap_or("");
|
||||||
|
fetch_openai_compatible_models(base_url, cached_key.as_deref()).await
|
||||||
|
}
|
||||||
};
|
};
|
||||||
self.select_from_model_list(&models)?;
|
|
||||||
}
|
// Apply models_filter from setup hint (e.g., Groq "chat" filters non-chat models)
|
||||||
"gemini_oauth" | "gemini-oauth" => {
|
let models =
|
||||||
let default_models: Vec<(String, String)> = vec![
|
if let Some(filter) = def.setup.as_ref().and_then(|s| s.models_filter()) {
|
||||||
(
|
let filter_lower = filter.to_lowercase();
|
||||||
"gemini-3.1-pro-preview".into(),
|
models
|
||||||
"Gemini 3.1 Pro (Latest, strongest reasoning)".into(),
|
.into_iter()
|
||||||
),
|
.filter(|(id, _)| id.to_lowercase().contains(&filter_lower))
|
||||||
(
|
.collect()
|
||||||
"gemini-3.1-pro-preview-customtools".into(),
|
} else {
|
||||||
"Gemini 3.1 Pro Custom Tools (Enhanced tool use)".into(),
|
models
|
||||||
),
|
};
|
||||||
(
|
|
||||||
"gemini-3-pro-preview".into(),
|
if models.is_empty() {
|
||||||
"Gemini 3 Pro (Preview)".into(),
|
// Fall back to manual entry
|
||||||
),
|
let default = &def.default_model;
|
||||||
(
|
let model_id = input(&format!("Model name (default: {default})"))
|
||||||
"gemini-3-flash-preview".into(),
|
|
||||||
"Gemini 3 Flash (Fast preview with thinking)".into(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"gemini-3.1-flash-lite-preview".into(),
|
|
||||||
"Gemini 3.1 Flash Lite (Preview, lightweight)".into(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"gemini-2.5-pro".into(),
|
|
||||||
"Gemini 2.5 Pro (Stable, strong reasoning)".into(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"gemini-2.5-flash".into(),
|
|
||||||
"Gemini 2.5 Flash (Fast, good quality)".into(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"gemini-2.5-flash-lite".into(),
|
|
||||||
"Gemini 2.5 Flash Lite (Fastest, lightweight)".into(),
|
|
||||||
),
|
|
||||||
];
|
|
||||||
self.select_from_model_list(&default_models)?;
|
|
||||||
}
|
|
||||||
"bedrock" => {
|
|
||||||
let model_id =
|
|
||||||
input("Bedrock model ID (e.g., anthropic.claude-v3-sonnet-20240229-v1:0)")
|
|
||||||
.map_err(SetupError::Io)?;
|
.map_err(SetupError::Io)?;
|
||||||
if model_id.is_empty() {
|
let model_id = if model_id.is_empty() {
|
||||||
return Err(SetupError::Config("Model ID is required".to_string()));
|
default.clone()
|
||||||
|
} else {
|
||||||
|
model_id
|
||||||
|
};
|
||||||
|
self.settings.selected_model = Some(model_id.clone());
|
||||||
|
print_success(&format!("Selected {}", model_id));
|
||||||
|
} else {
|
||||||
|
self.select_from_model_list(&models)?;
|
||||||
}
|
}
|
||||||
|
} else {
|
||||||
|
// Manual model entry
|
||||||
|
let default = &def.default_model;
|
||||||
|
let model_id =
|
||||||
|
input(&format!("Model name (default: {default})")).map_err(SetupError::Io)?;
|
||||||
|
let model_id = if model_id.is_empty() {
|
||||||
|
default.clone()
|
||||||
|
} else {
|
||||||
|
model_id
|
||||||
|
};
|
||||||
self.settings.selected_model = Some(model_id.clone());
|
self.settings.selected_model = Some(model_id.clone());
|
||||||
print_success(&format!("Selected {}", model_id));
|
print_success(&format!("Selected {}", model_id));
|
||||||
}
|
}
|
||||||
_ => {
|
} else if backend == "bedrock" {
|
||||||
if let Some(def) = registry.find(backend) {
|
let model_id = input("Bedrock model ID (e.g., anthropic.claude-opus-4-6-v1)")
|
||||||
let can_list = def
|
.map_err(SetupError::Io)?;
|
||||||
.setup
|
if model_id.is_empty() {
|
||||||
.as_ref()
|
return Err(SetupError::Config("Model ID is required".to_string()));
|
||||||
.map(|s| s.can_list_models())
|
|
||||||
.unwrap_or(false);
|
|
||||||
|
|
||||||
if can_list {
|
|
||||||
// Try to fetch models from the provider's /v1/models endpoint
|
|
||||||
let cached_key = self
|
|
||||||
.llm_api_key
|
|
||||||
.as_ref()
|
|
||||||
.map(|k| k.expose_secret().to_string());
|
|
||||||
|
|
||||||
let models = match backend {
|
|
||||||
"anthropic" => fetch_anthropic_models(cached_key.as_deref()).await,
|
|
||||||
"openai" => fetch_openai_models(cached_key.as_deref()).await,
|
|
||||||
"ollama" => {
|
|
||||||
let base_url = self
|
|
||||||
.settings
|
|
||||||
.ollama_base_url
|
|
||||||
.as_deref()
|
|
||||||
.or(def.default_base_url.as_deref())
|
|
||||||
.unwrap_or("http://localhost:11434");
|
|
||||||
let models = fetch_ollama_models(base_url).await;
|
|
||||||
if models.is_empty() {
|
|
||||||
print_info(
|
|
||||||
"No models found. Pull one first: ollama pull llama3",
|
|
||||||
);
|
|
||||||
}
|
|
||||||
models
|
|
||||||
}
|
|
||||||
_ => {
|
|
||||||
// Generic OpenAI-compatible model listing
|
|
||||||
let base_url = def.default_base_url.as_deref().unwrap_or("");
|
|
||||||
fetch_openai_compatible_models(base_url, cached_key.as_deref())
|
|
||||||
.await
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// Apply models_filter from setup hint
|
|
||||||
let models = if let Some(filter) =
|
|
||||||
def.setup.as_ref().and_then(|s| s.models_filter())
|
|
||||||
{
|
|
||||||
let filter_lower = filter.to_lowercase();
|
|
||||||
models
|
|
||||||
.into_iter()
|
|
||||||
.filter(|(id, _)| id.to_lowercase().contains(&filter_lower))
|
|
||||||
.collect()
|
|
||||||
} else {
|
|
||||||
models
|
|
||||||
};
|
|
||||||
|
|
||||||
if models.is_empty() {
|
|
||||||
// Fall back to manual entry
|
|
||||||
let default = &def.default_model;
|
|
||||||
let model_id = input(&format!("Model name (default: {default})"))
|
|
||||||
.map_err(SetupError::Io)?;
|
|
||||||
let model_id = if model_id.is_empty() {
|
|
||||||
default.clone()
|
|
||||||
} else {
|
|
||||||
model_id
|
|
||||||
};
|
|
||||||
self.settings.selected_model = Some(model_id.clone());
|
|
||||||
print_success(&format!("Selected {}", model_id));
|
|
||||||
} else {
|
|
||||||
self.select_from_model_list(&models)?;
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
// Manual model entry
|
|
||||||
let default = &def.default_model;
|
|
||||||
let model_id = input(&format!("Model name (default: {default})"))
|
|
||||||
.map_err(SetupError::Io)?;
|
|
||||||
let model_id = if model_id.is_empty() {
|
|
||||||
default.clone()
|
|
||||||
} else {
|
|
||||||
model_id
|
|
||||||
};
|
|
||||||
self.settings.selected_model = Some(model_id.clone());
|
|
||||||
print_success(&format!("Selected {}", model_id));
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
// Unknown provider, manual entry
|
|
||||||
let model_id = input("Model name (e.g., meta-llama/Llama-3-8b-chat-hf)")
|
|
||||||
.map_err(SetupError::Io)?;
|
|
||||||
if model_id.is_empty() {
|
|
||||||
return Err(SetupError::Config("Model name is required".to_string()));
|
|
||||||
}
|
|
||||||
self.settings.selected_model = Some(model_id.clone());
|
|
||||||
print_success(&format!("Selected {}", model_id));
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
self.settings.selected_model = Some(model_id.clone());
|
||||||
|
print_success(&format!("Selected {}", model_id));
|
||||||
|
} else {
|
||||||
|
// Unknown provider, manual entry
|
||||||
|
let model_id = input("Model name (e.g., meta-llama/Llama-3-8b-chat-hf)")
|
||||||
|
.map_err(SetupError::Io)?;
|
||||||
|
if model_id.is_empty() {
|
||||||
|
return Err(SetupError::Config("Model name is required".to_string()));
|
||||||
|
}
|
||||||
|
self.settings.selected_model = Some(model_id.clone());
|
||||||
|
print_success(&format!("Selected {}", model_id));
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ mod json;
|
|||||||
mod memory;
|
mod memory;
|
||||||
mod message;
|
mod message;
|
||||||
pub mod path_utils;
|
pub mod path_utils;
|
||||||
|
pub mod ptc_script;
|
||||||
mod restart;
|
mod restart;
|
||||||
pub mod routine;
|
pub mod routine;
|
||||||
pub mod secrets_tools;
|
pub mod secrets_tools;
|
||||||
@@ -31,6 +32,7 @@ pub use job::{
|
|||||||
pub use json::JsonTool;
|
pub use json::JsonTool;
|
||||||
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
|
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
|
||||||
pub use message::MessageTool;
|
pub use message::MessageTool;
|
||||||
|
pub use ptc_script::PtcScriptTool;
|
||||||
pub use restart::RestartTool;
|
pub use restart::RestartTool;
|
||||||
pub use routine::{
|
pub use routine::{
|
||||||
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
|
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
|
||||||
|
|||||||
@@ -0,0 +1,371 @@
|
|||||||
|
//! PTC script tool for running multi-step Python programs that call tools.
|
||||||
|
//!
|
||||||
|
//! Wraps user-provided Python code in a preamble that imports the IronClaw
|
||||||
|
//! SDK (`ironclaw_tools`), then executes it via `python3 -c`. The script
|
||||||
|
//! runs in the same environment as the worker container and can call any
|
||||||
|
//! registered tool through the SDK's `call_tool()` function.
|
||||||
|
|
||||||
|
use std::process::Stdio;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use tokio::io::AsyncReadExt;
|
||||||
|
use tokio::process::Command;
|
||||||
|
|
||||||
|
use crate::context::JobContext;
|
||||||
|
use crate::tools::tool::{
|
||||||
|
ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// Maximum output size before truncation (64KB).
|
||||||
|
const MAX_OUTPUT_SIZE: usize = 64 * 1024;
|
||||||
|
|
||||||
|
/// Default script timeout.
|
||||||
|
const DEFAULT_TIMEOUT_SECS: u64 = 120;
|
||||||
|
|
||||||
|
/// Maximum allowed timeout.
|
||||||
|
const MAX_TIMEOUT_SECS: u64 = 300;
|
||||||
|
|
||||||
|
/// Environment variables safe to forward to the Python subprocess.
|
||||||
|
const SAFE_ENV_VARS: &[&str] = &[
|
||||||
|
"PATH",
|
||||||
|
"HOME",
|
||||||
|
"USER",
|
||||||
|
"LOGNAME",
|
||||||
|
"SHELL",
|
||||||
|
"TERM",
|
||||||
|
"LANG",
|
||||||
|
"LC_ALL",
|
||||||
|
"LC_CTYPE",
|
||||||
|
"PWD",
|
||||||
|
"TMPDIR",
|
||||||
|
"TMP",
|
||||||
|
"TEMP",
|
||||||
|
"CARGO_HOME",
|
||||||
|
"RUSTUP_HOME",
|
||||||
|
"NODE_PATH",
|
||||||
|
"NPM_CONFIG_PREFIX",
|
||||||
|
];
|
||||||
|
|
||||||
|
/// PTC environment variables required by the ironclaw_tools SDK.
|
||||||
|
const PTC_ENV_VARS: &[&str] = &[
|
||||||
|
"IRONCLAW_ORCHESTRATOR_URL",
|
||||||
|
"IRONCLAW_JOB_ID",
|
||||||
|
"IRONCLAW_WORKER_TOKEN",
|
||||||
|
];
|
||||||
|
|
||||||
|
/// Python preamble injected before the user's script.
|
||||||
|
const PREAMBLE: &str = r#"
|
||||||
|
import json, sys, os
|
||||||
|
|
||||||
|
# Import IronClaw SDK
|
||||||
|
from ironclaw_tools import call_tool, shell, read_file, write_file, http_get
|
||||||
|
|
||||||
|
# Structured output collector
|
||||||
|
_ptc_outputs = {}
|
||||||
|
|
||||||
|
def ptc_output(key, value):
|
||||||
|
"""Register a named output value for structured results."""
|
||||||
|
_ptc_outputs[key] = value
|
||||||
|
|
||||||
|
try:
|
||||||
|
"#;
|
||||||
|
|
||||||
|
/// Python postamble appended after the user's script.
|
||||||
|
const POSTAMBLE: &str = r#"
|
||||||
|
except Exception as _ptc_err:
|
||||||
|
print(f"SCRIPT_ERROR: {type(_ptc_err).__name__}: {_ptc_err}", file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
# Print structured outputs if any were registered
|
||||||
|
if _ptc_outputs:
|
||||||
|
print("\n__PTC_OUTPUTS__")
|
||||||
|
print(json.dumps(_ptc_outputs))
|
||||||
|
"#;
|
||||||
|
|
||||||
|
pub struct PtcScriptTool;
|
||||||
|
|
||||||
|
impl Default for PtcScriptTool {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PtcScriptTool {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build the full Python program from user script + preamble/postamble.
|
||||||
|
fn build_program(script: &str) -> String {
|
||||||
|
let mut program =
|
||||||
|
String::with_capacity(PREAMBLE.len() + script.len() + POSTAMBLE.len() + 256);
|
||||||
|
program.push_str(PREAMBLE);
|
||||||
|
|
||||||
|
// Indent user script into the try: block
|
||||||
|
for line in script.lines() {
|
||||||
|
program.push_str(" ");
|
||||||
|
program.push_str(line);
|
||||||
|
program.push('\n');
|
||||||
|
}
|
||||||
|
|
||||||
|
program.push_str(POSTAMBLE);
|
||||||
|
program
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Truncate output to MAX_OUTPUT_SIZE with a truncation notice.
|
||||||
|
fn truncate_output(output: &str) -> String {
|
||||||
|
if output.len() <= MAX_OUTPUT_SIZE {
|
||||||
|
output.to_string()
|
||||||
|
} else {
|
||||||
|
let mut i = MAX_OUTPUT_SIZE;
|
||||||
|
while i > 0 && !output.is_char_boundary(i) {
|
||||||
|
i -= 1;
|
||||||
|
}
|
||||||
|
format!(
|
||||||
|
"{}\n\n[Output truncated at {} bytes]",
|
||||||
|
&output[..i],
|
||||||
|
MAX_OUTPUT_SIZE
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for PtcScriptTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"ptc_script"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Execute a Python script that can call IronClaw tools programmatically. \
|
||||||
|
The script has access to call_tool(), shell(), read_file(), write_file(), \
|
||||||
|
and http_get() from the ironclaw_tools SDK. Use ptc_output(key, value) \
|
||||||
|
to return structured results."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"script": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Python script to execute. Has access to call_tool(), shell(), read_file(), write_file(), http_get(), and ptc_output()."
|
||||||
|
},
|
||||||
|
"timeout_secs": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "Timeout in seconds (default 120, max 300).",
|
||||||
|
"default": 120,
|
||||||
|
"minimum": 1,
|
||||||
|
"maximum": 300
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["script"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
params: serde_json::Value,
|
||||||
|
ctx: &JobContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
let start = std::time::Instant::now();
|
||||||
|
|
||||||
|
let script = require_str(¶ms, "script")?;
|
||||||
|
let timeout_secs = params
|
||||||
|
.get("timeout_secs")
|
||||||
|
.and_then(|v| v.as_u64())
|
||||||
|
.unwrap_or(DEFAULT_TIMEOUT_SECS)
|
||||||
|
.min(MAX_TIMEOUT_SECS);
|
||||||
|
let timeout = Duration::from_secs(timeout_secs);
|
||||||
|
|
||||||
|
let program = Self::build_program(script);
|
||||||
|
|
||||||
|
// Build the subprocess command
|
||||||
|
let mut command = Command::new("python3");
|
||||||
|
command.args(["-c", &program]);
|
||||||
|
|
||||||
|
// Scrub environment -- only forward safe vars + PTC vars + extra_env
|
||||||
|
command.env_clear();
|
||||||
|
for var in SAFE_ENV_VARS {
|
||||||
|
if let Ok(val) = std::env::var(var) {
|
||||||
|
command.env(var, val);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for var in PTC_ENV_VARS {
|
||||||
|
if let Ok(val) = std::env::var(var) {
|
||||||
|
command.env(var, val);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Forward extra_env from JobContext (credentials fetched by worker runtime)
|
||||||
|
for (k, v) in ctx.extra_env.iter() {
|
||||||
|
command.env(k, v);
|
||||||
|
}
|
||||||
|
|
||||||
|
command
|
||||||
|
.stdin(Stdio::null())
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped());
|
||||||
|
|
||||||
|
// Spawn and drain stdout/stderr concurrently
|
||||||
|
let mut child = command
|
||||||
|
.spawn()
|
||||||
|
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to spawn python3: {}", e)))?;
|
||||||
|
|
||||||
|
let stdout_handle = child.stdout.take();
|
||||||
|
let stderr_handle = child.stderr.take();
|
||||||
|
|
||||||
|
let result = tokio::time::timeout(timeout, async {
|
||||||
|
let stdout_fut = async {
|
||||||
|
if let Some(mut out) = stdout_handle {
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
(&mut out)
|
||||||
|
.take(MAX_OUTPUT_SIZE as u64)
|
||||||
|
.read_to_end(&mut buf)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
tokio::io::copy(&mut out, &mut tokio::io::sink()).await.ok();
|
||||||
|
String::from_utf8_lossy(&buf).to_string()
|
||||||
|
} else {
|
||||||
|
String::new()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let stderr_fut = async {
|
||||||
|
if let Some(mut err) = stderr_handle {
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
(&mut err)
|
||||||
|
.take(MAX_OUTPUT_SIZE as u64)
|
||||||
|
.read_to_end(&mut buf)
|
||||||
|
.await
|
||||||
|
.ok();
|
||||||
|
tokio::io::copy(&mut err, &mut tokio::io::sink()).await.ok();
|
||||||
|
String::from_utf8_lossy(&buf).to_string()
|
||||||
|
} else {
|
||||||
|
String::new()
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let (stdout, stderr, wait_result) = tokio::join!(stdout_fut, stderr_fut, child.wait());
|
||||||
|
let status = wait_result?;
|
||||||
|
Ok::<_, std::io::Error>((stdout, stderr, status.code().unwrap_or(-1)))
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let duration = start.elapsed();
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(Ok((stdout, stderr, exit_code))) => {
|
||||||
|
if exit_code != 0 {
|
||||||
|
let error_msg = if stderr.is_empty() {
|
||||||
|
format!("Script exited with code {}", exit_code)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"Script exited with code {}:\n{}",
|
||||||
|
exit_code,
|
||||||
|
Self::truncate_output(&stderr)
|
||||||
|
)
|
||||||
|
};
|
||||||
|
return Err(ToolError::ExecutionFailed(error_msg));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Combine output
|
||||||
|
let output = if stderr.is_empty() {
|
||||||
|
stdout
|
||||||
|
} else {
|
||||||
|
format!("{}\n\n--- stderr ---\n{}", stdout, stderr)
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(ToolOutput::text(Self::truncate_output(&output), duration))
|
||||||
|
}
|
||||||
|
Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!(
|
||||||
|
"Script execution failed: {}",
|
||||||
|
e
|
||||||
|
))),
|
||||||
|
Err(_) => {
|
||||||
|
let _ = child.kill().await;
|
||||||
|
Err(ToolError::Timeout(timeout))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn requires_sanitization(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||||
|
ApprovalRequirement::Always
|
||||||
|
}
|
||||||
|
|
||||||
|
fn domain(&self) -> ToolDomain {
|
||||||
|
ToolDomain::Container
|
||||||
|
}
|
||||||
|
|
||||||
|
fn execution_timeout(&self) -> Duration {
|
||||||
|
Duration::from_secs(MAX_TIMEOUT_SECS)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_program_indents_script() {
|
||||||
|
let script = "x = 1\nprint(x)";
|
||||||
|
let program = PtcScriptTool::build_program(script);
|
||||||
|
assert!(program.contains(" x = 1\n"));
|
||||||
|
assert!(program.contains(" print(x)\n"));
|
||||||
|
assert!(program.contains("from ironclaw_tools import"));
|
||||||
|
assert!(program.contains("def ptc_output("));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_program_empty_script() {
|
||||||
|
let program = PtcScriptTool::build_program("");
|
||||||
|
// Empty script should still have preamble + postamble
|
||||||
|
assert!(program.contains("try:"));
|
||||||
|
assert!(program.contains("except Exception"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_output() {
|
||||||
|
let short = "hello";
|
||||||
|
assert_eq!(PtcScriptTool::truncate_output(short), "hello");
|
||||||
|
|
||||||
|
let long = "x".repeat(MAX_OUTPUT_SIZE + 100);
|
||||||
|
let truncated = PtcScriptTool::truncate_output(&long);
|
||||||
|
assert!(truncated.len() < long.len());
|
||||||
|
assert!(truncated.contains("[Output truncated"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_truncate_output_multibyte_boundary() {
|
||||||
|
// Build a string of multi-byte chars (emoji = 4 bytes each) that crosses MAX_OUTPUT_SIZE
|
||||||
|
let emoji = "\u{1F600}"; // 4 bytes
|
||||||
|
let count = MAX_OUTPUT_SIZE / emoji.len() + 10;
|
||||||
|
let long: String = emoji.repeat(count);
|
||||||
|
assert!(long.len() > MAX_OUTPUT_SIZE);
|
||||||
|
|
||||||
|
let truncated = PtcScriptTool::truncate_output(&long);
|
||||||
|
// Must not panic and must contain valid UTF-8
|
||||||
|
assert!(truncated.contains("[Output truncated"));
|
||||||
|
// The kept portion must end on a char boundary (valid UTF-8 guaranteed by compilation)
|
||||||
|
let kept = truncated.split("\n\n[Output truncated").next().unwrap();
|
||||||
|
assert!(kept.len() <= MAX_OUTPUT_SIZE);
|
||||||
|
// Every char should be complete (no partial emoji)
|
||||||
|
assert!(kept.chars().all(|c| c == '\u{1F600}'));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_tool_metadata() {
|
||||||
|
let tool = PtcScriptTool::new();
|
||||||
|
assert_eq!(tool.name(), "ptc_script");
|
||||||
|
assert_eq!(tool.domain(), ToolDomain::Container);
|
||||||
|
assert_eq!(
|
||||||
|
tool.requires_approval(&serde_json::json!({})),
|
||||||
|
ApprovalRequirement::Always
|
||||||
|
);
|
||||||
|
assert!(tool.requires_sanitization());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -45,23 +45,11 @@ impl ToolInfoDetail {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn schema_param_names(schema: &serde_json::Value) -> Vec<String> {
|
fn schema_param_names(schema: &serde_json::Value) -> Vec<String> {
|
||||||
let mut names = std::collections::BTreeSet::new();
|
schema
|
||||||
|
.get("properties")
|
||||||
if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
|
.and_then(|p| p.as_object())
|
||||||
names.extend(props.keys().cloned());
|
.map(|props| props.keys().cloned().collect())
|
||||||
}
|
.unwrap_or_default()
|
||||||
|
|
||||||
for key in ["allOf", "oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array()) {
|
|
||||||
for variant in variants {
|
|
||||||
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
|
||||||
names.extend(props.keys().cloned());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
names.into_iter().collect()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn fallback_summary(schema: &serde_json::Value) -> ToolDiscoverySummary {
|
fn fallback_summary(schema: &serde_json::Value) -> ToolDiscoverySummary {
|
||||||
|
|||||||
+12
-701
@@ -1,4 +1,4 @@
|
|||||||
pub fn prepare_tool_params(
|
pub(crate) fn prepare_tool_params(
|
||||||
tool: &dyn crate::tools::tool::Tool,
|
tool: &dyn crate::tools::tool::Tool,
|
||||||
params: &serde_json::Value,
|
params: &serde_json::Value,
|
||||||
) -> serde_json::Value {
|
) -> serde_json::Value {
|
||||||
@@ -9,87 +9,14 @@ pub(crate) fn prepare_params_for_schema(
|
|||||||
params: &serde_json::Value,
|
params: &serde_json::Value,
|
||||||
schema: &serde_json::Value,
|
schema: &serde_json::Value,
|
||||||
) -> serde_json::Value {
|
) -> serde_json::Value {
|
||||||
let resolved = resolve_refs(schema);
|
coerce_value(params, schema)
|
||||||
coerce_value(params, &resolved)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── $ref resolution ──────────────────────────────────────────────────
|
|
||||||
|
|
||||||
/// Inline all `$ref` pointers in a JSON Schema so downstream coercion
|
|
||||||
/// operates on a flat, self-contained schema tree.
|
|
||||||
///
|
|
||||||
/// Supports `#/definitions/<name>` and `#/$defs/<name>` (JSON Schema
|
|
||||||
/// draft-07 and 2020-12 respectively). Unknown `$ref` formats are left
|
|
||||||
/// unchanged. A depth limit prevents infinite recursion from circular refs.
|
|
||||||
fn resolve_refs(schema: &serde_json::Value) -> serde_json::Value {
|
|
||||||
let definitions = schema
|
|
||||||
.get("definitions")
|
|
||||||
.or_else(|| schema.get("$defs"))
|
|
||||||
.cloned()
|
|
||||||
.unwrap_or(serde_json::Value::Null);
|
|
||||||
resolve_refs_inner(schema, &definitions, 0)
|
|
||||||
}
|
|
||||||
|
|
||||||
const MAX_REF_DEPTH: usize = 16;
|
|
||||||
|
|
||||||
fn resolve_refs_inner(
|
|
||||||
schema: &serde_json::Value,
|
|
||||||
definitions: &serde_json::Value,
|
|
||||||
depth: usize,
|
|
||||||
) -> serde_json::Value {
|
|
||||||
if depth > MAX_REF_DEPTH {
|
|
||||||
return schema.clone();
|
|
||||||
}
|
|
||||||
match schema {
|
|
||||||
serde_json::Value::Object(obj) => {
|
|
||||||
// If this node is a $ref, resolve it and recurse into the target.
|
|
||||||
if let Some(ref_str) = obj.get("$ref").and_then(|v| v.as_str()) {
|
|
||||||
if let Some(target) = resolve_ref_pointer(ref_str, definitions) {
|
|
||||||
return resolve_refs_inner(&target, definitions, depth + 1);
|
|
||||||
}
|
|
||||||
return schema.clone();
|
|
||||||
}
|
|
||||||
|
|
||||||
// Recursively resolve refs in all values (skip definitions maps).
|
|
||||||
let resolved: serde_json::Map<String, serde_json::Value> = obj
|
|
||||||
.iter()
|
|
||||||
.map(|(k, v)| {
|
|
||||||
if k == "definitions" || k == "$defs" {
|
|
||||||
(k.clone(), v.clone())
|
|
||||||
} else {
|
|
||||||
(k.clone(), resolve_refs_inner(v, definitions, depth + 1))
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
serde_json::Value::Object(resolved)
|
|
||||||
}
|
|
||||||
serde_json::Value::Array(arr) => serde_json::Value::Array(
|
|
||||||
arr.iter()
|
|
||||||
.map(|v| resolve_refs_inner(v, definitions, depth + 1))
|
|
||||||
.collect(),
|
|
||||||
),
|
|
||||||
_ => schema.clone(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn resolve_ref_pointer(
|
|
||||||
ref_str: &str,
|
|
||||||
definitions: &serde_json::Value,
|
|
||||||
) -> Option<serde_json::Value> {
|
|
||||||
let path = ref_str.strip_prefix("#/")?;
|
|
||||||
let parts: Vec<&str> = path.split('/').collect();
|
|
||||||
if parts.len() == 2 && (parts[0] == "definitions" || parts[0] == "$defs") {
|
|
||||||
return definitions.get(parts[1]).cloned();
|
|
||||||
}
|
|
||||||
None
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Core coercion ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_json::Value {
|
fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_json::Value {
|
||||||
// This coercer handles concrete schema shapes including discriminated unions
|
// This coercer intentionally handles the concrete schema shapes we expose in
|
||||||
// (oneOf/anyOf with const or single-element enum discriminators), allOf
|
// discovery today. It does not resolve combinators like anyOf/oneOf/allOf or
|
||||||
// merges, and $ref references (resolved in a pre-pass).
|
// references via $ref; those schemas pass through unchanged unless they also
|
||||||
|
// advertise a directly coercible type/property shape.
|
||||||
if value.is_null() {
|
if value.is_null() {
|
||||||
return value.clone();
|
return value.clone();
|
||||||
}
|
}
|
||||||
@@ -120,35 +47,12 @@ fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_
|
|||||||
return value.clone();
|
return value.clone();
|
||||||
}
|
}
|
||||||
|
|
||||||
let resolved = resolve_effective_properties(schema, obj);
|
let properties = schema.get("properties").and_then(|p| p.as_object());
|
||||||
let properties = resolved
|
let additional_schema = schema.get("additionalProperties").filter(|v| v.is_object());
|
||||||
.as_ref()
|
|
||||||
.or_else(|| schema.get("properties").and_then(|p| p.as_object()));
|
|
||||||
let additional_schema = schema
|
|
||||||
.get("additionalProperties")
|
|
||||||
.filter(|v| v.is_object())
|
|
||||||
.or_else(|| resolve_additional_properties(schema, obj));
|
|
||||||
let required: std::collections::HashSet<&str> = schema
|
|
||||||
.get("required")
|
|
||||||
.and_then(|r| r.as_array())
|
|
||||||
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
|
|
||||||
.unwrap_or_default();
|
|
||||||
let mut coerced = obj.clone();
|
let mut coerced = obj.clone();
|
||||||
|
|
||||||
for (key, current) in &mut coerced {
|
for (key, current) in &mut coerced {
|
||||||
if let Some(prop_schema) = properties.and_then(|props| props.get(key)) {
|
if let Some(prop_schema) = properties.and_then(|props| props.get(key)) {
|
||||||
// LLMs send "" for optional fields instead of omitting them.
|
|
||||||
// Coerce to null only when the field is not required AND the schema
|
|
||||||
// allows null or doesn't allow string — a `type: "string"` field
|
|
||||||
// may legitimately accept "" as a meaningful value.
|
|
||||||
if current.as_str() == Some("")
|
|
||||||
&& !required.contains(key.as_str())
|
|
||||||
&& (schema_allows_type(prop_schema, "null")
|
|
||||||
|| !schema_allows_type(prop_schema, "string"))
|
|
||||||
{
|
|
||||||
*current = serde_json::Value::Null;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
*current = coerce_value(current, prop_schema);
|
*current = coerce_value(current, prop_schema);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@@ -164,179 +68,11 @@ fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_
|
|||||||
value.clone()
|
value.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// When the schema uses `oneOf`, `anyOf`, or `allOf` combinators, build a
|
|
||||||
/// merged property map that can be used for coercion.
|
|
||||||
///
|
|
||||||
/// - Top-level `properties` are included first (base properties).
|
|
||||||
/// - `allOf`: merge ALL variants' properties (last-wins on conflicts).
|
|
||||||
/// - `oneOf`/`anyOf`: find the discriminated match and merge its properties.
|
|
||||||
///
|
|
||||||
/// Returns `None` if no combinators are present or no match is found, so the
|
|
||||||
/// caller falls back to the existing top-level `properties` lookup.
|
|
||||||
fn resolve_effective_properties(
|
|
||||||
schema: &serde_json::Value,
|
|
||||||
obj: &serde_json::Map<String, serde_json::Value>,
|
|
||||||
) -> Option<serde_json::Map<String, serde_json::Value>> {
|
|
||||||
collect_properties(schema, obj, 0)
|
|
||||||
}
|
|
||||||
|
|
||||||
const MAX_COMBINATOR_DEPTH: usize = 4;
|
|
||||||
|
|
||||||
/// Recursively collect properties from a schema and its combinator variants.
|
|
||||||
fn collect_properties(
|
|
||||||
schema: &serde_json::Value,
|
|
||||||
obj: &serde_json::Map<String, serde_json::Value>,
|
|
||||||
depth: usize,
|
|
||||||
) -> Option<serde_json::Map<String, serde_json::Value>> {
|
|
||||||
if depth > MAX_COMBINATOR_DEPTH {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
let has_combinators = schema.get("allOf").is_some()
|
|
||||||
|| schema.get("oneOf").is_some()
|
|
||||||
|| schema.get("anyOf").is_some();
|
|
||||||
|
|
||||||
if !has_combinators {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut merged = serde_json::Map::new();
|
|
||||||
|
|
||||||
// Start with top-level properties
|
|
||||||
if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
|
|
||||||
merged.extend(props.iter().map(|(k, v)| (k.clone(), v.clone())));
|
|
||||||
}
|
|
||||||
|
|
||||||
// allOf: merge ALL variants' properties, recursing into nested combinators
|
|
||||||
if let Some(all_of) = schema.get("allOf").and_then(|a| a.as_array()) {
|
|
||||||
for variant in all_of {
|
|
||||||
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
|
||||||
merged.extend(props.iter().map(|(k, v)| (k.clone(), v.clone())));
|
|
||||||
}
|
|
||||||
// Recurse into variant if it has its own combinators
|
|
||||||
if let Some(nested) = collect_properties(variant, obj, depth + 1) {
|
|
||||||
merged.extend(nested);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// oneOf/anyOf: find discriminated match and merge its properties
|
|
||||||
for key in ["oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array())
|
|
||||||
&& let Some(variant) = find_discriminated_variant(variants, obj)
|
|
||||||
{
|
|
||||||
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
|
||||||
merged.extend(props.iter().map(|(k, v)| (k.clone(), v.clone())));
|
|
||||||
}
|
|
||||||
// Recurse into matched variant if it has its own combinators
|
|
||||||
if let Some(nested) = collect_properties(variant, obj, depth + 1) {
|
|
||||||
merged.extend(nested);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if merged.is_empty() {
|
|
||||||
None
|
|
||||||
} else {
|
|
||||||
Some(merged)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Find `additionalProperties` from a matched combinator variant.
|
|
||||||
///
|
|
||||||
/// Checks `allOf` variants first (last-wins), then the matched `oneOf`/`anyOf`
|
|
||||||
/// variant. Returns `None` if no variant defines `additionalProperties`.
|
|
||||||
fn resolve_additional_properties<'a>(
|
|
||||||
schema: &'a serde_json::Value,
|
|
||||||
obj: &serde_json::Map<String, serde_json::Value>,
|
|
||||||
) -> Option<&'a serde_json::Value> {
|
|
||||||
// allOf: last variant with additionalProperties wins
|
|
||||||
if let Some(all_of) = schema.get("allOf").and_then(|a| a.as_array()) {
|
|
||||||
for variant in all_of.iter().rev() {
|
|
||||||
if let Some(ap) = variant.get("additionalProperties")
|
|
||||||
&& ap.is_object()
|
|
||||||
{
|
|
||||||
return Some(ap);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// oneOf/anyOf: check matched variant
|
|
||||||
for key in ["oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array())
|
|
||||||
&& let Some(variant) = find_discriminated_variant(variants, obj)
|
|
||||||
&& let Some(ap) = variant.get("additionalProperties")
|
|
||||||
&& ap.is_object()
|
|
||||||
{
|
|
||||||
return Some(ap);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
None
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Find a `oneOf`/`anyOf` variant that matches the given object by checking
|
|
||||||
/// `const`-valued and single-element `enum`-valued properties (discriminators).
|
|
||||||
///
|
|
||||||
/// A variant matches when ALL its discriminator properties match the object's
|
|
||||||
/// values and at least one such discriminator exists. Returns `None` if no
|
|
||||||
/// variant matches (safe fallback — no coercion).
|
|
||||||
fn find_discriminated_variant<'a>(
|
|
||||||
variants: &'a [serde_json::Value],
|
|
||||||
obj: &serde_json::Map<String, serde_json::Value>,
|
|
||||||
) -> Option<&'a serde_json::Value> {
|
|
||||||
variants.iter().find(|variant| {
|
|
||||||
let Some(props) = variant.get("properties").and_then(|p| p.as_object()) else {
|
|
||||||
return false;
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut discriminator_count = 0;
|
|
||||||
|
|
||||||
for (key, prop_schema) in props {
|
|
||||||
// Check for const discriminator
|
|
||||||
if let Some(const_val) = prop_schema.get("const") {
|
|
||||||
discriminator_count += 1;
|
|
||||||
match obj.get(key) {
|
|
||||||
Some(v) if v == const_val => {}
|
|
||||||
_ => return false,
|
|
||||||
}
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Check for single-element enum discriminator
|
|
||||||
if let Some(enum_vals) = prop_schema.get("enum").and_then(|e| e.as_array())
|
|
||||||
&& enum_vals.len() == 1
|
|
||||||
{
|
|
||||||
discriminator_count += 1;
|
|
||||||
match obj.get(key) {
|
|
||||||
Some(v) if v == &enum_vals[0] => {}
|
|
||||||
_ => return false,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
discriminator_count > 0
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn coerce_string_value(s: &str, schema: &serde_json::Value) -> Option<serde_json::Value> {
|
fn coerce_string_value(s: &str, schema: &serde_json::Value) -> Option<serde_json::Value> {
|
||||||
// LLMs often send "" instead of null for optional fields. Coerce empty
|
|
||||||
// strings to null when the schema allows null but not string, or allows
|
|
||||||
// both but the value is empty (a string field with content "" is kept).
|
|
||||||
if s.is_empty() && schema_allows_type(schema, "null") && !schema_allows_type(schema, "string") {
|
|
||||||
return Some(serde_json::Value::Null);
|
|
||||||
}
|
|
||||||
|
|
||||||
if schema_allows_type(schema, "string") {
|
if schema_allows_type(schema, "string") {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Empty string with no type match — return unchanged since we can't
|
|
||||||
// determine the intended type.
|
|
||||||
if s.is_empty() {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
if schema_allows_type(schema, "integer")
|
if schema_allows_type(schema, "integer")
|
||||||
&& let Ok(v) = s.parse::<i64>()
|
&& let Ok(v) = s.parse::<i64>()
|
||||||
{
|
{
|
||||||
@@ -378,15 +114,10 @@ fn schema_allows_type(schema: &serde_json::Value, expected: &str) -> bool {
|
|||||||
Some(serde_json::Value::String(t)) => t == expected,
|
Some(serde_json::Value::String(t)) => t == expected,
|
||||||
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
|
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
|
||||||
_ => match expected {
|
_ => match expected {
|
||||||
"object" => {
|
"object" => schema
|
||||||
schema
|
.get("properties")
|
||||||
.get("properties")
|
.and_then(|p| p.as_object())
|
||||||
.and_then(|p| p.as_object())
|
.is_some(),
|
||||||
.is_some()
|
|
||||||
|| schema.get("oneOf").is_some()
|
|
||||||
|| schema.get("anyOf").is_some()
|
|
||||||
|| schema.get("allOf").is_some()
|
|
||||||
}
|
|
||||||
"array" => schema.get("items").is_some(),
|
"array" => schema.get("items").is_some(),
|
||||||
_ => false,
|
_ => false,
|
||||||
},
|
},
|
||||||
@@ -594,91 +325,6 @@ mod tests {
|
|||||||
assert_eq!(result["value"], serde_json::json!("{\"mode\":\"raw\"}")); // safety: test-only assertion
|
assert_eq!(result["value"], serde_json::json!("{\"mode\":\"raw\"}")); // safety: test-only assertion
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_empty_string_to_null_for_nullable_non_required_field() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"timezone": { "type": ["string", "null"] },
|
|
||||||
"schedule": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["schedule"]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"timezone": "",
|
|
||||||
"schedule": "0 9 * * *"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
// Non-required nullable "timezone" with empty string → null
|
|
||||||
assert_eq!(result["timezone"], serde_json::Value::Null);
|
|
||||||
// Required "schedule" keeps its value even if empty would be weird
|
|
||||||
assert_eq!(result["schedule"], serde_json::json!("0 9 * * *"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn keeps_empty_string_for_non_required_string_only_field() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"timezone": { "type": "string" },
|
|
||||||
"schedule": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["schedule"]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"timezone": "",
|
|
||||||
"schedule": "0 9 * * *"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
// Non-required string-only "timezone" keeps empty string (meaningful value)
|
|
||||||
assert_eq!(result["timezone"], serde_json::json!(""));
|
|
||||||
assert_eq!(result["schedule"], serde_json::json!("0 9 * * *"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_empty_string_to_null_for_explicit_nullable_type() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"from_timezone": { "type": ["string", "null"] },
|
|
||||||
"operation": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["operation"]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"from_timezone": "",
|
|
||||||
"operation": "now"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
// Nullable type with empty string → null (even if it were required,
|
|
||||||
// the per-value coercion in coerce_string_value handles this)
|
|
||||||
assert_eq!(result["from_timezone"], serde_json::Value::Null);
|
|
||||||
assert_eq!(result["operation"], serde_json::json!("now"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn keeps_empty_string_for_required_string_only_field() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"name": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["name"]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({ "name": "" });
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
// Required string-only field keeps empty string
|
|
||||||
assert_eq!(result["name"], serde_json::json!(""));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn permissive_schema_is_noop() {
|
fn permissive_schema_is_noop() {
|
||||||
let schema = serde_json::json!({
|
let schema = serde_json::json!({
|
||||||
@@ -693,341 +339,6 @@ mod tests {
|
|||||||
assert_eq!(result["count"], serde_json::json!("10")); // safety: test-only assertion
|
assert_eq!(result["count"], serde_json::json!("10")); // safety: test-only assertion
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_oneof_discriminated_variant() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"oneOf": [
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "list_repos" },
|
|
||||||
"limit": { "type": "integer" },
|
|
||||||
"sort": { "type": "string" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "get_repo" },
|
|
||||||
"repo": { "type": "string" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"action": "list_repos",
|
|
||||||
"limit": "100",
|
|
||||||
"sort": "stars"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["action"], serde_json::json!("list_repos"));
|
|
||||||
assert_eq!(result["limit"], serde_json::json!(100));
|
|
||||||
assert_eq!(result["sort"], serde_json::json!("stars"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_oneof_with_enum_discriminator() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"oneOf": [
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"mode": { "enum": ["fetch"] },
|
|
||||||
"count": { "type": "integer" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"mode": { "enum": ["push"] },
|
|
||||||
"force": { "type": "boolean" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"mode": "push",
|
|
||||||
"force": "true"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["mode"], serde_json::json!("push"));
|
|
||||||
assert_eq!(result["force"], serde_json::json!(true));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_allof_merged_properties() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"allOf": [
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"page": { "type": "integer" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"per_page": { "type": "integer" },
|
|
||||||
"verbose": { "type": "boolean" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"page": "2",
|
|
||||||
"per_page": "50",
|
|
||||||
"verbose": "false"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["page"], serde_json::json!(2));
|
|
||||||
assert_eq!(result["per_page"], serde_json::json!(50));
|
|
||||||
assert_eq!(result["verbose"], serde_json::json!(false));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn oneof_no_discriminator_match_is_noop() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"oneOf": [
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "list_repos" },
|
|
||||||
"limit": { "type": "integer" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "get_repo" },
|
|
||||||
"repo": { "type": "string" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"action": "unknown_action",
|
|
||||||
"limit": "100"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
// No variant matched, so no coercion happens
|
|
||||||
assert_eq!(result["limit"], serde_json::json!("100"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn anyof_without_discriminator_is_noop() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"anyOf": [
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"name": { "type": "string" }
|
|
||||||
},
|
|
||||||
"required": ["name"]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"id": { "type": "integer" }
|
|
||||||
},
|
|
||||||
"required": ["id"]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"id": "42"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
// No const/enum discriminators, so no variant matches, no coercion
|
|
||||||
assert_eq!(result["id"], serde_json::json!("42"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn resolves_ref_and_coerces_referenced_properties() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"definitions": {
|
|
||||||
"Pagination": {
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"page": { "type": "integer" },
|
|
||||||
"per_page": { "type": "integer" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"allOf": [
|
|
||||||
{ "$ref": "#/definitions/Pagination" },
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"query": { "type": "string" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"page": "2",
|
|
||||||
"per_page": "50",
|
|
||||||
"query": "test"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["page"], serde_json::json!(2));
|
|
||||||
assert_eq!(result["per_page"], serde_json::json!(50));
|
|
||||||
assert_eq!(result["query"], serde_json::json!("test"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn resolves_nested_refs_in_oneof_variants() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"$defs": {
|
|
||||||
"ListParams": {
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "list" },
|
|
||||||
"limit": { "type": "integer" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"oneOf": [
|
|
||||||
{ "$ref": "#/$defs/ListParams" },
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "get" },
|
|
||||||
"id": { "type": "integer" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"action": "list",
|
|
||||||
"limit": "25"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["limit"], serde_json::json!(25));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_nested_combinators_allof_containing_oneof() {
|
|
||||||
// allOf where one variant is itself a oneOf (nested combinator)
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"allOf": [
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"version": { "type": "integer" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"oneOf": [
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"mode": { "const": "fast" },
|
|
||||||
"threads": { "type": "integer" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"mode": { "const": "safe" },
|
|
||||||
"retries": { "type": "integer" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"version": "3",
|
|
||||||
"mode": "fast",
|
|
||||||
"threads": "8"
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["version"], serde_json::json!(3));
|
|
||||||
assert_eq!(result["threads"], serde_json::json!(8));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn coerces_array_items_with_oneof_discriminator() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"actions": {
|
|
||||||
"type": "array",
|
|
||||||
"items": {
|
|
||||||
"oneOf": [
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"type": { "const": "move" },
|
|
||||||
"distance": { "type": "integer" }
|
|
||||||
}
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"type": { "const": "wait" },
|
|
||||||
"seconds": { "type": "number" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"actions": [
|
|
||||||
{ "type": "move", "distance": "10" },
|
|
||||||
{ "type": "wait", "seconds": "2.5" }
|
|
||||||
]
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["actions"][0]["distance"], serde_json::json!(10));
|
|
||||||
assert_eq!(result["actions"][1]["seconds"], serde_json::json!(2.5));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn circular_ref_does_not_infinite_loop() {
|
|
||||||
let schema = serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"definitions": {
|
|
||||||
"Node": {
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"value": { "type": "integer" },
|
|
||||||
"child": { "$ref": "#/definitions/Node" }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"properties": {
|
|
||||||
"root": { "$ref": "#/definitions/Node" }
|
|
||||||
}
|
|
||||||
});
|
|
||||||
let params = serde_json::json!({
|
|
||||||
"root": { "value": "42" }
|
|
||||||
});
|
|
||||||
|
|
||||||
// Should not hang — depth limit stops the recursion
|
|
||||||
let result = prepare_params_for_schema(¶ms, &schema);
|
|
||||||
|
|
||||||
assert_eq!(result["root"]["value"], serde_json::json!(42));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn prepare_tool_params_uses_discovery_schema() {
|
fn prepare_tool_params_uses_discovery_schema() {
|
||||||
let tool = StubTool {
|
let tool = StubTool {
|
||||||
|
|||||||
@@ -0,0 +1,490 @@
|
|||||||
|
//! Tool executor for programmatic tool calling (PTC).
|
||||||
|
//!
|
||||||
|
//! Provides a standalone execution engine that can be used by both the
|
||||||
|
//! Docker HTTP RPC path (orchestrator endpoint) and the WASM host function
|
||||||
|
//! path (tool_invoke). Extracts the tool dispatch flow into a reusable struct.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use crate::context::JobContext;
|
||||||
|
use crate::safety::SafetyLayer;
|
||||||
|
use crate::tools::registry::ToolRegistry;
|
||||||
|
use crate::tools::tool::ToolDomain;
|
||||||
|
|
||||||
|
/// Maximum allowed nesting depth for tool-invokes-tool chains.
|
||||||
|
pub const MAX_NESTING_DEPTH: u32 = 5;
|
||||||
|
|
||||||
|
/// Maximum per-call timeout (5 minutes).
|
||||||
|
const MAX_TIMEOUT_SECS: u64 = 300;
|
||||||
|
|
||||||
|
/// Result of a programmatic tool call.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct PtcToolResult {
|
||||||
|
/// Tool output (potentially sanitized).
|
||||||
|
pub output: String,
|
||||||
|
/// Whether the output was modified by the safety layer.
|
||||||
|
pub was_sanitized: bool,
|
||||||
|
/// Wall-clock duration of the tool execution.
|
||||||
|
pub duration: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Errors that can occur during programmatic tool execution.
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum PtcError {
|
||||||
|
#[error("Tool not found: {name}")]
|
||||||
|
NotFound { name: String },
|
||||||
|
|
||||||
|
#[error("Tool execution failed: {name}: {reason}")]
|
||||||
|
ExecutionFailed { name: String, reason: String },
|
||||||
|
|
||||||
|
#[error("Tool execution timed out: {name} (timeout: {timeout:?})")]
|
||||||
|
Timeout { name: String, timeout: Duration },
|
||||||
|
|
||||||
|
#[error("Invalid parameters for tool {name}: {reason}")]
|
||||||
|
InvalidParameters { name: String, reason: String },
|
||||||
|
|
||||||
|
#[error("Tool {name} is rate limited")]
|
||||||
|
RateLimited { name: String },
|
||||||
|
|
||||||
|
#[error("Tool output blocked by safety layer: {reason}")]
|
||||||
|
SafetyBlocked { reason: String },
|
||||||
|
|
||||||
|
#[error("Nesting depth exceeded (max {max})")]
|
||||||
|
NestingDepthExceeded { max: u32 },
|
||||||
|
|
||||||
|
#[error("Tool {name} has domain Container and cannot be executed on the orchestrator")]
|
||||||
|
DomainBlocked { name: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Standalone tool execution engine for programmatic tool calling.
|
||||||
|
///
|
||||||
|
/// Used by:
|
||||||
|
/// - The orchestrator's `POST /worker/{job_id}/tools/call` endpoint
|
||||||
|
/// - The WASM `tool_invoke` host function
|
||||||
|
pub struct ToolExecutor {
|
||||||
|
tools: Arc<ToolRegistry>,
|
||||||
|
safety: Arc<SafetyLayer>,
|
||||||
|
default_timeout: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolExecutor {
|
||||||
|
/// Create a new tool executor.
|
||||||
|
pub fn new(
|
||||||
|
tools: Arc<ToolRegistry>,
|
||||||
|
safety: Arc<SafetyLayer>,
|
||||||
|
default_timeout: Duration,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
tools,
|
||||||
|
safety,
|
||||||
|
default_timeout,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a tool by name with the given parameters.
|
||||||
|
///
|
||||||
|
/// Flow: lookup -> execute with timeout -> sanitize output -> return.
|
||||||
|
pub async fn execute(
|
||||||
|
&self,
|
||||||
|
tool_name: &str,
|
||||||
|
params: serde_json::Value,
|
||||||
|
ctx: &JobContext,
|
||||||
|
timeout_override: Option<Duration>,
|
||||||
|
) -> Result<PtcToolResult, PtcError> {
|
||||||
|
// Enforce global nesting depth limit
|
||||||
|
if ctx.tool_nesting_depth >= MAX_NESTING_DEPTH {
|
||||||
|
return Err(PtcError::NestingDepthExceeded {
|
||||||
|
max: MAX_NESTING_DEPTH,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let start = Instant::now();
|
||||||
|
|
||||||
|
// Look up the tool
|
||||||
|
let tool = self
|
||||||
|
.tools
|
||||||
|
.get(tool_name)
|
||||||
|
.await
|
||||||
|
.ok_or_else(|| PtcError::NotFound {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Reject Container-domain tools — they must run inside a sandbox,
|
||||||
|
// not on the orchestrator host. Without this check a compromised
|
||||||
|
// worker could invoke shell/file tools on the host (sandbox escape).
|
||||||
|
if tool.domain() == ToolDomain::Container {
|
||||||
|
return Err(PtcError::DomainBlocked {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// Determine timeout: caller override -> tool's own timeout -> default,
|
||||||
|
// capped at MAX_TIMEOUT_SECS.
|
||||||
|
let timeout = timeout_override
|
||||||
|
.unwrap_or_else(|| tool.execution_timeout())
|
||||||
|
.min(Duration::from_secs(MAX_TIMEOUT_SECS));
|
||||||
|
|
||||||
|
// Execute with timeout
|
||||||
|
let tool_result = tokio::time::timeout(timeout, tool.execute(params, ctx))
|
||||||
|
.await
|
||||||
|
.map_err(|_| PtcError::Timeout {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
timeout,
|
||||||
|
})?
|
||||||
|
.map_err(|e| match e {
|
||||||
|
crate::tools::ToolError::InvalidParameters(reason) => PtcError::InvalidParameters {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
reason,
|
||||||
|
},
|
||||||
|
crate::tools::ToolError::RateLimited(_) => PtcError::RateLimited {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
},
|
||||||
|
other => PtcError::ExecutionFailed {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
reason: other.to_string(),
|
||||||
|
},
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Get output string
|
||||||
|
let raw_output = tool_result
|
||||||
|
.raw
|
||||||
|
.as_deref()
|
||||||
|
.or_else(|| tool_result.result.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
let raw_output = if raw_output.is_empty() {
|
||||||
|
serde_json::to_string(&tool_result.result).unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
raw_output
|
||||||
|
};
|
||||||
|
|
||||||
|
// Sanitize output if the tool requires it
|
||||||
|
let (output, was_sanitized) = if tool.requires_sanitization() {
|
||||||
|
let sanitized = self.safety.sanitize_tool_output(tool_name, &raw_output);
|
||||||
|
if sanitized.was_modified && sanitized.content.starts_with("[Output blocked") {
|
||||||
|
return Err(PtcError::SafetyBlocked {
|
||||||
|
reason: sanitized.content,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
(sanitized.content, sanitized.was_modified)
|
||||||
|
} else {
|
||||||
|
(raw_output, false)
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(PtcToolResult {
|
||||||
|
output,
|
||||||
|
was_sanitized,
|
||||||
|
duration: start.elapsed(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for ToolExecutor {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("ToolExecutor")
|
||||||
|
.field("default_timeout", &self.default_timeout)
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::safety::SafetyLayer;
|
||||||
|
use crate::tools::tool::{Tool, ToolDomain, ToolError, ToolOutput};
|
||||||
|
|
||||||
|
fn test_safety_config() -> crate::config::SafetyConfig {
|
||||||
|
crate::config::SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct SlowTool;
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl Tool for SlowTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"slow_tool"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"A tool that sleeps"
|
||||||
|
}
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
serde_json::json!({"type": "object"})
|
||||||
|
}
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
_params: serde_json::Value,
|
||||||
|
_ctx: &JobContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
tokio::time::sleep(Duration::from_secs(10)).await;
|
||||||
|
Ok(ToolOutput::text("done", Duration::from_secs(10)))
|
||||||
|
}
|
||||||
|
fn requires_sanitization(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_not_found() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
let result = executor
|
||||||
|
.execute("nonexistent", serde_json::json!({}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(PtcError::NotFound { .. })));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_echo() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
let result = executor
|
||||||
|
.execute("echo", serde_json::json!({"message": "hello"}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
let ptc_result = result.as_ref().ok();
|
||||||
|
assert!(ptc_result.is_some());
|
||||||
|
assert!(
|
||||||
|
ptc_result
|
||||||
|
.map(|r| r.output.contains("hello"))
|
||||||
|
.unwrap_or(false)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_timeout() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register(Arc::new(SlowTool)).await;
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
let result = executor
|
||||||
|
.execute(
|
||||||
|
"slow_tool",
|
||||||
|
serde_json::json!({}),
|
||||||
|
&ctx,
|
||||||
|
Some(Duration::from_millis(50)),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(PtcError::Timeout { .. })));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_nesting_depth_exceeded() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let mut ctx = JobContext::new("test", "test");
|
||||||
|
ctx.tool_nesting_depth = MAX_NESTING_DEPTH; // already at max
|
||||||
|
|
||||||
|
let result = executor
|
||||||
|
.execute("echo", serde_json::json!({"message": "hello"}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(PtcError::NestingDepthExceeded { .. })));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_nesting_depth_within_limit() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let mut ctx = JobContext::new("test", "test");
|
||||||
|
ctx.tool_nesting_depth = MAX_NESTING_DEPTH - 1; // one below max
|
||||||
|
|
||||||
|
let result = executor
|
||||||
|
.execute("echo", serde_json::json!({"message": "hello"}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(result.is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
struct LeakyTool;
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl Tool for LeakyTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"leaky_tool"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Returns output with fake bearer token"
|
||||||
|
}
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
serde_json::json!({"type": "object"})
|
||||||
|
}
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
_params: serde_json::Value,
|
||||||
|
_ctx: &JobContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
// Bearer token pattern triggers LeakAction::Redact (not Block),
|
||||||
|
// so the safety layer redacts it and returns sanitized output.
|
||||||
|
let output =
|
||||||
|
"Here is some data: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9_longtokenvalue end";
|
||||||
|
Ok(ToolOutput::text(output, Duration::from_millis(1)))
|
||||||
|
}
|
||||||
|
fn requires_sanitization(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct InvalidParamsTool;
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl Tool for InvalidParamsTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"invalid_params_tool"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Always fails with InvalidParameters"
|
||||||
|
}
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
serde_json::json!({"type": "object"})
|
||||||
|
}
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
_params: serde_json::Value,
|
||||||
|
_ctx: &JobContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
Err(ToolError::InvalidParameters("bad params".to_string()))
|
||||||
|
}
|
||||||
|
fn requires_sanitization(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_safety_sanitization() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register(Arc::new(LeakyTool)).await;
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
let result = executor
|
||||||
|
.execute("leaky_tool", serde_json::json!({}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// The safety layer should detect the API key pattern and modify the output
|
||||||
|
assert!(result.is_ok());
|
||||||
|
let ptc_result = result.unwrap();
|
||||||
|
assert!(
|
||||||
|
ptc_result.was_sanitized,
|
||||||
|
"Output with API key should be sanitized"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_invalid_params() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register(Arc::new(InvalidParamsTool)).await;
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
let result = executor
|
||||||
|
.execute("invalid_params_tool", serde_json::json!({}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Err(PtcError::InvalidParameters { name, reason }) => {
|
||||||
|
assert_eq!(name, "invalid_params_tool");
|
||||||
|
assert!(reason.contains("bad params"));
|
||||||
|
}
|
||||||
|
other => panic!("Expected InvalidParameters, got {:?}", other),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A tool that declares Container domain — must be blocked by the executor.
|
||||||
|
struct ContainerDomainTool;
|
||||||
|
|
||||||
|
#[async_trait::async_trait]
|
||||||
|
impl Tool for ContainerDomainTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"container_tool"
|
||||||
|
}
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Simulates a container-domain tool"
|
||||||
|
}
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
serde_json::json!({"type": "object"})
|
||||||
|
}
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
_params: serde_json::Value,
|
||||||
|
_ctx: &JobContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
Ok(ToolOutput::text(
|
||||||
|
"should not reach here",
|
||||||
|
Duration::from_millis(1),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
fn domain(&self) -> ToolDomain {
|
||||||
|
ToolDomain::Container
|
||||||
|
}
|
||||||
|
fn requires_sanitization(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_container_domain_blocked() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register(Arc::new(ContainerDomainTool)).await;
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
let result = executor
|
||||||
|
.execute("container_tool", serde_json::json!({}), &ctx, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
matches!(result, Err(PtcError::DomainBlocked { .. })),
|
||||||
|
"Container-domain tools must be rejected: {:?}",
|
||||||
|
result
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_sequential_calls() {
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
|
||||||
|
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
|
||||||
|
|
||||||
|
let ctx = JobContext::new("test", "test");
|
||||||
|
|
||||||
|
let messages = ["alpha", "beta", "gamma"];
|
||||||
|
for msg in &messages {
|
||||||
|
let result = executor
|
||||||
|
.execute("echo", serde_json::json!({"message": msg}), &ctx, None)
|
||||||
|
.await
|
||||||
|
.expect("echo should succeed");
|
||||||
|
assert!(
|
||||||
|
result.output.contains(msg),
|
||||||
|
"Output should contain '{}'",
|
||||||
|
msg
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -18,6 +18,7 @@ pub mod redaction;
|
|||||||
pub mod schema_validator;
|
pub mod schema_validator;
|
||||||
pub mod wasm;
|
pub mod wasm;
|
||||||
|
|
||||||
|
mod executor;
|
||||||
mod registry;
|
mod registry;
|
||||||
mod tool;
|
mod tool;
|
||||||
|
|
||||||
@@ -31,6 +32,7 @@ pub use builder::{
|
|||||||
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
||||||
};
|
};
|
||||||
pub(crate) use coercion::prepare_tool_params;
|
pub(crate) use coercion::prepare_tool_params;
|
||||||
|
pub use executor::{PtcError, PtcToolResult, ToolExecutor};
|
||||||
pub use rate_limiter::RateLimiter;
|
pub use rate_limiter::RateLimiter;
|
||||||
pub use registry::ToolRegistry;
|
pub use registry::ToolRegistry;
|
||||||
pub use tool::{
|
pub use tool::{
|
||||||
|
|||||||
+73
-5
@@ -19,11 +19,12 @@ use crate::tools::builder::{
|
|||||||
use crate::tools::builtin::{
|
use crate::tools::builtin::{
|
||||||
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
||||||
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
||||||
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool,
|
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, PtcScriptTool,
|
||||||
ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool,
|
ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool,
|
||||||
ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool,
|
TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool,
|
||||||
ToolUpgradeTool, WriteFileTool,
|
ToolSearchTool, ToolUpgradeTool, WriteFileTool,
|
||||||
};
|
};
|
||||||
|
use crate::tools::executor::ToolExecutor;
|
||||||
use crate::tools::rate_limiter::RateLimiter;
|
use crate::tools::rate_limiter::RateLimiter;
|
||||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain};
|
use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain};
|
||||||
use crate::tools::wasm::{
|
use crate::tools::wasm::{
|
||||||
@@ -78,6 +79,7 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[
|
|||||||
"image_edit",
|
"image_edit",
|
||||||
"image_analyze",
|
"image_analyze",
|
||||||
"tool_info",
|
"tool_info",
|
||||||
|
"ptc_script",
|
||||||
];
|
];
|
||||||
|
|
||||||
/// Registry of available tools.
|
/// Registry of available tools.
|
||||||
@@ -93,6 +95,14 @@ pub struct ToolRegistry {
|
|||||||
rate_limiter: RateLimiter,
|
rate_limiter: RateLimiter,
|
||||||
/// Reference to the message tool for setting context per-turn.
|
/// Reference to the message tool for setting context per-turn.
|
||||||
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
|
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
|
||||||
|
/// Shared slot for the tool executor (enables PTC via tool_invoke).
|
||||||
|
///
|
||||||
|
/// Uses `std::sync::RwLock` (not tokio) because reads happen inside
|
||||||
|
/// `spawn_blocking` closures in WASM tool execution. The slot is
|
||||||
|
/// populated lazily after `AppBuilder::build_all()` completes, so
|
||||||
|
/// WASM tools registered during startup still get access to the
|
||||||
|
/// executor when they execute later.
|
||||||
|
tool_executor_slot: Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ToolRegistry {
|
impl ToolRegistry {
|
||||||
@@ -114,6 +124,7 @@ impl ToolRegistry {
|
|||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
rate_limiter: RateLimiter::new(),
|
rate_limiter: RateLimiter::new(),
|
||||||
message_tool: RwLock::new(None),
|
message_tool: RwLock::new(None),
|
||||||
|
tool_executor_slot: Arc::new(std::sync::RwLock::new(None)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -138,6 +149,27 @@ impl ToolRegistry {
|
|||||||
&self.rate_limiter
|
&self.rate_limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Set the tool executor for programmatic tool calling (PTC).
|
||||||
|
///
|
||||||
|
/// Writes the executor into the shared slot so all WASM tools --
|
||||||
|
/// including those registered before this call -- can resolve it
|
||||||
|
/// lazily at execution time.
|
||||||
|
pub fn set_tool_executor(&self, executor: Arc<ToolExecutor>) {
|
||||||
|
if let Ok(mut guard) = self.tool_executor_slot.write() {
|
||||||
|
*guard = Some(executor);
|
||||||
|
} else {
|
||||||
|
tracing::error!("tool_executor_slot RwLock is poisoned; PTC will be unavailable");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a clone of the shared tool executor slot.
|
||||||
|
///
|
||||||
|
/// WASM wrappers hold this slot and read from it at execution time,
|
||||||
|
/// allowing the executor to be set after tool registration.
|
||||||
|
pub fn tool_executor_slot(&self) -> Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>> {
|
||||||
|
Arc::clone(&self.tool_executor_slot)
|
||||||
|
}
|
||||||
|
|
||||||
/// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name.
|
/// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name.
|
||||||
pub async fn register(&self, tool: Arc<dyn Tool>) {
|
pub async fn register(&self, tool: Arc<dyn Tool>) {
|
||||||
let name = tool.name().to_string();
|
let name = tool.name().to_string();
|
||||||
@@ -330,8 +362,9 @@ impl ToolRegistry {
|
|||||||
self.register_sync(Arc::new(WriteFileTool::new()));
|
self.register_sync(Arc::new(WriteFileTool::new()));
|
||||||
self.register_sync(Arc::new(ListDirTool::new()));
|
self.register_sync(Arc::new(ListDirTool::new()));
|
||||||
self.register_sync(Arc::new(ApplyPatchTool::new()));
|
self.register_sync(Arc::new(ApplyPatchTool::new()));
|
||||||
|
self.register_sync(Arc::new(PtcScriptTool::new()));
|
||||||
|
|
||||||
tracing::debug!("Registered 5 development tools");
|
tracing::debug!("Registered 6 development tools");
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Register memory tools with a workspace.
|
/// Register memory tools with a workspace.
|
||||||
@@ -659,6 +692,11 @@ impl ToolRegistry {
|
|||||||
wrapper = wrapper.with_oauth_refresh(oauth);
|
wrapper = wrapper.with_oauth_refresh(oauth);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Inject shared tool executor slot for PTC (lazy resolution).
|
||||||
|
// The WASM wrapper reads from this slot at execution time, so the
|
||||||
|
// executor can be set after tool registration.
|
||||||
|
wrapper = wrapper.with_tool_executor_slot(Arc::clone(&self.tool_executor_slot));
|
||||||
|
|
||||||
// Register the tool
|
// Register the tool
|
||||||
self.register(Arc::new(wrapper)).await;
|
self.register(Arc::new(wrapper)).await;
|
||||||
|
|
||||||
@@ -889,6 +927,36 @@ mod tests {
|
|||||||
assert!(def.parameters.get("extra").is_none());
|
assert!(def.parameters.get("extra").is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_tool_executor_slot_lazy_resolution() {
|
||||||
|
let registry = ToolRegistry::new();
|
||||||
|
|
||||||
|
// Get the slot BEFORE setting the executor (simulates startup order)
|
||||||
|
let slot = registry.tool_executor_slot();
|
||||||
|
|
||||||
|
// Slot should be empty
|
||||||
|
assert!(slot.read().unwrap().is_none());
|
||||||
|
|
||||||
|
// Set the executor (simulates main.rs wiring after build_all)
|
||||||
|
let tools = Arc::new(ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(crate::safety::SafetyLayer::new(
|
||||||
|
&crate::config::SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
},
|
||||||
|
));
|
||||||
|
let executor = Arc::new(crate::tools::ToolExecutor::new(
|
||||||
|
tools,
|
||||||
|
safety,
|
||||||
|
std::time::Duration::from_secs(60),
|
||||||
|
));
|
||||||
|
registry.set_tool_executor(Arc::clone(&executor));
|
||||||
|
|
||||||
|
// Slot should now contain the executor
|
||||||
|
assert!(slot.read().unwrap().is_some());
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_builtin_tool_cannot_be_shadowed() {
|
async fn test_builtin_tool_cannot_be_shadowed() {
|
||||||
let registry = ToolRegistry::new();
|
let registry = ToolRegistry::new();
|
||||||
|
|||||||
@@ -42,38 +42,11 @@ pub fn validate_strict_schema(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns true if the schema uses `oneOf`, `anyOf`, or `allOf` combinators
|
|
||||||
/// where at least one variant is an object type (has `type: "object"` or `properties`).
|
|
||||||
fn has_object_combinator_variants(schema: &serde_json::Value) -> bool {
|
|
||||||
for key in ["oneOf", "anyOf", "allOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array())
|
|
||||||
&& variants.iter().any(|v| {
|
|
||||||
v.get("type").and_then(|t| t.as_str()) == Some("object")
|
|
||||||
|| v.get("properties").is_some()
|
|
||||||
})
|
|
||||||
{
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
false
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Recursively validate an object-typed schema node.
|
/// Recursively validate an object-typed schema node.
|
||||||
fn check_object_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
|
fn check_object_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
|
||||||
let mut errors = Vec::new();
|
let mut errors = Vec::new();
|
||||||
|
|
||||||
// Report non-array combinator values as errors.
|
// Rule 1: must have "type": "object"
|
||||||
for key in ["oneOf", "anyOf", "allOf"] {
|
|
||||||
if let Some(val) = schema.get(key)
|
|
||||||
&& !val.is_array()
|
|
||||||
{
|
|
||||||
errors.push(format!("{path}: \"{key}\" must be an array"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let has_combinators = has_object_combinator_variants(schema);
|
|
||||||
|
|
||||||
// Rule 1: must have "type": "object" (unless combinators define the structure)
|
|
||||||
match schema.get("type").and_then(|t| t.as_str()) {
|
match schema.get("type").and_then(|t| t.as_str()) {
|
||||||
Some("object") => {}
|
Some("object") => {}
|
||||||
Some(other) => {
|
Some(other) => {
|
||||||
@@ -81,67 +54,16 @@ fn check_object_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
|
|||||||
return errors;
|
return errors;
|
||||||
}
|
}
|
||||||
None => {
|
None => {
|
||||||
if !has_combinators {
|
errors.push(format!("{path}: missing \"type\": \"object\""));
|
||||||
errors.push(format!("{path}: missing \"type\": \"object\""));
|
return errors;
|
||||||
return errors;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Validate combinator variants recursively
|
// Rule 2: must have "properties" as an object
|
||||||
for key in ["allOf", "oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array()) {
|
|
||||||
for (i, variant) in variants.iter().enumerate() {
|
|
||||||
if variant.get("type").and_then(|t| t.as_str()) == Some("object")
|
|
||||||
|| variant.get("properties").is_some()
|
|
||||||
{
|
|
||||||
let variant_path = format!("{path}.{key}[{i}]");
|
|
||||||
errors.extend(check_object_schema(variant, &variant_path));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Rule 2: must have "properties" as an object (unless combinators define them)
|
|
||||||
let properties = match schema.get("properties").and_then(|p| p.as_object()) {
|
let properties = match schema.get("properties").and_then(|p| p.as_object()) {
|
||||||
Some(p) => p,
|
Some(p) => p,
|
||||||
None => {
|
None => {
|
||||||
if !has_combinators {
|
errors.push(format!("{path}: missing or non-object \"properties\""));
|
||||||
errors.push(format!("{path}: missing or non-object \"properties\""));
|
|
||||||
return errors;
|
|
||||||
}
|
|
||||||
// Combinators define the structure — validate top-level `required` keys
|
|
||||||
// against merged properties from all combinator variants.
|
|
||||||
if let Some(required) = schema.get("required").and_then(|r| r.as_array()) {
|
|
||||||
let mut merged_keys = std::collections::HashSet::new();
|
|
||||||
if let Some(all_of) = schema.get("allOf").and_then(|a| a.as_array()) {
|
|
||||||
for variant in all_of {
|
|
||||||
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
|
||||||
merged_keys.extend(props.keys().cloned());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for key in ["oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array()) {
|
|
||||||
for variant in variants {
|
|
||||||
if let Some(props) =
|
|
||||||
variant.get("properties").and_then(|p| p.as_object())
|
|
||||||
{
|
|
||||||
merged_keys.extend(props.keys().cloned());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for req in required {
|
|
||||||
if let Some(key) = req.as_str()
|
|
||||||
&& !merged_keys.contains(key)
|
|
||||||
{
|
|
||||||
errors.push(format!(
|
|
||||||
"{path}: required key \"{key}\" not found in any combinator variant properties"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return errors;
|
return errors;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
+5
-87
@@ -462,22 +462,6 @@ pub fn redact_params(params: &serde_json::Value, sensitive: &[&str]) -> serde_js
|
|||||||
/// on maliciously crafted schemas.
|
/// on maliciously crafted schemas.
|
||||||
const MAX_SCHEMA_DEPTH: usize = 16;
|
const MAX_SCHEMA_DEPTH: usize = 16;
|
||||||
|
|
||||||
/// Returns true if the schema uses `oneOf`, `anyOf`, or `allOf` combinators
|
|
||||||
/// where at least one variant is an object type (has `type: "object"` or `properties`).
|
|
||||||
fn has_object_combinator_variants(schema: &serde_json::Value) -> bool {
|
|
||||||
for key in ["oneOf", "anyOf", "allOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array())
|
|
||||||
&& variants.iter().any(|v| {
|
|
||||||
v.get("type").and_then(|t| t.as_str()) == Some("object")
|
|
||||||
|| v.get("properties").is_some()
|
|
||||||
})
|
|
||||||
{
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
false
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn validate_tool_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
|
pub fn validate_tool_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
|
||||||
validate_tool_schema_inner(schema, path, 0)
|
validate_tool_schema_inner(schema, path, 0)
|
||||||
}
|
}
|
||||||
@@ -492,18 +476,7 @@ fn validate_tool_schema_inner(schema: &serde_json::Value, path: &str, depth: usi
|
|||||||
return errors;
|
return errors;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Report non-array combinator values as errors.
|
// Rule 1: must have "type": "object" at this level
|
||||||
for key in ["oneOf", "anyOf", "allOf"] {
|
|
||||||
if let Some(val) = schema.get(key)
|
|
||||||
&& !val.is_array()
|
|
||||||
{
|
|
||||||
errors.push(format!("{path}: \"{key}\" must be an array"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let has_combinators = has_object_combinator_variants(schema);
|
|
||||||
|
|
||||||
// Rule 1: must have "type": "object" at this level (unless combinators define the structure)
|
|
||||||
match schema.get("type").and_then(|t| t.as_str()) {
|
match schema.get("type").and_then(|t| t.as_str()) {
|
||||||
Some("object") => {}
|
Some("object") => {}
|
||||||
Some(other) => {
|
Some(other) => {
|
||||||
@@ -511,71 +484,16 @@ fn validate_tool_schema_inner(schema: &serde_json::Value, path: &str, depth: usi
|
|||||||
return errors; // Can't check further
|
return errors; // Can't check further
|
||||||
}
|
}
|
||||||
None => {
|
None => {
|
||||||
if !has_combinators {
|
errors.push(format!("{path}: missing \"type\": \"object\""));
|
||||||
errors.push(format!("{path}: missing \"type\": \"object\""));
|
return errors;
|
||||||
return errors;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Validate combinator variants recursively
|
// Rule 2: must have "properties" as an object
|
||||||
for key in ["allOf", "oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array()) {
|
|
||||||
for (i, variant) in variants.iter().enumerate() {
|
|
||||||
if variant.get("type").and_then(|t| t.as_str()) == Some("object")
|
|
||||||
|| variant.get("properties").is_some()
|
|
||||||
{
|
|
||||||
let variant_path = format!("{path}.{key}[{i}]");
|
|
||||||
errors.extend(validate_tool_schema_inner(
|
|
||||||
variant,
|
|
||||||
&variant_path,
|
|
||||||
depth + 1,
|
|
||||||
));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Rule 2: must have "properties" as an object (unless combinators define them)
|
|
||||||
let properties = match schema.get("properties").and_then(|p| p.as_object()) {
|
let properties = match schema.get("properties").and_then(|p| p.as_object()) {
|
||||||
Some(p) => p,
|
Some(p) => p,
|
||||||
None => {
|
None => {
|
||||||
if !has_combinators {
|
errors.push(format!("{path}: missing or non-object \"properties\""));
|
||||||
errors.push(format!("{path}: missing or non-object \"properties\""));
|
|
||||||
return errors;
|
|
||||||
}
|
|
||||||
// Combinators define the structure — validate top-level `required` keys
|
|
||||||
// against merged properties from all combinator variants.
|
|
||||||
if let Some(required) = schema.get("required").and_then(|r| r.as_array()) {
|
|
||||||
let mut merged_keys = std::collections::HashSet::new();
|
|
||||||
if let Some(all_of) = schema.get("allOf").and_then(|a| a.as_array()) {
|
|
||||||
for variant in all_of {
|
|
||||||
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
|
||||||
merged_keys.extend(props.keys().cloned());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for key in ["oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array()) {
|
|
||||||
for variant in variants {
|
|
||||||
if let Some(props) =
|
|
||||||
variant.get("properties").and_then(|p| p.as_object())
|
|
||||||
{
|
|
||||||
merged_keys.extend(props.keys().cloned());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for req in required {
|
|
||||||
if let Some(key) = req.as_str()
|
|
||||||
&& !merged_keys.contains(key)
|
|
||||||
{
|
|
||||||
errors.push(format!(
|
|
||||||
"{path}: required key \"{key}\" not found in any combinator variant properties"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return errors;
|
return errors;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -708,9 +708,6 @@ pub struct ToolSetupSchema {
|
|||||||
/// Secrets the user must provide before the tool can be used.
|
/// Secrets the user must provide before the tool can be used.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub required_secrets: Vec<ToolSecretSetupSchema>,
|
pub required_secrets: Vec<ToolSecretSetupSchema>,
|
||||||
/// Non-secret fields the user can configure in the setup modal.
|
|
||||||
#[serde(default)]
|
|
||||||
pub required_fields: Vec<ToolFieldSetupSchema>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A single secret required during tool setup.
|
/// A single secret required during tool setup.
|
||||||
@@ -725,46 +722,6 @@ pub struct ToolSecretSetupSchema {
|
|||||||
pub optional: bool,
|
pub optional: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A non-secret field required during tool setup.
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
||||||
pub struct ToolFieldSetupSchema {
|
|
||||||
/// Field name in setup payload.
|
|
||||||
pub name: String,
|
|
||||||
/// User-facing prompt shown in the setup modal.
|
|
||||||
pub prompt: String,
|
|
||||||
/// If true, the user may skip this field.
|
|
||||||
#[serde(default)]
|
|
||||||
pub optional: bool,
|
|
||||||
/// Input type used in the setup modal.
|
|
||||||
#[serde(default = "default_tool_setup_field_input_type")]
|
|
||||||
pub input_type: ToolSetupFieldInputType,
|
|
||||||
/// Optional dotted setting path to persist this value to.
|
|
||||||
///
|
|
||||||
/// Restricted by the host to extension-owned namespaces and a small
|
|
||||||
/// allowlist of approved global settings.
|
|
||||||
///
|
|
||||||
/// Example: `extensions.switch-llm.provider`, `llm_backend`, or
|
|
||||||
/// `selected_model`.
|
|
||||||
#[serde(default)]
|
|
||||||
pub setting_path: Option<String>,
|
|
||||||
/// Whether changing this field requires a restart to fully apply.
|
|
||||||
#[serde(default)]
|
|
||||||
pub restart_required: bool,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Input widget type for a setup field.
|
|
||||||
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
|
|
||||||
#[serde(rename_all = "snake_case")]
|
|
||||||
pub enum ToolSetupFieldInputType {
|
|
||||||
#[default]
|
|
||||||
Text,
|
|
||||||
Password,
|
|
||||||
}
|
|
||||||
|
|
||||||
fn default_tool_setup_field_input_type() -> ToolSetupFieldInputType {
|
|
||||||
ToolSetupFieldInputType::Text
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use crate::tools::wasm::capabilities_schema::{CapabilitiesFile, CredentialLocationSchema};
|
use crate::tools::wasm::capabilities_schema::{CapabilitiesFile, CredentialLocationSchema};
|
||||||
@@ -1261,20 +1218,6 @@ mod tests {
|
|||||||
"prompt": "Google OAuth Client Secret",
|
"prompt": "Google OAuth Client Secret",
|
||||||
"optional": true
|
"optional": true
|
||||||
}
|
}
|
||||||
],
|
|
||||||
"required_fields": [
|
|
||||||
{
|
|
||||||
"name": "llm_backend",
|
|
||||||
"prompt": "LLM Provider",
|
|
||||||
"setting_path": "llm_backend",
|
|
||||||
"restart_required": true
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"name": "selected_model",
|
|
||||||
"prompt": "Model Name",
|
|
||||||
"input_type": "text",
|
|
||||||
"setting_path": "selected_model"
|
|
||||||
}
|
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
}"#;
|
}"#;
|
||||||
@@ -1287,48 +1230,6 @@ mod tests {
|
|||||||
assert!(!setup.required_secrets[0].optional);
|
assert!(!setup.required_secrets[0].optional);
|
||||||
assert_eq!(setup.required_secrets[1].name, "google_oauth_client_secret");
|
assert_eq!(setup.required_secrets[1].name, "google_oauth_client_secret");
|
||||||
assert!(setup.required_secrets[1].optional);
|
assert!(setup.required_secrets[1].optional);
|
||||||
assert_eq!(setup.required_fields.len(), 2);
|
|
||||||
assert_eq!(setup.required_fields[0].name, "llm_backend");
|
|
||||||
assert_eq!(
|
|
||||||
setup.required_fields[0].setting_path.as_deref(),
|
|
||||||
Some("llm_backend")
|
|
||||||
);
|
|
||||||
assert!(setup.required_fields[0].restart_required);
|
|
||||||
assert_eq!(
|
|
||||||
setup.required_fields[0].input_type,
|
|
||||||
crate::tools::wasm::capabilities_schema::ToolSetupFieldInputType::Text
|
|
||||||
);
|
|
||||||
assert_eq!(setup.required_fields[1].name, "selected_model");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_tool_setup_field_input_type_defaults_to_text() {
|
|
||||||
let json = r#"{
|
|
||||||
"setup": {
|
|
||||||
"required_fields": [
|
|
||||||
{
|
|
||||||
"name": "provider",
|
|
||||||
"prompt": "Provider"
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"name": "token_hint",
|
|
||||||
"prompt": "Token Hint",
|
|
||||||
"input_type": "password"
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
}"#;
|
|
||||||
|
|
||||||
let caps = CapabilitiesFile::from_json(json).unwrap();
|
|
||||||
let setup = caps.setup.unwrap();
|
|
||||||
assert_eq!(
|
|
||||||
setup.required_fields[0].input_type,
|
|
||||||
crate::tools::wasm::capabilities_schema::ToolSetupFieldInputType::Text
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
setup.required_fields[1].input_type,
|
|
||||||
crate::tools::wasm::capabilities_schema::ToolSetupFieldInputType::Password
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -139,5 +139,5 @@ pub use loader::{
|
|||||||
// Capabilities schema (for parsing *.capabilities.json files)
|
// Capabilities schema (for parsing *.capabilities.json files)
|
||||||
pub use capabilities_schema::{
|
pub use capabilities_schema::{
|
||||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, RateLimitSchema,
|
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, RateLimitSchema,
|
||||||
ToolFieldSetupSchema, ToolSetupFieldInputType, ToolSetupSchema, ValidationEndpointSchema,
|
ValidationEndpointSchema,
|
||||||
};
|
};
|
||||||
|
|||||||
+402
-193
@@ -17,9 +17,9 @@ use wasmtime::component::Linker;
|
|||||||
use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
||||||
|
|
||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor};
|
|
||||||
use crate::safety::LeakDetector;
|
use crate::safety::LeakDetector;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
|
use crate::tools::ToolExecutor;
|
||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||||
use crate::tools::wasm::capabilities::Capabilities;
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::credential_injector::{
|
use crate::tools::wasm::credential_injector::{
|
||||||
@@ -30,6 +30,26 @@ use crate::tools::wasm::host::{HostState, LogLevel};
|
|||||||
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
|
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
|
||||||
use crate::tools::wasm::runtime::{EPOCH_TICK_INTERVAL, PreparedModule, WasmToolRuntime};
|
use crate::tools::wasm::runtime::{EPOCH_TICK_INTERVAL, PreparedModule, WasmToolRuntime};
|
||||||
|
|
||||||
|
/// Synchronous tool resolver callable from within a WASM host function.
|
||||||
|
/// The closure internally creates a tokio runtime to bridge async tool execution.
|
||||||
|
/// Closure that resolves a tool call by name. The `u32` parameter is the current
|
||||||
|
/// nesting depth so the executor can enforce the global depth limit across
|
||||||
|
/// WASM->executor->WASM chains.
|
||||||
|
pub type ToolResolver =
|
||||||
|
Arc<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + Send + Sync>;
|
||||||
|
|
||||||
|
/// RAII guard that decrements the nesting depth counter on drop, ensuring the
|
||||||
|
/// counter is restored even if the code between increment and decrement panics.
|
||||||
|
struct NestingGuard<'a> {
|
||||||
|
depth: &'a mut u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for NestingGuard<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
*self.depth = self.depth.saturating_sub(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Generate component model bindings from the WIT file.
|
// Generate component model bindings from the WIT file.
|
||||||
//
|
//
|
||||||
// This creates:
|
// This creates:
|
||||||
@@ -100,9 +120,11 @@ struct StoreData {
|
|||||||
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
|
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
|
||||||
/// Reused across multiple `http_request` calls within one execution.
|
/// Reused across multiple `http_request` calls within one execution.
|
||||||
http_runtime: Option<tokio::runtime::Runtime>,
|
http_runtime: Option<tokio::runtime::Runtime>,
|
||||||
/// Optional HTTP interceptor for testing — returns canned responses
|
/// Optional tool resolver for programmatic tool calling (PTC).
|
||||||
/// instead of making real requests when set.
|
/// When set, WASM tools can invoke other tools via the `tool_invoke` host function.
|
||||||
http_interceptor: Option<Arc<dyn HttpInterceptor>>,
|
tool_resolver: Option<ToolResolver>,
|
||||||
|
/// Current nesting depth for tool_invoke calls. Prevents infinite recursion.
|
||||||
|
tool_nesting_depth: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl StoreData {
|
impl StoreData {
|
||||||
@@ -111,6 +133,7 @@ impl StoreData {
|
|||||||
capabilities: Capabilities,
|
capabilities: Capabilities,
|
||||||
credentials: HashMap<String, String>,
|
credentials: HashMap<String, String>,
|
||||||
host_credentials: Vec<ResolvedHostCredential>,
|
host_credentials: Vec<ResolvedHostCredential>,
|
||||||
|
tool_resolver: Option<ToolResolver>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
// Minimal WASI context: no filesystem, no env vars (security)
|
// Minimal WASI context: no filesystem, no env vars (security)
|
||||||
let wasi = WasiCtxBuilder::new().build();
|
let wasi = WasiCtxBuilder::new().build();
|
||||||
@@ -123,7 +146,8 @@ impl StoreData {
|
|||||||
credentials,
|
credentials,
|
||||||
host_credentials,
|
host_credentials,
|
||||||
http_runtime: None,
|
http_runtime: None,
|
||||||
http_interceptor: None,
|
tool_resolver,
|
||||||
|
tool_nesting_depth: 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -349,59 +373,6 @@ impl near::agent::host::Host for StoreData {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
let rt = self.http_runtime.as_ref().expect("just initialized"); // safety: is_none branch above guarantees Some
|
let rt = self.http_runtime.as_ref().expect("just initialized"); // safety: is_none branch above guarantees Some
|
||||||
|
|
||||||
// If an HTTP interceptor is set (testing), short-circuit with a canned response.
|
|
||||||
if let Some(interceptor) = &self.http_interceptor {
|
|
||||||
let interceptor = Arc::clone(interceptor);
|
|
||||||
let intercept_url = url.clone();
|
|
||||||
let intercept_method = method.clone();
|
|
||||||
let mut intercept_headers: Vec<(String, String)> = headers
|
|
||||||
.iter()
|
|
||||||
.map(|(k, v)| (k.clone(), v.clone()))
|
|
||||||
.collect();
|
|
||||||
intercept_headers.sort_by(|a, b| a.0.cmp(&b.0));
|
|
||||||
let intercept_body = body
|
|
||||||
.as_ref()
|
|
||||||
.map(|b| String::from_utf8_lossy(b).to_string());
|
|
||||||
let intercepted = rt.block_on(async {
|
|
||||||
let req = HttpExchangeRequest {
|
|
||||||
method: intercept_method,
|
|
||||||
url: intercept_url,
|
|
||||||
headers: intercept_headers,
|
|
||||||
body: intercept_body,
|
|
||||||
};
|
|
||||||
interceptor.before_request(&req).await
|
|
||||||
});
|
|
||||||
if let Some(resp) = intercepted {
|
|
||||||
let resp_headers: HashMap<String, String> = resp
|
|
||||||
.headers
|
|
||||||
.iter()
|
|
||||||
.map(|(k, v)| (k.clone(), v.clone()))
|
|
||||||
.collect();
|
|
||||||
let resp_headers_json =
|
|
||||||
serde_json::to_string(&resp_headers).unwrap_or_else(|_| "{}".to_string());
|
|
||||||
return Ok(near::agent::host::HttpResponse {
|
|
||||||
status: resp.status,
|
|
||||||
headers_json: resp_headers_json,
|
|
||||||
body: resp.body.into_bytes(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Capture request metadata before headers/body are consumed by the reqwest
|
|
||||||
// builder. Used for after_response callback when a recording interceptor is set.
|
|
||||||
let interceptor_req = self.http_interceptor.as_ref().map(|_| HttpExchangeRequest {
|
|
||||||
method: method.clone(),
|
|
||||||
url: url.clone(),
|
|
||||||
headers: headers
|
|
||||||
.iter()
|
|
||||||
.map(|(k, v)| (k.clone(), v.clone()))
|
|
||||||
.collect(),
|
|
||||||
body: body
|
|
||||||
.as_ref()
|
|
||||||
.map(|b| String::from_utf8_lossy(b).to_string()),
|
|
||||||
});
|
|
||||||
|
|
||||||
let result = rt.block_on(async {
|
let result = rt.block_on(async {
|
||||||
let client = reqwest::Client::builder()
|
let client = reqwest::Client::builder()
|
||||||
.connect_timeout(Duration::from_secs(10))
|
.connect_timeout(Duration::from_secs(10))
|
||||||
@@ -492,63 +463,43 @@ impl near::agent::host::Host for StoreData {
|
|||||||
})
|
})
|
||||||
});
|
});
|
||||||
|
|
||||||
// Notify the interceptor about the completed response (recording mode).
|
|
||||||
// RecordingHttpInterceptor returns None from before_request and captures
|
|
||||||
// exchanges via after_response, so this path is exercised during trace recording.
|
|
||||||
if let (Some(interceptor), Some(req), Ok(resp)) =
|
|
||||||
(&self.http_interceptor, &interceptor_req, &result)
|
|
||||||
{
|
|
||||||
let interceptor = Arc::clone(interceptor);
|
|
||||||
|
|
||||||
// Redact credentials from request before passing to the interceptor
|
|
||||||
// to prevent credential leakage into recorded traces.
|
|
||||||
let mut redacted_req = req.clone();
|
|
||||||
redacted_req.url = self.redact_credentials(&redacted_req.url);
|
|
||||||
redacted_req.headers = redacted_req
|
|
||||||
.headers
|
|
||||||
.into_iter()
|
|
||||||
.map(|(k, v)| (k, self.redact_credentials(&v)))
|
|
||||||
.collect();
|
|
||||||
redacted_req.body = redacted_req.body.map(|b| self.redact_credentials(&b));
|
|
||||||
|
|
||||||
let resp_headers: Vec<(String, String)> =
|
|
||||||
serde_json::from_str::<HashMap<String, String>>(&resp.headers_json)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.into_iter()
|
|
||||||
.collect();
|
|
||||||
let resp_body = String::from_utf8_lossy(&resp.body).to_string();
|
|
||||||
|
|
||||||
// Redact credentials from response as well
|
|
||||||
let redacted_headers: Vec<(String, String)> = resp_headers
|
|
||||||
.into_iter()
|
|
||||||
.map(|(k, v)| (k, self.redact_credentials(&v)))
|
|
||||||
.collect();
|
|
||||||
let redacted_body = self.redact_credentials(&resp_body);
|
|
||||||
|
|
||||||
let exchange_resp = HttpExchangeResponse {
|
|
||||||
status: resp.status,
|
|
||||||
headers: redacted_headers,
|
|
||||||
body: redacted_body,
|
|
||||||
};
|
|
||||||
rt.block_on(async {
|
|
||||||
interceptor
|
|
||||||
.after_response(&redacted_req, &exchange_resp)
|
|
||||||
.await;
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
// Redact credentials from error messages before returning to WASM
|
// Redact credentials from error messages before returning to WASM
|
||||||
result.map_err(|e| self.redact_credentials(&e))
|
result.map_err(|e| self.redact_credentials(&e))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn tool_invoke(&mut self, alias: String, _params_json: String) -> Result<String, String> {
|
fn tool_invoke(&mut self, alias: String, params_json: String) -> Result<String, String> {
|
||||||
|
use crate::tools::executor::MAX_NESTING_DEPTH;
|
||||||
|
|
||||||
// Validate capability and resolve alias
|
// Validate capability and resolve alias
|
||||||
let _real_name = self.host_state.check_tool_invoke_allowed(&alias)?;
|
let real_name = self.host_state.check_tool_invoke_allowed(&alias)?;
|
||||||
self.host_state.record_tool_invoke()?;
|
self.host_state.record_tool_invoke()?;
|
||||||
|
|
||||||
// Tool invocation requires async context and access to the tool registry,
|
// Check nesting depth
|
||||||
// which aren't available inside a synchronous WASM callback.
|
if self.tool_nesting_depth >= MAX_NESTING_DEPTH {
|
||||||
Err("Tool invocation from WASM tools is not yet supported".to_string())
|
return Err(format!(
|
||||||
|
"Tool invoke nesting depth exceeded (max {})",
|
||||||
|
MAX_NESTING_DEPTH
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get the resolver
|
||||||
|
let resolver = self
|
||||||
|
.tool_resolver
|
||||||
|
.as_ref()
|
||||||
|
.ok_or("Tool invocation not available: no tool executor configured")?;
|
||||||
|
|
||||||
|
// Parse parameters
|
||||||
|
let params: serde_json::Value = serde_json::from_str(¶ms_json)
|
||||||
|
.map_err(|e| format!("Invalid tool parameters JSON: {}", e))?;
|
||||||
|
|
||||||
|
// Increment depth with RAII guard to ensure decrement even on panic
|
||||||
|
self.tool_nesting_depth += 1;
|
||||||
|
let current_depth = self.tool_nesting_depth;
|
||||||
|
let _guard = NestingGuard {
|
||||||
|
depth: &mut self.tool_nesting_depth,
|
||||||
|
};
|
||||||
|
// _guard drops at end of scope (or on panic), decrementing depth
|
||||||
|
resolver(&real_name, params, current_depth)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn secret_exists(&mut self, name: String) -> bool {
|
fn secret_exists(&mut self, name: String) -> bool {
|
||||||
@@ -579,9 +530,11 @@ pub struct WasmToolWrapper {
|
|||||||
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
||||||
oauth_refresh: Option<OAuthRefreshConfig>,
|
oauth_refresh: Option<OAuthRefreshConfig>,
|
||||||
/// Optional HTTP interceptor for testing — returns canned responses
|
/// Direct tool executor reference (for tests that wire it explicitly).
|
||||||
/// instead of making real requests when set.
|
tool_executor: Option<Arc<ToolExecutor>>,
|
||||||
http_interceptor: Option<Arc<dyn HttpInterceptor>>,
|
/// Shared slot for lazy executor resolution (production path).
|
||||||
|
/// Reads happen inside `spawn_blocking`, so this uses `std::sync::RwLock`.
|
||||||
|
tool_executor_slot: Option<Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -608,51 +561,23 @@ impl WasmToolSchemas {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn is_permissive_schema(schema: &serde_json::Value) -> bool {
|
fn is_permissive_schema(schema: &serde_json::Value) -> bool {
|
||||||
if schema
|
schema
|
||||||
.get("properties")
|
.get("properties")
|
||||||
.and_then(|p| p.as_object())
|
.and_then(|p| p.as_object())
|
||||||
.is_some_and(|p| !p.is_empty())
|
.is_none_or(|p| p.is_empty())
|
||||||
{
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Schemas with combinator variants containing properties are not permissive
|
|
||||||
for key in ["oneOf", "anyOf", "allOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array())
|
|
||||||
&& variants.iter().any(|v| {
|
|
||||||
v.get("properties")
|
|
||||||
.and_then(|p| p.as_object())
|
|
||||||
.is_some_and(|p| !p.is_empty())
|
|
||||||
})
|
|
||||||
{
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
true
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn typed_property_count(schema: &serde_json::Value) -> usize {
|
fn typed_property_count(schema: &serde_json::Value) -> usize {
|
||||||
let mut all_props = serde_json::Map::new();
|
schema
|
||||||
|
.get("properties")
|
||||||
if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
|
.and_then(|p| p.as_object())
|
||||||
all_props.extend(props.iter().map(|(k, v)| (k.clone(), v.clone())));
|
.map(|props| {
|
||||||
}
|
props
|
||||||
|
.values()
|
||||||
for key in ["allOf", "oneOf", "anyOf"] {
|
.filter(|prop| schema_is_typed_property(prop))
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array()) {
|
.count()
|
||||||
for variant in variants {
|
})
|
||||||
if let Some(props) = variant.get("properties").and_then(|p| p.as_object()) {
|
.unwrap_or(0)
|
||||||
all_props.extend(props.iter().map(|(k, v)| (k.clone(), v.clone())));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
all_props
|
|
||||||
.values()
|
|
||||||
.filter(|prop| schema_is_typed_property(prop))
|
|
||||||
.count()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn new(discovery: serde_json::Value) -> Self {
|
fn new(discovery: serde_json::Value) -> Self {
|
||||||
@@ -698,20 +623,11 @@ impl WasmToolWrapper {
|
|||||||
credentials: HashMap::new(),
|
credentials: HashMap::new(),
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
oauth_refresh: None,
|
oauth_refresh: None,
|
||||||
http_interceptor: None,
|
tool_executor: None,
|
||||||
|
tool_executor_slot: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set an HTTP interceptor for testing.
|
|
||||||
///
|
|
||||||
/// When set, WASM tool HTTP requests are routed through the interceptor
|
|
||||||
/// instead of making real network calls. This allows tests to verify the
|
|
||||||
/// exact HTTP requests a WASM tool constructs.
|
|
||||||
pub fn with_http_interceptor(mut self, interceptor: Arc<dyn HttpInterceptor>) -> Self {
|
|
||||||
self.http_interceptor = Some(interceptor);
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Override the tool description.
|
/// Override the tool description.
|
||||||
pub fn with_description(mut self, description: impl Into<String>) -> Self {
|
pub fn with_description(mut self, description: impl Into<String>) -> Self {
|
||||||
self.description = description.into();
|
self.description = description.into();
|
||||||
@@ -763,6 +679,28 @@ impl WasmToolWrapper {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Set the tool executor for programmatic tool calling (direct reference).
|
||||||
|
///
|
||||||
|
/// When set, the WASM `tool_invoke` host function can call other
|
||||||
|
/// registered tools synchronously via a bridged resolver closure.
|
||||||
|
/// Prefer `with_tool_executor_slot()` for production use.
|
||||||
|
pub fn with_tool_executor(mut self, executor: Arc<ToolExecutor>) -> Self {
|
||||||
|
self.tool_executor = Some(executor);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Set the shared tool executor slot for lazy resolution.
|
||||||
|
///
|
||||||
|
/// The executor is read from this slot at execution time, allowing
|
||||||
|
/// it to be set after tool registration (production startup order).
|
||||||
|
pub fn with_tool_executor_slot(
|
||||||
|
mut self,
|
||||||
|
slot: Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>,
|
||||||
|
) -> Self {
|
||||||
|
self.tool_executor_slot = Some(slot);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
/// Get the resource limits for this tool.
|
/// Get the resource limits for this tool.
|
||||||
pub fn limits(&self) -> &ResourceLimits {
|
pub fn limits(&self) -> &ResourceLimits {
|
||||||
&self.prepared.limits
|
&self.prepared.limits
|
||||||
@@ -791,18 +729,19 @@ impl WasmToolWrapper {
|
|||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
context_json: Option<String>,
|
context_json: Option<String>,
|
||||||
host_credentials: Vec<ResolvedHostCredential>,
|
host_credentials: Vec<ResolvedHostCredential>,
|
||||||
|
tool_resolver: Option<ToolResolver>,
|
||||||
) -> Result<(String, Vec<crate::tools::wasm::host::LogEntry>), WasmError> {
|
) -> Result<(String, Vec<crate::tools::wasm::host::LogEntry>), WasmError> {
|
||||||
let engine = self.runtime.engine();
|
let engine = self.runtime.engine();
|
||||||
let limits = &self.prepared.limits;
|
let limits = &self.prepared.limits;
|
||||||
|
|
||||||
// Create store with fresh state (NEAR pattern: fresh instance per call)
|
// Create store with fresh state (NEAR pattern: fresh instance per call)
|
||||||
let mut store_data = StoreData::new(
|
let store_data = StoreData::new(
|
||||||
limits.memory_bytes,
|
limits.memory_bytes,
|
||||||
self.capabilities.clone(),
|
self.capabilities.clone(),
|
||||||
self.credentials.clone(),
|
self.credentials.clone(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
|
tool_resolver,
|
||||||
);
|
);
|
||||||
store_data.http_interceptor = self.http_interceptor.clone();
|
|
||||||
let mut store = Store::new(engine, store_data);
|
let mut store = Store::new(engine, store_data);
|
||||||
|
|
||||||
// Configure fuel if enabled
|
// Configure fuel if enabled
|
||||||
@@ -900,6 +839,7 @@ pub(super) fn extract_wasm_metadata(
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
vec![],
|
vec![],
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
let mut store = Store::new(engine, store_data);
|
let mut store = Store::new(engine, store_data);
|
||||||
|
|
||||||
@@ -999,6 +939,48 @@ impl Tool for WasmToolWrapper {
|
|||||||
// Serialize context for WASM
|
// Serialize context for WASM
|
||||||
let context_json = serde_json::to_string(ctx).ok();
|
let context_json = serde_json::to_string(ctx).ok();
|
||||||
|
|
||||||
|
// Resolve the tool executor: direct reference takes priority, then shared slot.
|
||||||
|
let resolved_executor: Option<Arc<ToolExecutor>> =
|
||||||
|
self.tool_executor.as_ref().cloned().or_else(|| {
|
||||||
|
self.tool_executor_slot
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|slot| slot.read().ok())
|
||||||
|
.and_then(|guard| guard.clone())
|
||||||
|
});
|
||||||
|
|
||||||
|
// Build a tool resolver closure if we have a tool executor.
|
||||||
|
// The resolver creates a single-threaded tokio runtime (same pattern
|
||||||
|
// as http_request) to bridge the sync WASM callback to async tool execution.
|
||||||
|
let tool_resolver: Option<ToolResolver> = resolved_executor.as_ref().map(|executor| {
|
||||||
|
let executor = Arc::clone(executor);
|
||||||
|
let user_id = ctx.user_id.clone();
|
||||||
|
Arc::new(move |name: &str, params: serde_json::Value, depth: u32| {
|
||||||
|
let executor = Arc::clone(&executor);
|
||||||
|
let name = name.to_string();
|
||||||
|
let mut ctx = JobContext::with_user(
|
||||||
|
user_id.clone(),
|
||||||
|
format!("WASM PTC: {}", name),
|
||||||
|
"Programmatic tool call from WASM tool".to_string(),
|
||||||
|
);
|
||||||
|
// Propagate depth so the executor enforces the global limit
|
||||||
|
ctx.tool_nesting_depth = depth;
|
||||||
|
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.map_err(|e| format!("Failed to create runtime: {}", e))?;
|
||||||
|
|
||||||
|
rt.block_on(async {
|
||||||
|
executor
|
||||||
|
.execute(&name, params, &ctx, None)
|
||||||
|
.await
|
||||||
|
.map(|r| r.output)
|
||||||
|
.map_err(|e| e.to_string())
|
||||||
|
})
|
||||||
|
})
|
||||||
|
as Arc<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + Send + Sync>
|
||||||
|
});
|
||||||
|
|
||||||
// Clone what we need for the blocking task
|
// Clone what we need for the blocking task
|
||||||
let runtime = Arc::clone(&self.runtime);
|
let runtime = Arc::clone(&self.runtime);
|
||||||
let prepared = Arc::clone(&self.prepared);
|
let prepared = Arc::clone(&self.prepared);
|
||||||
@@ -1016,13 +998,14 @@ impl Tool for WasmToolWrapper {
|
|||||||
description,
|
description,
|
||||||
schemas,
|
schemas,
|
||||||
credentials,
|
credentials,
|
||||||
secrets_store: None, // Not needed in blocking task
|
secrets_store: None, // Not needed in blocking task
|
||||||
oauth_refresh: None, // Already used above for pre-refresh
|
oauth_refresh: None, // Already used above for pre-refresh
|
||||||
http_interceptor: self.http_interceptor.clone(),
|
tool_executor: None, // Resolver closure captures the executor
|
||||||
|
tool_executor_slot: None, // Resolver closure captures the executor
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
wrapper.execute_sync(params, context_json, host_credentials)
|
wrapper.execute_sync(params, context_json, host_credentials, tool_resolver)
|
||||||
})
|
})
|
||||||
.await
|
.await
|
||||||
.map_err(|e| WasmError::ExecutionPanicked(e.to_string()))?
|
.map_err(|e| WasmError::ExecutionPanicked(e.to_string()))?
|
||||||
@@ -1467,33 +1450,15 @@ fn is_private_ip(ip: std::net::IpAddr) -> bool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn schema_contains_container_properties(schema: &serde_json::Value) -> bool {
|
fn schema_contains_container_properties(schema: &serde_json::Value) -> bool {
|
||||||
let has_container = |props: &serde_json::Map<String, serde_json::Value>| {
|
schema
|
||||||
props
|
|
||||||
.values()
|
|
||||||
.any(|prop| schema_declares_type(prop, "array") || schema_declares_type(prop, "object"))
|
|
||||||
};
|
|
||||||
|
|
||||||
if schema
|
|
||||||
.get("properties")
|
.get("properties")
|
||||||
.and_then(|p| p.as_object())
|
.and_then(|p| p.as_object())
|
||||||
.is_some_and(has_container)
|
.map(|props| {
|
||||||
{
|
props.values().any(|prop| {
|
||||||
return true;
|
schema_declares_type(prop, "array") || schema_declares_type(prop, "object")
|
||||||
}
|
|
||||||
|
|
||||||
for key in ["allOf", "oneOf", "anyOf"] {
|
|
||||||
if let Some(variants) = schema.get(key).and_then(|v| v.as_array())
|
|
||||||
&& variants.iter().any(|v| {
|
|
||||||
v.get("properties")
|
|
||||||
.and_then(|p| p.as_object())
|
|
||||||
.is_some_and(has_container)
|
|
||||||
})
|
})
|
||||||
{
|
})
|
||||||
return true;
|
.unwrap_or(false)
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
false
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn schema_declares_type(schema: &serde_json::Value, expected: &str) -> bool {
|
fn schema_declares_type(schema: &serde_json::Value, expected: &str) -> bool {
|
||||||
@@ -1551,6 +1516,7 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
@@ -1567,10 +1533,12 @@ mod tests {
|
|||||||
TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET,
|
TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET,
|
||||||
test_secrets_store,
|
test_secrets_store,
|
||||||
};
|
};
|
||||||
use crate::tools::tool::Tool;
|
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||||
use crate::tools::wasm::capabilities::Capabilities;
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
||||||
|
|
||||||
|
use super::WasmToolWrapper;
|
||||||
|
|
||||||
struct RecordingSecretsStore {
|
struct RecordingSecretsStore {
|
||||||
inner: InMemorySecretsStore,
|
inner: InMemorySecretsStore,
|
||||||
get_decrypted_lookups: Mutex<Vec<(String, String)>>,
|
get_decrypted_lookups: Mutex<Vec<(String, String)>>,
|
||||||
@@ -1798,6 +1766,7 @@ mod tests {
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
|
|
||||||
// Should inject for matching host
|
// Should inject for matching host
|
||||||
@@ -1837,6 +1806,7 @@ mod tests {
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
|
|
||||||
let mut headers = HashMap::new();
|
let mut headers = HashMap::new();
|
||||||
@@ -1863,6 +1833,7 @@ mod tests {
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
|
None,
|
||||||
);
|
);
|
||||||
|
|
||||||
let text = "Error: request to https://api.example.com?key=super-secret-token failed";
|
let text = "Error: request to https://api.example.com?key=super-secret-token failed";
|
||||||
@@ -2349,6 +2320,244 @@ mod tests {
|
|||||||
assert!(hint.contains("native JSON arrays/objects")); // safety: test-only assertion
|
assert!(hint.contains("native JSON arrays/objects")); // safety: test-only assertion
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_coerce_params_already_correct_type() {
|
||||||
|
let schema = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"count": { "type": "number" }
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let params = serde_json::json!({"count": 5});
|
||||||
|
let result = crate::tools::coercion::prepare_params_for_schema(¶ms, &schema);
|
||||||
|
assert_eq!(result["count"], serde_json::json!(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_coerce_params_invalid_string_not_coerced() {
|
||||||
|
let schema = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"count": { "type": "number" }
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let params = serde_json::json!({"count": "not-a-number"});
|
||||||
|
let result = crate::tools::coercion::prepare_params_for_schema(¶ms, &schema);
|
||||||
|
// Should remain as string since it can't be parsed
|
||||||
|
assert_eq!(result["count"], serde_json::json!("not-a-number"));
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Programmatic Tool Calling (PTC) integration tests ===
|
||||||
|
//
|
||||||
|
// These tests require the test-ptc WASM binary to be pre-built:
|
||||||
|
// cargo build --target wasm32-wasip2 --release --manifest-path tools-src/test-ptc/Cargo.toml
|
||||||
|
|
||||||
|
use crate::config::SafetyConfig;
|
||||||
|
use crate::tools::executor::ToolExecutor;
|
||||||
|
|
||||||
|
fn wasm_binary_path() -> std::path::PathBuf {
|
||||||
|
let manifest_dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR"));
|
||||||
|
manifest_dir.join("tools-src/test-ptc/target/wasm32-wasip2/release/test_ptc_tool.wasm")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn load_wasm_binary() -> Option<Vec<u8>> {
|
||||||
|
let path = wasm_binary_path();
|
||||||
|
if !path.exists() {
|
||||||
|
eprintln!(
|
||||||
|
"WASM test binary not found at {:?}. Build with: \
|
||||||
|
cargo build --target wasm32-wasip2 --release --manifest-path tools-src/test-ptc/Cargo.toml",
|
||||||
|
path
|
||||||
|
);
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(std::fs::read(&path).expect("failed to read WASM binary"))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore]
|
||||||
|
async fn test_wasm_tool_invoke_echo() {
|
||||||
|
let wasm_bytes = match load_wasm_binary() {
|
||||||
|
Some(b) => b,
|
||||||
|
None => return, // Skip if binary not built
|
||||||
|
};
|
||||||
|
|
||||||
|
// Set up runtime
|
||||||
|
let runtime = Arc::new(
|
||||||
|
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
||||||
|
.expect("failed to create WASM runtime"),
|
||||||
|
);
|
||||||
|
|
||||||
|
// Set up tool registry with echo
|
||||||
|
let tools = Arc::new(crate::tools::registry::ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(crate::safety::SafetyLayer::new(&SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
}));
|
||||||
|
let executor = Arc::new(ToolExecutor::new(
|
||||||
|
tools,
|
||||||
|
safety,
|
||||||
|
std::time::Duration::from_secs(60),
|
||||||
|
));
|
||||||
|
|
||||||
|
// Prepare WASM module
|
||||||
|
let prepared = runtime
|
||||||
|
.prepare("test_ptc", &wasm_bytes, None)
|
||||||
|
.await
|
||||||
|
.expect("failed to prepare WASM module");
|
||||||
|
|
||||||
|
// Build capabilities with echo_alias -> echo
|
||||||
|
let mut aliases = HashMap::new();
|
||||||
|
aliases.insert("echo_alias".to_string(), "echo".to_string());
|
||||||
|
let capabilities = Capabilities::default().with_tool_invoke(aliases);
|
||||||
|
|
||||||
|
// Create wrapper with executor
|
||||||
|
let wrapper =
|
||||||
|
WasmToolWrapper::new(runtime, prepared, capabilities).with_tool_executor(executor);
|
||||||
|
|
||||||
|
// Execute
|
||||||
|
let ctx = crate::context::JobContext::new("test", "WASM PTC test");
|
||||||
|
let result: Result<ToolOutput, ToolError> = wrapper
|
||||||
|
.execute(serde_json::json!({"message": "hello"}), &ctx)
|
||||||
|
.await;
|
||||||
|
let result = result.expect("WASM tool execution should succeed");
|
||||||
|
|
||||||
|
let output = result.result.as_str().unwrap_or("");
|
||||||
|
assert!(
|
||||||
|
output.contains("via_wasm:"),
|
||||||
|
"Output should contain 'via_wasm:' prefix, got: {}",
|
||||||
|
output
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
output.contains("hello"),
|
||||||
|
"Output should contain 'hello', got: {}",
|
||||||
|
output
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore]
|
||||||
|
async fn test_wasm_tool_invoke_alias_not_granted() {
|
||||||
|
let wasm_bytes = match load_wasm_binary() {
|
||||||
|
Some(b) => b,
|
||||||
|
None => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
let runtime = Arc::new(
|
||||||
|
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
||||||
|
.expect("failed to create WASM runtime"),
|
||||||
|
);
|
||||||
|
|
||||||
|
let tools = Arc::new(crate::tools::registry::ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(crate::safety::SafetyLayer::new(&SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
}));
|
||||||
|
let executor = Arc::new(ToolExecutor::new(
|
||||||
|
tools,
|
||||||
|
safety,
|
||||||
|
std::time::Duration::from_secs(60),
|
||||||
|
));
|
||||||
|
|
||||||
|
let prepared = runtime
|
||||||
|
.prepare("test_ptc", &wasm_bytes, None)
|
||||||
|
.await
|
||||||
|
.expect("failed to prepare WASM module");
|
||||||
|
|
||||||
|
// Only grant a DIFFERENT alias, not "echo_alias"
|
||||||
|
let mut aliases = HashMap::new();
|
||||||
|
aliases.insert("other_alias".to_string(), "echo".to_string());
|
||||||
|
let capabilities = Capabilities::default().with_tool_invoke(aliases);
|
||||||
|
|
||||||
|
let wrapper =
|
||||||
|
WasmToolWrapper::new(runtime, prepared, capabilities).with_tool_executor(executor);
|
||||||
|
|
||||||
|
let ctx = crate::context::JobContext::new("test", "WASM PTC test");
|
||||||
|
let result: Result<ToolOutput, ToolError> = wrapper
|
||||||
|
.execute(serde_json::json!({"message": "hello"}), &ctx)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Should fail because "echo_alias" is not in the aliases
|
||||||
|
assert!(result.is_err(), "Should fail when alias not granted");
|
||||||
|
let err_msg = format!("{:?}", result.unwrap_err());
|
||||||
|
assert!(
|
||||||
|
err_msg.contains("Unknown tool alias") || err_msg.contains("echo_alias"),
|
||||||
|
"Error should mention unknown alias, got: {}",
|
||||||
|
err_msg
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore]
|
||||||
|
async fn test_wasm_tool_invoke_no_capability() {
|
||||||
|
let wasm_bytes = match load_wasm_binary() {
|
||||||
|
Some(b) => b,
|
||||||
|
None => return,
|
||||||
|
};
|
||||||
|
|
||||||
|
let runtime = Arc::new(
|
||||||
|
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
||||||
|
.expect("failed to create WASM runtime"),
|
||||||
|
);
|
||||||
|
|
||||||
|
let tools = Arc::new(crate::tools::registry::ToolRegistry::new());
|
||||||
|
tools.register_builtin_tools();
|
||||||
|
let safety = Arc::new(crate::safety::SafetyLayer::new(&SafetyConfig {
|
||||||
|
max_output_length: 100_000,
|
||||||
|
injection_check_enabled: true,
|
||||||
|
}));
|
||||||
|
let executor = Arc::new(ToolExecutor::new(
|
||||||
|
tools,
|
||||||
|
safety,
|
||||||
|
std::time::Duration::from_secs(60),
|
||||||
|
));
|
||||||
|
|
||||||
|
let prepared = runtime
|
||||||
|
.prepare("test_ptc", &wasm_bytes, None)
|
||||||
|
.await
|
||||||
|
.expect("failed to prepare WASM module");
|
||||||
|
|
||||||
|
// No tool_invoke capability at all
|
||||||
|
let capabilities = Capabilities::default();
|
||||||
|
|
||||||
|
let wrapper =
|
||||||
|
WasmToolWrapper::new(runtime, prepared, capabilities).with_tool_executor(executor);
|
||||||
|
|
||||||
|
let ctx = crate::context::JobContext::new("test", "WASM PTC test");
|
||||||
|
let result: Result<ToolOutput, ToolError> = wrapper
|
||||||
|
.execute(serde_json::json!({"message": "hello"}), &ctx)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
result.is_err(),
|
||||||
|
"Should fail when no tool_invoke capability"
|
||||||
|
);
|
||||||
|
let err_msg = format!("{:?}", result.unwrap_err());
|
||||||
|
assert!(
|
||||||
|
err_msg.contains("not granted") || err_msg.contains("capability"),
|
||||||
|
"Error should mention capability not granted, got: {}",
|
||||||
|
err_msg
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Regression: permissive fallback schema (empty properties) must NOT coerce.
|
||||||
|
/// This documents the bug where WASM tools with no sidecar `parameters` field
|
||||||
|
/// got the permissive fallback, causing coercion to be a no-op and LLM-provided
|
||||||
|
/// string integers to reach the WASM tool un-coerced.
|
||||||
|
#[test]
|
||||||
|
fn test_coerce_noop_with_permissive_schema() {
|
||||||
|
let permissive = serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {},
|
||||||
|
"additionalProperties": true
|
||||||
|
});
|
||||||
|
let params = serde_json::json!({"query": "test", "count": "10"});
|
||||||
|
let result = crate::tools::coercion::prepare_params_for_schema(¶ms, &permissive);
|
||||||
|
// With empty properties, no coercion happens — string stays string
|
||||||
|
assert_eq!(result["count"], serde_json::json!("10"));
|
||||||
|
}
|
||||||
|
|
||||||
/// Regression test: leak scan must run on raw headers (before credential
|
/// Regression test: leak scan must run on raw headers (before credential
|
||||||
/// injection), not after. If it ran post-injection, the host-injected
|
/// injection), not after. If it ran post-injection, the host-injected
|
||||||
/// Slack bot token (`xoxb-...`) would trigger a Block and reject the
|
/// Slack bot token (`xoxb-...`) would trigger a Block and reject the
|
||||||
|
|||||||
@@ -114,6 +114,36 @@ pub struct CredentialResponse {
|
|||||||
pub value: String,
|
pub value: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Request to call a tool programmatically via the orchestrator.
|
||||||
|
#[derive(Debug, Serialize, Deserialize)]
|
||||||
|
pub struct ToolCallRequest {
|
||||||
|
/// Name of the tool to invoke.
|
||||||
|
pub tool_name: String,
|
||||||
|
/// JSON parameters to pass to the tool.
|
||||||
|
pub parameters: serde_json::Value,
|
||||||
|
/// Optional timeout in seconds (capped at 300s by the orchestrator).
|
||||||
|
pub timeout_secs: Option<u64>,
|
||||||
|
/// Current nesting depth for tool-invokes-tool chains.
|
||||||
|
/// Defaults to 0 for top-level calls (backward compatible).
|
||||||
|
#[serde(default)]
|
||||||
|
pub nesting_depth: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Response from a programmatic tool call.
|
||||||
|
#[derive(Debug, Serialize, Deserialize)]
|
||||||
|
pub struct ToolCallResponse {
|
||||||
|
/// Whether the tool call succeeded.
|
||||||
|
pub success: bool,
|
||||||
|
/// Tool output (present on success).
|
||||||
|
pub output: Option<String>,
|
||||||
|
/// Error message (present on failure).
|
||||||
|
pub error: Option<String>,
|
||||||
|
/// Execution duration in milliseconds.
|
||||||
|
pub duration_ms: u64,
|
||||||
|
/// Whether the output was modified by the safety layer.
|
||||||
|
pub was_sanitized: bool,
|
||||||
|
}
|
||||||
|
|
||||||
impl WorkerHttpClient {
|
impl WorkerHttpClient {
|
||||||
/// Create a new client from environment.
|
/// Create a new client from environment.
|
||||||
///
|
///
|
||||||
@@ -399,6 +429,14 @@ impl WorkerHttpClient {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Call a tool programmatically via the orchestrator (PTC).
|
||||||
|
///
|
||||||
|
/// This bypasses the LLM round-trip and invokes a tool directly on the
|
||||||
|
/// orchestrator side. Useful for scripted multi-step sequences.
|
||||||
|
pub async fn call_tool(&self, req: &ToolCallRequest) -> Result<ToolCallResponse, WorkerError> {
|
||||||
|
self.post_json("tools/call", req, "tool call").await
|
||||||
|
}
|
||||||
|
|
||||||
/// Signal job completion to the orchestrator.
|
/// Signal job completion to the orchestrator.
|
||||||
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
|
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
|
||||||
let _: serde_json::Value = self
|
let _: serde_json::Value = self
|
||||||
|
|||||||
@@ -343,412 +343,4 @@ mod tests {
|
|||||||
|
|
||||||
rig.shutdown();
|
rig.shutdown();
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Fixture tool that mirrors the github WASM tool's `oneOf` discriminated
|
|
||||||
/// union schema. Uses `#[serde(tag = "action")]` deserialization — exactly
|
|
||||||
/// what the real tool does — so if coercion fails the test reproduces:
|
|
||||||
/// `invalid type: string "100", expected u32`
|
|
||||||
struct GitHubFixtureTool;
|
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
|
||||||
#[serde(tag = "action")]
|
|
||||||
enum GitHubFixtureAction {
|
|
||||||
#[serde(rename = "list_issues")]
|
|
||||||
ListIssues {
|
|
||||||
owner: String,
|
|
||||||
repo: String,
|
|
||||||
#[serde(default)]
|
|
||||||
state: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
limit: Option<u32>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "get_issue")]
|
|
||||||
GetIssue {
|
|
||||||
owner: String,
|
|
||||||
repo: String,
|
|
||||||
issue_number: u32,
|
|
||||||
},
|
|
||||||
#[serde(rename = "list_pull_requests")]
|
|
||||||
ListPullRequests {
|
|
||||||
owner: String,
|
|
||||||
repo: String,
|
|
||||||
#[serde(default)]
|
|
||||||
limit: Option<u32>,
|
|
||||||
#[serde(default)]
|
|
||||||
page: Option<u32>,
|
|
||||||
},
|
|
||||||
#[serde(rename = "create_pull_request")]
|
|
||||||
CreatePullRequest {
|
|
||||||
owner: String,
|
|
||||||
repo: String,
|
|
||||||
title: String,
|
|
||||||
head: String,
|
|
||||||
base: String,
|
|
||||||
#[serde(default)]
|
|
||||||
draft: Option<bool>,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
use serde::Deserialize;
|
|
||||||
|
|
||||||
#[async_trait]
|
|
||||||
impl Tool for GitHubFixtureTool {
|
|
||||||
fn name(&self) -> &str {
|
|
||||||
"github_fixture"
|
|
||||||
}
|
|
||||||
|
|
||||||
fn description(&self) -> &str {
|
|
||||||
"Fixture mirroring the github WASM tool's oneOf schema"
|
|
||||||
}
|
|
||||||
|
|
||||||
fn parameters_schema(&self) -> serde_json::Value {
|
|
||||||
json!({
|
|
||||||
"type": "object",
|
|
||||||
"required": ["action"],
|
|
||||||
"oneOf": [
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "list_issues" },
|
|
||||||
"owner": { "type": "string" },
|
|
||||||
"repo": { "type": "string" },
|
|
||||||
"state": { "type": "string", "enum": ["open", "closed", "all"] },
|
|
||||||
"limit": { "type": "integer", "default": 30 }
|
|
||||||
},
|
|
||||||
"required": ["action", "owner", "repo"]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "get_issue" },
|
|
||||||
"owner": { "type": "string" },
|
|
||||||
"repo": { "type": "string" },
|
|
||||||
"issue_number": { "type": "integer" }
|
|
||||||
},
|
|
||||||
"required": ["action", "owner", "repo", "issue_number"]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "list_pull_requests" },
|
|
||||||
"owner": { "type": "string" },
|
|
||||||
"repo": { "type": "string" },
|
|
||||||
"limit": { "type": "integer", "default": 30 },
|
|
||||||
"page": { "type": "integer" }
|
|
||||||
},
|
|
||||||
"required": ["action", "owner", "repo"]
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"properties": {
|
|
||||||
"action": { "const": "create_pull_request" },
|
|
||||||
"owner": { "type": "string" },
|
|
||||||
"repo": { "type": "string" },
|
|
||||||
"title": { "type": "string" },
|
|
||||||
"head": { "type": "string" },
|
|
||||||
"base": { "type": "string" },
|
|
||||||
"draft": { "type": "boolean", "default": false }
|
|
||||||
},
|
|
||||||
"required": ["action", "owner", "repo", "title", "head", "base"]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn execute(
|
|
||||||
&self,
|
|
||||||
params: serde_json::Value,
|
|
||||||
_ctx: &JobContext,
|
|
||||||
) -> Result<ToolOutput, ToolError> {
|
|
||||||
// Deserialize exactly like the real github WASM tool does.
|
|
||||||
// Without coercion, this fails: `invalid type: string "100", expected u32`
|
|
||||||
let action: GitHubFixtureAction = serde_json::from_value(params).map_err(|e| {
|
|
||||||
ToolError::InvalidParameters(format!("serde deserialization failed: {e}"))
|
|
||||||
})?;
|
|
||||||
|
|
||||||
let result = match action {
|
|
||||||
GitHubFixtureAction::ListIssues {
|
|
||||||
owner,
|
|
||||||
repo,
|
|
||||||
state,
|
|
||||||
limit,
|
|
||||||
} => json!({
|
|
||||||
"action": "list_issues",
|
|
||||||
"owner": owner,
|
|
||||||
"repo": repo,
|
|
||||||
"state": state.unwrap_or_else(|| "open".to_string()),
|
|
||||||
"limit": limit.unwrap_or(30),
|
|
||||||
}),
|
|
||||||
GitHubFixtureAction::GetIssue {
|
|
||||||
owner,
|
|
||||||
repo,
|
|
||||||
issue_number,
|
|
||||||
} => json!({
|
|
||||||
"action": "get_issue",
|
|
||||||
"owner": owner,
|
|
||||||
"repo": repo,
|
|
||||||
"issue_number": issue_number,
|
|
||||||
}),
|
|
||||||
GitHubFixtureAction::ListPullRequests {
|
|
||||||
owner,
|
|
||||||
repo,
|
|
||||||
limit,
|
|
||||||
page,
|
|
||||||
} => json!({
|
|
||||||
"action": "list_pull_requests",
|
|
||||||
"owner": owner,
|
|
||||||
"repo": repo,
|
|
||||||
"limit": limit.unwrap_or(30),
|
|
||||||
"page": page.unwrap_or(1),
|
|
||||||
}),
|
|
||||||
GitHubFixtureAction::CreatePullRequest {
|
|
||||||
owner,
|
|
||||||
repo,
|
|
||||||
title,
|
|
||||||
head,
|
|
||||||
base,
|
|
||||||
draft,
|
|
||||||
} => json!({
|
|
||||||
"action": "create_pull_request",
|
|
||||||
"owner": owner,
|
|
||||||
"repo": repo,
|
|
||||||
"title": title,
|
|
||||||
"head": head,
|
|
||||||
"base": base,
|
|
||||||
"draft": draft.unwrap_or(false),
|
|
||||||
}),
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(ToolOutput::success(result, Duration::from_millis(1)))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn requires_sanitization(&self) -> bool {
|
|
||||||
false
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Reproduces the exact bug: LLM sends `limit: "100"` and `issue_number: "42"`
|
|
||||||
/// as strings to a `oneOf` discriminated union schema. Without coercion support
|
|
||||||
/// for combinators, serde fails with `invalid type: string "100", expected u32`.
|
|
||||||
#[tokio::test]
|
|
||||||
async fn e2e_coerces_oneof_discriminated_union_params() {
|
|
||||||
let trace = LlmTrace {
|
|
||||||
model_name: "test-coercion-oneof".to_string(),
|
|
||||||
turns: vec![crate::support::trace_llm::TraceTurn {
|
|
||||||
user_input: "List issues in nearai/ironclaw with limit 100".to_string(),
|
|
||||||
steps: vec![
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::ToolCalls {
|
|
||||||
tool_calls: vec![TraceToolCall {
|
|
||||||
id: "call_gh_list".to_string(),
|
|
||||||
name: "github_fixture".to_string(),
|
|
||||||
// LLM sends numeric params as strings — the exact bug
|
|
||||||
arguments: json!({
|
|
||||||
"action": "list_issues",
|
|
||||||
"owner": "nearai",
|
|
||||||
"repo": "ironclaw",
|
|
||||||
"state": "open",
|
|
||||||
"limit": "100"
|
|
||||||
}),
|
|
||||||
}],
|
|
||||||
input_tokens: 100,
|
|
||||||
output_tokens: 30,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::Text {
|
|
||||||
content: "Found issues in nearai/ironclaw with limit 100.".to_string(),
|
|
||||||
input_tokens: 150,
|
|
||||||
output_tokens: 20,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
],
|
|
||||||
expects: TraceExpects::default(),
|
|
||||||
}],
|
|
||||||
memory_snapshot: Vec::new(),
|
|
||||||
http_exchanges: Vec::new(),
|
|
||||||
expects: TraceExpects {
|
|
||||||
tools_used: vec!["github_fixture".to_string()],
|
|
||||||
all_tools_succeeded: Some(true),
|
|
||||||
max_tool_calls: Some(1),
|
|
||||||
min_responses: Some(1),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
steps: Vec::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
let rig = TestRigBuilder::new()
|
|
||||||
.with_trace(trace.clone())
|
|
||||||
.with_extra_tools(vec![Arc::new(GitHubFixtureTool)])
|
|
||||||
.build()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
rig.send_message("List issues in nearai/ironclaw with limit 100")
|
|
||||||
.await;
|
|
||||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
|
||||||
|
|
||||||
rig.verify_trace_expects(&trace, &responses);
|
|
||||||
let tool_results = rig.tool_results();
|
|
||||||
assert!(
|
|
||||||
tool_results
|
|
||||||
.iter()
|
|
||||||
.any(|(name, preview)| name == "github_fixture"
|
|
||||||
&& preview.contains("\"limit\"")
|
|
||||||
&& preview.contains("100")),
|
|
||||||
"expected coerced list_issues result, got {tool_results:?}"
|
|
||||||
);
|
|
||||||
|
|
||||||
rig.shutdown();
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Tests a second oneOf variant with different string-to-integer coercions:
|
|
||||||
/// `issue_number: "42"` must be coerced to match the `get_issue` variant.
|
|
||||||
#[tokio::test]
|
|
||||||
async fn e2e_coerces_oneof_get_issue_variant() {
|
|
||||||
let trace = LlmTrace {
|
|
||||||
model_name: "test-coercion-oneof-issue".to_string(),
|
|
||||||
turns: vec![crate::support::trace_llm::TraceTurn {
|
|
||||||
user_input: "Get issue 42 from nearai/ironclaw".to_string(),
|
|
||||||
steps: vec![
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::ToolCalls {
|
|
||||||
tool_calls: vec![TraceToolCall {
|
|
||||||
id: "call_gh_issue".to_string(),
|
|
||||||
name: "github_fixture".to_string(),
|
|
||||||
arguments: json!({
|
|
||||||
"action": "get_issue",
|
|
||||||
"owner": "nearai",
|
|
||||||
"repo": "ironclaw",
|
|
||||||
"issue_number": "42"
|
|
||||||
}),
|
|
||||||
}],
|
|
||||||
input_tokens: 80,
|
|
||||||
output_tokens: 20,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::Text {
|
|
||||||
content: "Issue 42 retrieved.".to_string(),
|
|
||||||
input_tokens: 100,
|
|
||||||
output_tokens: 10,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
],
|
|
||||||
expects: TraceExpects::default(),
|
|
||||||
}],
|
|
||||||
memory_snapshot: Vec::new(),
|
|
||||||
http_exchanges: Vec::new(),
|
|
||||||
expects: TraceExpects {
|
|
||||||
tools_used: vec!["github_fixture".to_string()],
|
|
||||||
all_tools_succeeded: Some(true),
|
|
||||||
max_tool_calls: Some(1),
|
|
||||||
min_responses: Some(1),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
steps: Vec::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
let rig = TestRigBuilder::new()
|
|
||||||
.with_trace(trace.clone())
|
|
||||||
.with_extra_tools(vec![Arc::new(GitHubFixtureTool)])
|
|
||||||
.build()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
rig.send_message("Get issue 42 from nearai/ironclaw").await;
|
|
||||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
|
||||||
|
|
||||||
rig.verify_trace_expects(&trace, &responses);
|
|
||||||
let tool_results = rig.tool_results();
|
|
||||||
assert!(
|
|
||||||
tool_results
|
|
||||||
.iter()
|
|
||||||
.any(|(name, preview)| name == "github_fixture"
|
|
||||||
&& preview.contains("\"issue_number\"")
|
|
||||||
&& preview.contains("42")),
|
|
||||||
"expected coerced get_issue result, got {tool_results:?}"
|
|
||||||
);
|
|
||||||
|
|
||||||
rig.shutdown();
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Tests boolean coercion in a oneOf variant: `draft: "true"` must become
|
|
||||||
/// a boolean for the `create_pull_request` variant.
|
|
||||||
#[tokio::test]
|
|
||||||
async fn e2e_coerces_oneof_boolean_in_variant() {
|
|
||||||
let trace = LlmTrace {
|
|
||||||
model_name: "test-coercion-oneof-bool".to_string(),
|
|
||||||
turns: vec![crate::support::trace_llm::TraceTurn {
|
|
||||||
user_input: "Create a draft PR".to_string(),
|
|
||||||
steps: vec![
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::ToolCalls {
|
|
||||||
tool_calls: vec![TraceToolCall {
|
|
||||||
id: "call_gh_pr".to_string(),
|
|
||||||
name: "github_fixture".to_string(),
|
|
||||||
arguments: json!({
|
|
||||||
"action": "create_pull_request",
|
|
||||||
"owner": "nearai",
|
|
||||||
"repo": "ironclaw",
|
|
||||||
"title": "Fix coercion",
|
|
||||||
"head": "fix/coercion",
|
|
||||||
"base": "main",
|
|
||||||
"draft": "true"
|
|
||||||
}),
|
|
||||||
}],
|
|
||||||
input_tokens: 90,
|
|
||||||
output_tokens: 25,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::Text {
|
|
||||||
content: "Draft PR created.".to_string(),
|
|
||||||
input_tokens: 110,
|
|
||||||
output_tokens: 10,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
],
|
|
||||||
expects: TraceExpects::default(),
|
|
||||||
}],
|
|
||||||
memory_snapshot: Vec::new(),
|
|
||||||
http_exchanges: Vec::new(),
|
|
||||||
expects: TraceExpects {
|
|
||||||
tools_used: vec!["github_fixture".to_string()],
|
|
||||||
all_tools_succeeded: Some(true),
|
|
||||||
max_tool_calls: Some(1),
|
|
||||||
min_responses: Some(1),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
steps: Vec::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
let rig = TestRigBuilder::new()
|
|
||||||
.with_trace(trace.clone())
|
|
||||||
.with_extra_tools(vec![Arc::new(GitHubFixtureTool)])
|
|
||||||
.build()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
rig.send_message("Create a draft PR").await;
|
|
||||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
|
||||||
|
|
||||||
rig.verify_trace_expects(&trace, &responses);
|
|
||||||
let tool_results = rig.tool_results();
|
|
||||||
assert!(
|
|
||||||
tool_results
|
|
||||||
.iter()
|
|
||||||
.any(|(name, preview)| name == "github_fixture"
|
|
||||||
&& preview.contains("\"draft\"")
|
|
||||||
&& preview.contains("true")),
|
|
||||||
"expected coerced create_pull_request result with draft=true, got {tool_results:?}"
|
|
||||||
);
|
|
||||||
|
|
||||||
rig.shutdown();
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,277 +0,0 @@
|
|||||||
//! E2E test: real github WASM tool with parameter coercion via TestRig.
|
|
||||||
//!
|
|
||||||
//! Loads the compiled github WASM binary into the test rig, replays an LLM
|
|
||||||
//! trace that sends string-typed numeric params, and verifies the WASM tool
|
|
||||||
//! constructs the correct HTTP API call via `http_exchanges` in the trace.
|
|
||||||
//!
|
|
||||||
//! These tests are `#[ignore]` by default because they require a pre-compiled
|
|
||||||
//! WASM binary. Build it with:
|
|
||||||
//! cargo build -p github-tool --target wasm32-wasip2 --release
|
|
||||||
//! Then run with:
|
|
||||||
//! cargo test --features libsql --test e2e_wasm_github_coercion -- --ignored
|
|
||||||
|
|
||||||
#[cfg(feature = "libsql")]
|
|
||||||
mod support;
|
|
||||||
|
|
||||||
/// Note on URL verification: the `ReplayingHttpInterceptor` logs warnings on
|
|
||||||
/// URL mismatch but still returns the canned response. The real verification is
|
|
||||||
/// that the tool succeeds end-to-end: coercion produced the correct typed
|
|
||||||
/// parameters, serde deserialization succeeded, and the WASM tool constructed a
|
|
||||||
/// valid HTTP request. A URL mismatch warning in logs does not indicate test
|
|
||||||
/// failure — it is a soft check only.
|
|
||||||
#[cfg(feature = "libsql")]
|
|
||||||
mod tests {
|
|
||||||
use std::time::Duration;
|
|
||||||
|
|
||||||
use serde_json::json;
|
|
||||||
|
|
||||||
use ironclaw::llm::recording::{HttpExchange, HttpExchangeRequest, HttpExchangeResponse};
|
|
||||||
|
|
||||||
use crate::support::test_rig::TestRigBuilder;
|
|
||||||
use crate::support::trace_llm::{
|
|
||||||
LlmTrace, TraceExpects, TraceResponse, TraceStep, TraceToolCall,
|
|
||||||
};
|
|
||||||
|
|
||||||
const GITHUB_WASM: &str = "tools-src/github/target/wasm32-wasip2/release/github_tool.wasm";
|
|
||||||
const GITHUB_CAPS: &str = "tools-src/github/github-tool.capabilities.json";
|
|
||||||
|
|
||||||
fn github_ok(body: &str) -> HttpExchangeResponse {
|
|
||||||
HttpExchangeResponse {
|
|
||||||
status: 200,
|
|
||||||
headers: vec![
|
|
||||||
("content-type".to_string(), "application/json".to_string()),
|
|
||||||
("x-ratelimit-remaining".to_string(), "100".to_string()),
|
|
||||||
],
|
|
||||||
body: body.to_string(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// LLM sends `limit: "50"` (string) to `list_issues`. Coercion converts it
|
|
||||||
/// to integer, and the WASM tool must call `GET /repos/.../issues?...&per_page=50`.
|
|
||||||
#[tokio::test]
|
|
||||||
#[ignore] // requires pre-compiled WASM binary
|
|
||||||
async fn wasm_github_list_issues_coerces_string_limit() {
|
|
||||||
let expected_url =
|
|
||||||
"https://api.github.com/repos/nearai/ironclaw/issues?state=open&per_page=50";
|
|
||||||
|
|
||||||
let trace = LlmTrace {
|
|
||||||
model_name: "test-wasm-coercion-list-issues".to_string(),
|
|
||||||
turns: vec![crate::support::trace_llm::TraceTurn {
|
|
||||||
user_input: "List issues in nearai/ironclaw with limit 50".to_string(),
|
|
||||||
steps: vec![
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::ToolCalls {
|
|
||||||
tool_calls: vec![TraceToolCall {
|
|
||||||
id: "call_gh_1".to_string(),
|
|
||||||
name: "github".to_string(),
|
|
||||||
arguments: json!({
|
|
||||||
"action": "list_issues",
|
|
||||||
"owner": "nearai",
|
|
||||||
"repo": "ironclaw",
|
|
||||||
"state": "open",
|
|
||||||
"limit": "50"
|
|
||||||
}),
|
|
||||||
}],
|
|
||||||
input_tokens: 100,
|
|
||||||
output_tokens: 30,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::Text {
|
|
||||||
content: "Found 1 issue.".to_string(),
|
|
||||||
input_tokens: 150,
|
|
||||||
output_tokens: 10,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
],
|
|
||||||
expects: TraceExpects::default(),
|
|
||||||
}],
|
|
||||||
memory_snapshot: Vec::new(),
|
|
||||||
http_exchanges: vec![HttpExchange {
|
|
||||||
request: HttpExchangeRequest {
|
|
||||||
method: "GET".to_string(),
|
|
||||||
url: expected_url.to_string(),
|
|
||||||
headers: vec![],
|
|
||||||
body: None,
|
|
||||||
},
|
|
||||||
response: github_ok(r#"[{"number":1,"title":"Test issue","state":"open"}]"#),
|
|
||||||
}],
|
|
||||||
expects: TraceExpects {
|
|
||||||
tools_used: vec!["github".to_string()],
|
|
||||||
all_tools_succeeded: Some(true),
|
|
||||||
max_tool_calls: Some(1),
|
|
||||||
min_responses: Some(1),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
steps: Vec::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
let rig = TestRigBuilder::new()
|
|
||||||
.with_trace(trace.clone())
|
|
||||||
.with_wasm_tool("github", GITHUB_WASM, Some(GITHUB_CAPS.into()))
|
|
||||||
.build()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
rig.send_message("List issues in nearai/ironclaw with limit 50")
|
|
||||||
.await;
|
|
||||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
|
||||||
rig.verify_trace_expects(&trace, &responses);
|
|
||||||
|
|
||||||
rig.shutdown();
|
|
||||||
}
|
|
||||||
|
|
||||||
/// LLM sends `issue_number: "42"` (string) to `get_issue`. Coercion converts
|
|
||||||
/// it to integer, and the URL must contain `/issues/42`.
|
|
||||||
#[tokio::test]
|
|
||||||
#[ignore] // requires pre-compiled WASM binary
|
|
||||||
async fn wasm_github_get_issue_coerces_string_issue_number() {
|
|
||||||
let expected_url = "https://api.github.com/repos/nearai/ironclaw/issues/42";
|
|
||||||
|
|
||||||
let trace = LlmTrace {
|
|
||||||
model_name: "test-wasm-coercion-get-issue".to_string(),
|
|
||||||
turns: vec![crate::support::trace_llm::TraceTurn {
|
|
||||||
user_input: "Get issue 42 from nearai/ironclaw".to_string(),
|
|
||||||
steps: vec![
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::ToolCalls {
|
|
||||||
tool_calls: vec![TraceToolCall {
|
|
||||||
id: "call_gh_2".to_string(),
|
|
||||||
name: "github".to_string(),
|
|
||||||
arguments: json!({
|
|
||||||
"action": "get_issue",
|
|
||||||
"owner": "nearai",
|
|
||||||
"repo": "ironclaw",
|
|
||||||
"issue_number": "42"
|
|
||||||
}),
|
|
||||||
}],
|
|
||||||
input_tokens: 80,
|
|
||||||
output_tokens: 20,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::Text {
|
|
||||||
content: "Issue 42 retrieved.".to_string(),
|
|
||||||
input_tokens: 100,
|
|
||||||
output_tokens: 10,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
],
|
|
||||||
expects: TraceExpects::default(),
|
|
||||||
}],
|
|
||||||
memory_snapshot: Vec::new(),
|
|
||||||
http_exchanges: vec![HttpExchange {
|
|
||||||
request: HttpExchangeRequest {
|
|
||||||
method: "GET".to_string(),
|
|
||||||
url: expected_url.to_string(),
|
|
||||||
headers: vec![],
|
|
||||||
body: None,
|
|
||||||
},
|
|
||||||
response: github_ok(r#"{"number":42,"title":"Test","state":"open","body":"desc"}"#),
|
|
||||||
}],
|
|
||||||
expects: TraceExpects {
|
|
||||||
tools_used: vec!["github".to_string()],
|
|
||||||
all_tools_succeeded: Some(true),
|
|
||||||
max_tool_calls: Some(1),
|
|
||||||
min_responses: Some(1),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
steps: Vec::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
let rig = TestRigBuilder::new()
|
|
||||||
.with_trace(trace.clone())
|
|
||||||
.with_wasm_tool("github", GITHUB_WASM, Some(GITHUB_CAPS.into()))
|
|
||||||
.build()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
rig.send_message("Get issue 42 from nearai/ironclaw").await;
|
|
||||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
|
||||||
rig.verify_trace_expects(&trace, &responses);
|
|
||||||
|
|
||||||
rig.shutdown();
|
|
||||||
}
|
|
||||||
|
|
||||||
/// LLM sends `limit: "25"` (string) to `list_pull_requests`. URL must
|
|
||||||
/// contain `per_page=25`.
|
|
||||||
#[tokio::test]
|
|
||||||
#[ignore] // requires pre-compiled WASM binary
|
|
||||||
async fn wasm_github_list_prs_coerces_string_limit() {
|
|
||||||
let expected_url =
|
|
||||||
"https://api.github.com/repos/nearai/ironclaw/pulls?state=open&per_page=25";
|
|
||||||
|
|
||||||
let trace = LlmTrace {
|
|
||||||
model_name: "test-wasm-coercion-list-prs".to_string(),
|
|
||||||
turns: vec![crate::support::trace_llm::TraceTurn {
|
|
||||||
user_input: "List PRs in nearai/ironclaw".to_string(),
|
|
||||||
steps: vec![
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::ToolCalls {
|
|
||||||
tool_calls: vec![TraceToolCall {
|
|
||||||
id: "call_gh_3".to_string(),
|
|
||||||
name: "github".to_string(),
|
|
||||||
arguments: json!({
|
|
||||||
"action": "list_pull_requests",
|
|
||||||
"owner": "nearai",
|
|
||||||
"repo": "ironclaw",
|
|
||||||
"limit": "25"
|
|
||||||
}),
|
|
||||||
}],
|
|
||||||
input_tokens: 80,
|
|
||||||
output_tokens: 20,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
TraceStep {
|
|
||||||
request_hint: None,
|
|
||||||
response: TraceResponse::Text {
|
|
||||||
content: "Found PRs.".to_string(),
|
|
||||||
input_tokens: 100,
|
|
||||||
output_tokens: 10,
|
|
||||||
},
|
|
||||||
expected_tool_results: Vec::new(),
|
|
||||||
},
|
|
||||||
],
|
|
||||||
expects: TraceExpects::default(),
|
|
||||||
}],
|
|
||||||
memory_snapshot: Vec::new(),
|
|
||||||
http_exchanges: vec![HttpExchange {
|
|
||||||
request: HttpExchangeRequest {
|
|
||||||
method: "GET".to_string(),
|
|
||||||
url: expected_url.to_string(),
|
|
||||||
headers: vec![],
|
|
||||||
body: None,
|
|
||||||
},
|
|
||||||
response: github_ok(r#"[{"number":1,"title":"Test PR","state":"open"}]"#),
|
|
||||||
}],
|
|
||||||
expects: TraceExpects {
|
|
||||||
tools_used: vec!["github".to_string()],
|
|
||||||
all_tools_succeeded: Some(true),
|
|
||||||
max_tool_calls: Some(1),
|
|
||||||
min_responses: Some(1),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
steps: Vec::new(),
|
|
||||||
};
|
|
||||||
|
|
||||||
let rig = TestRigBuilder::new()
|
|
||||||
.with_trace(trace.clone())
|
|
||||||
.with_wasm_tool("github", GITHUB_WASM, Some(GITHUB_CAPS.into()))
|
|
||||||
.build()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
rig.send_message("List PRs in nearai/ironclaw").await;
|
|
||||||
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
|
|
||||||
rig.verify_trace_expects(&trace, &responses);
|
|
||||||
|
|
||||||
rig.shutdown();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,99 +0,0 @@
|
|||||||
use ironclaw::llm::ChatMessage;
|
|
||||||
use ironclaw::llm::gemini_oauth::GeminiOauthProvider;
|
|
||||||
|
|
||||||
/// Regression: Cloud Code API routing for Gemini 2.0+ models.
|
|
||||||
/// Gemini 1.x → legacy generativelanguage.googleapis.com
|
|
||||||
/// Gemini 2.0+ → Cloud Code API (cloudcode-pa.googleapis.com)
|
|
||||||
#[test]
|
|
||||||
fn test_regression_cloud_code_api_routing() {
|
|
||||||
// Legacy models (1.x) → false
|
|
||||||
assert!(!GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-1.5-pro"
|
|
||||||
));
|
|
||||||
assert!(!GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-1.5-flash"
|
|
||||||
));
|
|
||||||
|
|
||||||
// 2.0+ models → true
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-2.0-flash"
|
|
||||||
));
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-2.5-pro"
|
|
||||||
));
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-2.5-flash"
|
|
||||||
));
|
|
||||||
|
|
||||||
// Preview models with hyphen → true
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-3.1-pro-preview"
|
|
||||||
));
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-3-flash-preview"
|
|
||||||
));
|
|
||||||
|
|
||||||
// Gemini 3 family → true
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"gemini-3-pro"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Regression: "preview" false-positive fix.
|
|
||||||
/// `model.contains("-preview")` (with hyphen) prevents models whose name
|
|
||||||
/// happens to include "preview" without a hyphen prefix from being
|
|
||||||
/// mis-routed to Cloud Code API.
|
|
||||||
#[test]
|
|
||||||
fn test_regression_preview_false_positive_fix() {
|
|
||||||
// "my-preview-custom" still matches (contains "-preview")
|
|
||||||
assert!(GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"my-preview-custom"
|
|
||||||
));
|
|
||||||
|
|
||||||
// "mypreviewcustom" does NOT match (no hyphen before "preview")
|
|
||||||
assert!(!GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"mypreviewcustom"
|
|
||||||
));
|
|
||||||
|
|
||||||
// Non-Gemini models without "-preview" → false
|
|
||||||
assert!(!GeminiOauthProvider::model_uses_cloud_code_api(
|
|
||||||
"not-a-gemini-model"
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Regression: model list consistency.
|
|
||||||
/// Wizard, list_models(), and LLM_PROVIDERS.md all return the same 8 models.
|
|
||||||
#[test]
|
|
||||||
fn test_regression_standardized_model_list() {
|
|
||||||
let expected_models = [
|
|
||||||
"gemini-3.1-pro-preview",
|
|
||||||
"gemini-3.1-pro-preview-customtools",
|
|
||||||
"gemini-3-pro-preview",
|
|
||||||
"gemini-3-flash-preview",
|
|
||||||
"gemini-3.1-flash-lite-preview",
|
|
||||||
"gemini-2.5-pro",
|
|
||||||
"gemini-2.5-flash",
|
|
||||||
"gemini-2.5-flash-lite",
|
|
||||||
];
|
|
||||||
|
|
||||||
// All standardized models must route to Cloud Code API (all are >= 2.0)
|
|
||||||
for model in &expected_models {
|
|
||||||
assert!(
|
|
||||||
GeminiOauthProvider::model_uses_cloud_code_api(model),
|
|
||||||
"Standardized model '{}' should route to Cloud Code API",
|
|
||||||
model
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Regression: ChatMessage helper constructors.
|
|
||||||
#[test]
|
|
||||||
fn test_regression_chat_message_helpers() {
|
|
||||||
let user_msg = ChatMessage::user("hello");
|
|
||||||
assert_eq!(user_msg.role, ironclaw::llm::Role::User);
|
|
||||||
assert_eq!(user_msg.content, "hello");
|
|
||||||
|
|
||||||
let system_msg = ChatMessage::system("you are helpful");
|
|
||||||
assert_eq!(system_msg.role, ironclaw::llm::Role::System);
|
|
||||||
assert_eq!(system_msg.content, "you are helpful");
|
|
||||||
}
|
|
||||||
+15
-112
@@ -23,7 +23,7 @@ use crate::support::metrics::{ToolInvocation, TraceMetrics};
|
|||||||
use crate::support::test_channel::{TestChannel, TestChannelHandle};
|
use crate::support::test_channel::{TestChannel, TestChannelHandle};
|
||||||
use crate::support::trace_llm::{LlmTrace, TraceLlm};
|
use crate::support::trace_llm::{LlmTrace, TraceLlm};
|
||||||
|
|
||||||
use ironclaw::llm::recording::{HttpExchange, HttpInterceptor, ReplayingHttpInterceptor};
|
use ironclaw::llm::recording::{HttpExchange, ReplayingHttpInterceptor};
|
||||||
|
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
// TestRig
|
// TestRig
|
||||||
@@ -343,13 +343,6 @@ impl Drop for TestRig {
|
|||||||
// TestRigBuilder
|
// TestRigBuilder
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
/// Specification for loading a real WASM tool in the test rig.
|
|
||||||
pub struct WasmToolSpec {
|
|
||||||
pub name: String,
|
|
||||||
pub wasm_path: std::path::PathBuf,
|
|
||||||
pub capabilities_path: Option<std::path::PathBuf>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Builder for constructing a `TestRig`.
|
/// Builder for constructing a `TestRig`.
|
||||||
pub struct TestRigBuilder {
|
pub struct TestRigBuilder {
|
||||||
trace: Option<LlmTrace>,
|
trace: Option<LlmTrace>,
|
||||||
@@ -361,7 +354,6 @@ pub struct TestRigBuilder {
|
|||||||
enable_routines: bool,
|
enable_routines: bool,
|
||||||
http_exchanges: Vec<HttpExchange>,
|
http_exchanges: Vec<HttpExchange>,
|
||||||
extra_tools: Vec<Arc<dyn Tool>>,
|
extra_tools: Vec<Arc<dyn Tool>>,
|
||||||
wasm_tools: Vec<WasmToolSpec>,
|
|
||||||
keep_bootstrap: bool,
|
keep_bootstrap: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -378,34 +370,10 @@ impl TestRigBuilder {
|
|||||||
enable_routines: false,
|
enable_routines: false,
|
||||||
http_exchanges: Vec::new(),
|
http_exchanges: Vec::new(),
|
||||||
extra_tools: Vec::new(),
|
extra_tools: Vec::new(),
|
||||||
wasm_tools: Vec::new(),
|
|
||||||
keep_bootstrap: false,
|
keep_bootstrap: false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Load a real WASM tool binary into the test rig.
|
|
||||||
///
|
|
||||||
/// The tool will be compiled, registered, and wired with the same HTTP
|
|
||||||
/// interceptor used for `with_http_exchanges()`, so `http_exchanges` in
|
|
||||||
/// the trace can specify expected requests/responses for WASM tool HTTP calls.
|
|
||||||
///
|
|
||||||
/// If the WASM binary does not exist at build time, the tool is silently
|
|
||||||
/// skipped (logged as a warning). Tests should use `#[ignore]` or check
|
|
||||||
/// for the binary in a preamble if the tool is required.
|
|
||||||
pub fn with_wasm_tool(
|
|
||||||
mut self,
|
|
||||||
name: impl Into<String>,
|
|
||||||
wasm_path: impl Into<std::path::PathBuf>,
|
|
||||||
capabilities_path: Option<std::path::PathBuf>,
|
|
||||||
) -> Self {
|
|
||||||
self.wasm_tools.push(WasmToolSpec {
|
|
||||||
name: name.into(),
|
|
||||||
wasm_path: wasm_path.into(),
|
|
||||||
capabilities_path,
|
|
||||||
});
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Set the LLM trace to replay.
|
/// Set the LLM trace to replay.
|
||||||
pub fn with_trace(mut self, trace: LlmTrace) -> Self {
|
pub fn with_trace(mut self, trace: LlmTrace) -> Self {
|
||||||
self.trace = Some(trace);
|
self.trace = Some(trace);
|
||||||
@@ -497,7 +465,6 @@ impl TestRigBuilder {
|
|||||||
enable_routines,
|
enable_routines,
|
||||||
http_exchanges: explicit_http_exchanges,
|
http_exchanges: explicit_http_exchanges,
|
||||||
extra_tools,
|
extra_tools,
|
||||||
wasm_tools,
|
|
||||||
keep_bootstrap,
|
keep_bootstrap,
|
||||||
} = self;
|
} = self;
|
||||||
|
|
||||||
@@ -593,20 +560,6 @@ impl TestRigBuilder {
|
|||||||
let scheduler_slot: ironclaw::tools::builtin::SchedulerSlot =
|
let scheduler_slot: ironclaw::tools::builtin::SchedulerSlot =
|
||||||
Arc::new(tokio::sync::RwLock::new(None));
|
Arc::new(tokio::sync::RwLock::new(None));
|
||||||
|
|
||||||
// Build HTTP interceptor once — shared by both AgentDeps and WASM tools.
|
|
||||||
let http_interceptor: Option<Arc<dyn HttpInterceptor>> = {
|
|
||||||
let exchanges = if explicit_http_exchanges.is_empty() {
|
|
||||||
trace_http_exchanges
|
|
||||||
} else {
|
|
||||||
explicit_http_exchanges
|
|
||||||
};
|
|
||||||
if exchanges.is_empty() {
|
|
||||||
None
|
|
||||||
} else {
|
|
||||||
Some(Arc::new(ReplayingHttpInterceptor::new(exchanges)) as Arc<dyn HttpInterceptor>)
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
// 6. Register job tools, routine tools, and extra tools.
|
// 6. Register job tools, routine tools, and extra tools.
|
||||||
{
|
{
|
||||||
// Ensure filesystem/shell dev tools are always available in the
|
// Ensure filesystem/shell dev tools are always available in the
|
||||||
@@ -667,69 +620,6 @@ impl TestRigBuilder {
|
|||||||
for tool in extra_tools {
|
for tool in extra_tools {
|
||||||
components.tools.register(tool).await;
|
components.tools.register(tool).await;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Register WASM tools with the shared HTTP interceptor.
|
|
||||||
if !wasm_tools.is_empty() {
|
|
||||||
use ironclaw::tools::wasm::{
|
|
||||||
Capabilities, CapabilitiesFile, WasmRuntimeConfig, WasmToolRuntime,
|
|
||||||
WasmToolWrapper,
|
|
||||||
};
|
|
||||||
|
|
||||||
let runtime = Arc::new(
|
|
||||||
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
|
||||||
.expect("create WASM runtime for test rig"),
|
|
||||||
);
|
|
||||||
|
|
||||||
for spec in wasm_tools {
|
|
||||||
if !spec.wasm_path.exists() {
|
|
||||||
tracing::warn!(
|
|
||||||
name = %spec.name,
|
|
||||||
path = %spec.wasm_path.display(),
|
|
||||||
"WASM tool binary not found, skipping"
|
|
||||||
);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
let wasm_bytes = tokio::fs::read(&spec.wasm_path)
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| panic!("read {}: {e}", spec.wasm_path.display()));
|
|
||||||
let (capabilities, description, schema) =
|
|
||||||
if let Some(cap_path) = &spec.capabilities_path {
|
|
||||||
if cap_path.exists() {
|
|
||||||
let cap_bytes = tokio::fs::read(cap_path)
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| panic!("read {}: {e}", cap_path.display()));
|
|
||||||
let cap_file = CapabilitiesFile::from_bytes(&cap_bytes)
|
|
||||||
.expect("parse capabilities.json");
|
|
||||||
(
|
|
||||||
cap_file.to_capabilities(),
|
|
||||||
cap_file.description.clone(),
|
|
||||||
cap_file.parameters.clone(),
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
(Capabilities::default(), None, None)
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
(Capabilities::default(), None, None)
|
|
||||||
};
|
|
||||||
|
|
||||||
let prepared = runtime
|
|
||||||
.prepare(&spec.name, &wasm_bytes, None)
|
|
||||||
.await
|
|
||||||
.unwrap_or_else(|e| panic!("prepare WASM tool '{}': {e}", spec.name));
|
|
||||||
let mut wrapper =
|
|
||||||
WasmToolWrapper::new(Arc::clone(&runtime), prepared, capabilities);
|
|
||||||
if let Some(desc) = description {
|
|
||||||
wrapper = wrapper.with_description(desc);
|
|
||||||
}
|
|
||||||
if let Some(s) = schema {
|
|
||||||
wrapper = wrapper.with_schema(s);
|
|
||||||
}
|
|
||||||
if let Some(interceptor) = &http_interceptor {
|
|
||||||
wrapper = wrapper.with_http_interceptor(Arc::clone(interceptor));
|
|
||||||
}
|
|
||||||
components.tools.register(Arc::new(wrapper)).await;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Save references for test accessors.
|
// Save references for test accessors.
|
||||||
@@ -753,7 +643,20 @@ impl TestRigBuilder {
|
|||||||
hooks: components.hooks,
|
hooks: components.hooks,
|
||||||
cost_guard: components.cost_guard,
|
cost_guard: components.cost_guard,
|
||||||
sse_tx: None,
|
sse_tx: None,
|
||||||
http_interceptor,
|
http_interceptor: {
|
||||||
|
// Prefer explicit exchanges from with_http_exchanges(), fall back to trace.
|
||||||
|
let exchanges = if explicit_http_exchanges.is_empty() {
|
||||||
|
trace_http_exchanges
|
||||||
|
} else {
|
||||||
|
explicit_http_exchanges
|
||||||
|
};
|
||||||
|
if exchanges.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(Arc::new(ReplayingHttpInterceptor::new(exchanges))
|
||||||
|
as Arc<dyn ironclaw::llm::recording::HttpInterceptor>)
|
||||||
|
}
|
||||||
|
},
|
||||||
transcription: None,
|
transcription: None,
|
||||||
document_extraction: None,
|
document_extraction: None,
|
||||||
sandbox_readiness: ironclaw::agent::SandboxReadiness::Available, // tests don't use real Docker
|
sandbox_readiness: ironclaw::agent::SandboxReadiness::Available, // tests don't use real Docker
|
||||||
|
|||||||
@@ -65,6 +65,7 @@ async fn core_registration_covers_expected_tools() {
|
|||||||
"http",
|
"http",
|
||||||
"json",
|
"json",
|
||||||
"list_dir",
|
"list_dir",
|
||||||
|
"ptc_script",
|
||||||
"read_file",
|
"read_file",
|
||||||
"shell",
|
"shell",
|
||||||
"time",
|
"time",
|
||||||
|
|||||||
@@ -0,0 +1,20 @@
|
|||||||
|
[package]
|
||||||
|
name = "test-ptc-tool"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2021"
|
||||||
|
publish = false
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
wit-bindgen = "0.41.0"
|
||||||
|
serde_json = "1.0"
|
||||||
|
|
||||||
|
[lib]
|
||||||
|
crate-type = ["cdylib"]
|
||||||
|
|
||||||
|
[profile.release]
|
||||||
|
opt-level = "s"
|
||||||
|
lto = true
|
||||||
|
strip = true
|
||||||
|
codegen-units = 1
|
||||||
|
|
||||||
|
[workspace]
|
||||||
@@ -0,0 +1,52 @@
|
|||||||
|
wit_bindgen::generate!({
|
||||||
|
world: "sandboxed-tool",
|
||||||
|
path: "../../wit/tool.wit",
|
||||||
|
});
|
||||||
|
|
||||||
|
struct TestPtcTool;
|
||||||
|
|
||||||
|
impl exports::near::agent::tool::Guest for TestPtcTool {
|
||||||
|
fn execute(req: exports::near::agent::tool::Request) -> exports::near::agent::tool::Response {
|
||||||
|
match execute_inner(&req.params) {
|
||||||
|
Ok(result) => exports::near::agent::tool::Response {
|
||||||
|
output: Some(result),
|
||||||
|
error: None,
|
||||||
|
},
|
||||||
|
Err(e) => exports::near::agent::tool::Response {
|
||||||
|
output: None,
|
||||||
|
error: Some(e),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn schema() -> String {
|
||||||
|
r#"{"type":"object","properties":{"message":{"type":"string"}},"required":["message"]}"#.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description() -> String {
|
||||||
|
"Test tool for PTC: calls echo via tool_invoke".to_string()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn execute_inner(params: &str) -> Result<String, String> {
|
||||||
|
let parsed: serde_json::Value = serde_json::from_str(params)
|
||||||
|
.map_err(|e| format!("Invalid params: {}", e))?;
|
||||||
|
|
||||||
|
let message = parsed.get("message")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.ok_or("Missing 'message' parameter")?;
|
||||||
|
|
||||||
|
// Build the parameters for the echo tool
|
||||||
|
let echo_params = serde_json::json!({"message": message});
|
||||||
|
|
||||||
|
// Call tool_invoke with alias "echo_alias" which should resolve to "echo"
|
||||||
|
let result = near::agent::host::tool_invoke(
|
||||||
|
"echo_alias",
|
||||||
|
&echo_params.to_string(),
|
||||||
|
)?;
|
||||||
|
|
||||||
|
// Prefix to prove it went through WASM
|
||||||
|
Ok(format!("via_wasm:{}", result))
|
||||||
|
}
|
||||||
|
|
||||||
|
export!(TestPtcTool);
|
||||||
Reference in New Issue
Block a user