Compare commits

..
Author SHA1 Message Date
[email protected]andClaude Opus 4.6 86853244ca docs: add Gemini OAuth env vars to .env.example [skip-regression-check]
Document GEMINI_MODEL, GEMINI_CREDENTIALS_PATH, GEMINI_API_KEY, and
all extended generation config env vars in the example config file.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 21:59:21 -07:00
[email protected]andClaude Opus 4.6 dd9057d069 fix(llm): support smart routing cheap model for gemini_oauth backend
Add explicit gemini_oauth handling in create_cheap_provider_for_backend()
to create a GeminiOauthProvider with the cheap model swapped in. Without
this, setting LLM_CHEAP_MODEL with gemini_oauth backend would fail with
a confusing "no registry provider config available" error.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 20:58:57 -07:00
[email protected]andClaude Opus 4.6 3a72b71d97 style(gemini_oauth): rustfmt formatting [skip-regression-check]
Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 15:37:42 -07:00
[email protected]andClaude Opus 4.6 94f88231a1 fix(gemini_oauth): curate_contents per-part filtering and dead code removal
Fix curate_contents to filter invalid parts individually instead of
dropping entire model turn sequences. Previously a single empty text
part would discard all consecutive model turns including valid
functionCall parts, breaking the tool-call flow.

Also remove unused MID_STREAM_* constants.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 15:37:10 -07:00
[email protected]andClaude Opus 4.6 99566bb997 fix(gemini_oauth): align header parser doc with implementation [skip-regression-check]
Update parse_custom_headers doc comments to include underscore in the
header-name character class, matching the actual implementation.
Also fix formatting from merge.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 15:26:50 -07:00
[email protected]andClaude Opus 4.6 cbf96f25c9 merge: resolve origin/staging into feat/gemini-cli-oauth
Merge staging to pick up GitHub Copilot provider, OpenAI Codex provider,
and other recent changes. Both gemini_oauth and openai_codex backends are
now registered as dedicated configs with proper credential guards.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 15:18:13 -07:00
b97d82dbe6 feat(extensions): support text setup fields in web configure modal (#496)
* feat(extensions): support text setup fields in web configure modal

* fix(extensions): use exported wasm setup schema types

* fix(extensions): validate extension name in setup APIs

* fix(extensions): restrict setup setting_path writes

* refactor(web): use enum for setup field input type

* fix: restore registry versions reverted during merge [skip-regression-check]

The merge auto-resolved registry JSON conflicts in favor of the PR's
older 0.2.0 versions. Restore discord, github, and web-search to
0.2.1 from staging.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: 您的GitHub用户名 <[email protected]>
Co-authored-by: [email protected] <[email protected]>
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 15:10:09 -07:00
9d538136b5 fix(oauth): reject malformed ic2.* states in decode_hosted_oauth_state (#1441) (#1454)
* fix(oauth): reject malformed ic2.* states instead of falling through to legacy handler (#1441)

When decode_hosted_oauth_state() encountered a versioned state (ic2.*)
that failed to fully parse (bad base64, invalid JSON, missing separator),
it silently fell through to legacy handling which used the full malformed
envelope as the flow_id. This never matched the raw nonce stored in
pending_oauth_flows, breaking the OAuth callback.

Restructure the versioned decode path so any ic2.* state must parse as a
valid envelope or return Err — never fall through to legacy handling.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(oauth): address PR review — avoid alloc in strip_prefix, strengthen JSON parse test

- Replace `strip_prefix(&format!(...))` with a `HOSTED_STATE_PREFIX_DOT`
  constant to avoid per-call allocation.
- Fix "valid base64 but not JSON" test to compute the correct checksum so
  it actually exercises the JSON parse error path instead of stopping at
  the checksum check.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: add missing fallback_deliverable field in job_monitor tests

The SseEvent::JobResult struct gained a fallback_deliverable field in
the structured fallback deliverables feature, but the job_monitor test
constructors were not updated.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix(oauth): remove HOSTED_STATE_PREFIX_DOT to avoid drift with HOSTED_STATE_PREFIX

concat! requires literals and cannot reference const items, so a
separate _DOT constant would duplicate the prefix string. Revert to
deriving the dotted prefix via format!() — both encode and decode now
use the same single HOSTED_STATE_PREFIX constant, keeping them
mechanically consistent.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-21 14:39:52 -07:00
8ad7d78a70 fix: parameter coercion and validation for oneOf/anyOf/allOf schemas (#1397)
* fix: parameter coercion and validation for oneOf/anyOf/allOf schemas

WASM extension tools with multi-action schemas (e.g. github extension)
fail when the LLM passes numeric parameters as strings because the
coercion layer skips JSON Schema combinators. This causes serde
deserialization errors like `invalid type: string "100", expected u32`.

Add discriminated-union resolution to the coercion layer: for oneOf/anyOf,
match the active variant by const or single-element enum discriminators;
for allOf, merge all variants' properties. Also propagate combinator
awareness to schema validators, WASM wrapper helpers, and tool discovery
so they no longer reject or ignore valid combinator-based schemas.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* test: add e2e tests for oneOf discriminated union parameter coercion

Add three end-to-end tests using a fixture tool that mirrors the github
WASM tool's oneOf schema with #[serde(tag = "action")] deserialization.
Each test sends string-typed numeric/boolean params through the full
agent loop, verifying that coercion resolves them before serde runs:

- list_issues: limit "100" → 100 (integer in oneOf variant)
- get_issue: issue_number "42" → 42 (integer in different variant)
- create_pull_request: draft "true" → true (boolean in variant)

Without the coercion fix these fail with:
  invalid type: string "100", expected u32

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* test: add real WASM github tool e2e tests with HTTP interception

Load the actual compiled github WASM binary, send params with string-typed
numbers through the coercion layer, and verify the WASM tool constructs
correct HTTP API calls via a new HTTP interceptor in the WASM wrapper.

Changes:
- Add `http_interceptor` field to `StoreData` and `WasmToolWrapper` so
  WASM tool HTTP requests can be captured/mocked in tests
- Make `prepare_tool_params` and `coercion` module public for integration tests
- Add 3 e2e tests loading the real github WASM binary:
  - list_issues: `limit: "50"` → URL contains `per_page=50`
  - get_issue: `issue_number: "42"` → URL contains `/issues/42`
  - list_pull_requests: `limit: "25"` → URL contains `per_page=25`

Tests gracefully skip if the WASM binary isn't compiled.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* refactor: simplify WASM e2e tests to use TestRig with with_wasm_tool()

Replace the manual WasmToolWrapper construction with TestRig integration:

- Add `with_wasm_tool(name, wasm_path, capabilities_path)` to TestRigBuilder
  that loads real WASM binaries and wires the shared HTTP interceptor
- Build the HTTP interceptor before tool registration so it can be shared
  between AgentDeps and WASM tool wrappers
- Rewrite github WASM e2e tests to use the standard trace pattern:
  TraceLlm sends tool calls with string params, http_exchanges specify
  expected outgoing requests and canned responses

The test code is now identical to other trace-based e2e tests — no custom
interceptors or manual WASM construction needed.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address review comments on combinator schema support

- Validate `has_combinators` checks array type (`.as_array().is_some()`)
  instead of bare `.is_some()` to reject malformed `{ "oneOf": {} }`
- Validate top-level `required` keys against merged combinator variant
  properties when no top-level `properties` exists (both validators)
- Deduplicate oneOf/anyOf handling into single loop in coercion.rs
- Revert `pub mod coercion` to private; only re-export `prepare_tool_params`
- Call `after_response` on interceptor after real HTTP when `before_request`
  returns None (recording mode correctness)
- Fix formatting (CI failure)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address second round of review comments

- Fix headers deserialization bug: deserialize resp.headers_json as
  HashMap<String, String> then convert to Vec, not directly as Vec
- Sort interceptor headers for deterministic trace fixtures
- Update after_response comment: RecordingHttpInterceptor does exercise
  this path (returns None from before_request)
- Mark WASM tests #[ignore] instead of silent skip — avoids false-green
  CI while keeping them runnable with --ignored
- Fix with_wasm_tool signature: Option<PathBuf> instead of
  Option<impl Into<PathBuf>> which doesn't compile in nested position
- Fix with_wasm_tool doc comment to match actual behavior
- Revert prepare_tool_params to pub(crate) — no longer needed publicly

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: coerce empty strings to null for optional tool parameters

LLMs often send "" instead of null/omitting optional parameters, causing
parse errors in tools that expect typed values (e.g., timezone, schedule).

PR #1127 fixed this per-field in the time tool. This commit adds
dispatcher-level coercion so all tools benefit:

- Non-required properties with value "" are coerced to null at the
  object level (based on the schema's `required` array)
- Explicitly nullable schemas (`type: ["string", "null"]`) coerce ""
  to null in the per-value coercion path
- Required string-only fields keep "" unchanged

Closes #755

Co-Authored-By: spiritj <[email protected]>
Co-Authored-By: Xing Ji <[email protected]>
Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* feat: complete coercion coverage for $ref, nested combinators, and additionalProperties

Close remaining coercion gaps so 3rd-party tools (MCP servers, complex
WASM tools) work correctly:

- $ref resolution: inline all #/definitions/<name> and #/$defs/<name>
  references in a pre-pass before coercion, with depth limit (16) for
  circular ref safety
- Nested combinators: resolve_effective_properties now recurses into
  variants that themselves contain allOf/oneOf/anyOf (depth limit 4)
- additionalProperties inheritance: check allOf variants and matched
  oneOf/anyOf variant for additionalProperties schemas

New tests:
- resolves_ref_and_coerces_referenced_properties
- resolves_nested_refs_in_oneof_variants
- coerces_nested_combinators_allof_containing_oneof
- coerces_array_items_with_oneof_discriminator
- circular_ref_does_not_infinite_loop

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address third round of review comments

- Validators: tighten has_combinators to require at least one object-typed
  variant (has type:"object" or properties), rejecting non-object combinator
  schemas like { "oneOf": [{"type":"integer"}] }
- Empty-string coercion: only coerce "" → null when schema allows null or
  doesn't allow string; pure type:"string" fields keep "" as meaningful
- Fix comment: "coerce to null" → "return unchanged" for empty strings
  with no type match (code returns None, not null)
- Redact credentials before passing to after_response interceptor to
  prevent secret leakage into recorded trace files
- Switch to tokio::fs::read for async WASM binary loading in test rig
- Add doc comment explaining soft URL check in WASM e2e tests

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* ci: retrigger after staging merge [skip-regression-check]

* fix: merge staging, report non-array combinator values as errors

Merge latest staging to fix CI (missing fallback_deliverable field).
Add explicit error reporting when oneOf/anyOf/allOf values are not
arrays in both strict and lenient validators.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: recurse into combinator variants that have properties but no explicit type

Both validators only recursed into variants with `type: "object"`,
missing variants that define `properties` without an explicit type
(common in allOf patterns). Now recurse when variant has either.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: spiritj <[email protected]>
Co-authored-by: Xing Ji <[email protected]>
2026-03-21 12:41:46 -07:00
Artem a070c069cb fix: add missing allow_always field after staging merge 2026-03-19 22:05:56 +03:00
Artem abf13a3ee2 Merge remote-tracking branch 'origin/staging' into feat/gemini-cli-oauth 2026-03-19 21:57:29 +03:00
Artem 21abbe5691 fix: address Copilot PR review feedback
- Fix empty text part for assistant messages with tool calls
  (curate_contents could drop entire model turn)
- Propagate cache_read/creation_input_tokens in complete_with_tools
- Log warning on save_credential failure instead of silently ignoring
- Fix doc comment to mention underscore in header name pattern
- Handle gemini-oauth (hyphen variant) in setup wizard display
- Fix docs: thinkingConfig uses thinkingBudget/thinkingLevel, not
  includeThoughts
2026-03-19 07:26:55 +03:00
Artem 9be29b2c22 fix: address PR review feedback from gemini-code-assist
- Fix parse_custom_headers to preserve commas in values by splitting
  only on commas followed by a header-name:colon pattern (manual scan
  instead of simple split(','))
- Use matches! macro for backend exclusion check in app.rs
- Merge SSE metadata extraction into single pass (was iterating twice)
- Replace fragile substring-based context_length with explicit match
  on known Gemini model IDs via gemini_context_length()
- Add missing models to regression test (8 models, not 5)
2026-03-18 14:27:12 +03:00
Artem 09a2320e64 fix: CI violations — add safety comment on expect, fix fmt
- Add '// safety: hardcoded literal' to regex .expect() to satisfy
  the no-panic-in-prod CI check
- Fix cargo fmt whitespace in collapsible if-let chain
2026-03-18 13:57:19 +03:00
Artem d4bfc2db58 feat(gemini_oauth): full Cloud Code API integration with project discovery
- Register gemini_oauth as a dedicated backend in config/llm.rs (skip
  registry fallback, preserve backend name, suppress unknown-backend warning)
- Fix app.rs credential guard to exclude backends with dedicated configs
  (gemini_oauth, bedrock) from the provider.is_none() check
- Auto-discover Cloud Code project_id via loadCodeAssist when credentials
  lack it (e.g. created by the original Gemini CLI)
- Persist discovered project_id to credentials file for subsequent runs
- Add safety settings (BLOCK_NONE), gated behind GEMINI_SAFETY_BLOCK_NONE env
- Add thinkingConfig: budget-based for Gemini 2.5, level-based for Gemini 3.x
  (without includeThoughts to avoid empty responses from reasoning.rs stripping)
- Add thought signature injection for Gemini 3.x preview APIs
- Add history curation to filter invalid model outputs before re-sending
- Add extended generationConfig env vars (topP, topK, seed, penalties,
  responseMimeType, responseJsonSchema, cachedContent)
- Add custom headers support via GEMINI_CLI_CUSTOM_HEADERS
- Add API key auth mode (GEMINI_API_KEY + GEMINI_API_KEY_AUTH_MECHANISM)
- Add SSE metadata extraction (modelVersion, credits, promptFeedback,
  groundingMetadata, citationMetadata, cachedContentTokenCount)
- Add countTokens API support
- Add new models to wizard (gemini-3.1-pro-preview-customtools,
  gemini-3-pro-preview, gemini-3.1-flash-lite-preview)
- Update docs/LLM_PROVIDERS.md with new models and routing rules
- Rewrite regression tests with comprehensive coverage (23 unit tests pass)
2026-03-18 13:53:58 +03:00
Artem 6e86f50bc5 Merge remote-tracking branch 'origin/main' into feat/gemini-cli-oauth 2026-03-17 18:21:02 +03:00
Artem 12b0e90a7a feat(gemini-oauth): implement code review v3 refinements
- Add force_refresh() for 401 retry (bypass timestamp check)
- Standardize Gemini model list across docs, wizard, and provider
- Restore gemini-3 check for thinkingConfig
- Redact sensitive tokens in GoogleTokenRefreshResponse Debug output
- Use dynamic version for GOOG_API_CLIENT
- Improve model_metadata() context length heuristics
- Use strip_prefix("data:") for safer SSE parsing
- Skip re-auth in wizard if keeping existing provider
2026-03-10 13:15:05 +03:00
Artem 4a2950e777 style: fix formatting in Gemini OAuth regression tests 2026-03-09 16:47:09 +03:00
Artem b35771d505 Add dedicated regression tests for Gemini OAuth fixes 2026-03-09 16:44:30 +03:00
Artem e4e747ba54 fix: address code review issues in gemini-cli OAuth integration
- Add cache_read_input_tokens/cache_creation_input_tokens fields (value 0)
- Implement manual Debug for OAuthCredential to redact tokens
- Fix hardcoded /tmp: use GeminiOauthConfig::default_credentials_path()
- Replace emoji output with plain text markers
- Propagate Client::builder() errors instead of silent fallback
- Use tokio::fs for all file I/O in CredentialManager (was std::fs)
- Use if let Some(ref pid) to avoid consuming credential.project_id
- Extract uses_cloud_code_api() helper; route by major version (gemini-2+)
- Concatenate multiple system messages into systemInstruction
- Include functionCall parts in assistant message conversion
- Add 401 retry loop with allow_retry flag for auth failures
- Remove biased from tokio::select! in OAuth callback handler
- Remove hardcoded context_length 1M; vary by model family
- Change GOOG_API_CLIENT from Node.js spoof to gl-rust/1.0.0
- Implement list_models() with static model list
- Move create_gemini_oauth_provider() before test module (clippy)
- Fix 9 additional clippy warnings (collapsible_if, map_or, needless_borrow)
- Run cargo fmt
2026-03-09 16:31:15 +03:00
Artem fdb0077736 Merge origin/main into feat/gemini-cli-oauth and resolve conflicts 2026-03-09 15:20:57 +03:00
Artem 8452102454 feat(gemini): implement function calling, generationConfig, and update models
- Implement function calling support (functionDeclarations, functionResponse)
- Add functionCall SSE parsing and empty stream retry support
- Add generationConfig (temperature, maxOutputTokens)
- Add thinkingConfig for Gemini 3 and thinking models
- Add toolConfig (functionCallingConfig.mode)
- Fix .expect() panics with .ok_or_else()
- Restrict oauth credentials file permissions to 0600
- Update docs and FEATURE_PARITY.md
- Update wizard to current Gemini 3.1 and 2.5 models
2026-03-04 23:05:16 +03:00
Artem 727283afe3 feat: integrate Gemini CLI OAuth with Cloud Code API
- Add gemini_oauth.rs: full OAuth flow with PKCE, token refresh,
  and Cloud Code project discovery (loadCodeAssist + onboardUser)
- Route preview/gemini-3 models through cloudcode-pa.googleapis.com
  with proper project ID injection in request payload
- Trigger OAuth login during onboarding wizard (not first chat message)
- Support manual redirect URL paste as fallback (tokio::select race)
- Parse 429 rate-limit errors with retry_after from Google response
- Add static model list: gemini-1.5/2.0/2.5/3.0/3.1 variants
- Add GeminiOauthConfig with default credentials path (~/.gemini/)
2026-03-02 23:50:45 +03:00
47 changed files with 5781 additions and 2591 deletions
+18 -1
View File
@@ -4,7 +4,7 @@ DATABASE_POOL_SIZE=10
# LLM Provider
# LLM_BACKEND=nearai # default
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, github_copilot, tinfoil, openai_codex
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, github_copilot, tinfoil, openai_codex, gemini_oauth
# LLM_REQUEST_TIMEOUT_SECS=120 # Increase for local LLMs (Ollama, vLLM, LM Studio)
# === Anthropic Direct ===
@@ -110,6 +110,23 @@ NEARAI_AUTH_URL=https://private.near.ai
# OPENAI_CODEX_AUTH_URL=https://auth.openai.com # 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
# Channel Configuration
-3
View File
@@ -55,9 +55,6 @@ RUN npm install -g @anthropic-ai/claude-code@latest
# Copy the binary
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)
RUN useradd -m -u 1000 -s /bin/bash sandbox \
&& mkdir -p /workspace \
+14 -5
View File
@@ -3,6 +3,7 @@
This document tracks feature parity between IronClaw (Rust implementation) and OpenClaw (TypeScript reference implementation). Use this to coordinate work across developers.
**Legend:**
- ✅ Implemented
- 🚧 Partial (in progress or incomplete)
- ❌ Not implemented
@@ -204,7 +205,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 |
| Skill routing blocks | ✅ | 🚧 | ActivationCriteria (keywords, patterns, tags) but no "Use when / Don't use when" blocks |
| Skill path compaction | ✅ | ❌ | ~ prefix to reduce prompt tokens |
| Thinking modes (off/minimal/low/medium/high/xhigh/adaptive) | ✅ | | Configurable reasoning depth |
| Thinking modes (off/minimal/low/medium/high/xhigh/adaptive) | ✅ | 🚧 | thinkingConfig for Gemini models (thinkingBudget/thinkingLevel); no per-level control yet |
| Per-model thinkingDefault override | ✅ | ❌ | Override thinking level per model; Anthropic Claude 4.6 defaults to adaptive |
| Block-level streaming | ✅ | ❌ | |
| Tool-level streaming | ✅ | ❌ | |
@@ -236,9 +237,13 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| NEAR AI | ✅ | ✅ | - | Primary provider |
| 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 |
| AWS Bedrock | ✅ | ❌ | P3 | |
| Google Gemini | ✅ | | P3 | |
| NVIDIA API | ✅ | | P3 | New provider |
| AWS Bedrock | ✅ | ✅ | - | Native Converse API via aws-sdk-bedrockruntime (requires `--features bedrock`) |
| Google Gemini | ✅ | | - | OAuth (PKCE + S256), function calling, thinkingConfig, generationConfig |
| io.net | ✅ | | P3 | Via `ionet` adapter |
| 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) |
| Tinfoil | ❌ | ✅ | - | Private inference provider (IronClaw-only) |
| OpenAI-compatible | ❌ | ✅ | - | Generic OpenAI-compatible endpoint (RigAdapter) |
@@ -466,7 +471,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Device pairing | ✅ | ❌ | |
| Tailscale identity | ✅ | ❌ | |
| Trusted-proxy auth | ✅ | ❌ | Header-based reverse proxy auth |
| OAuth flows | ✅ | 🚧 | NEAR AI OAuth plus hosted extension/MCP OAuth broker; external auth-proxy rollout still pending |
| OAuth flows | ✅ | 🚧 | NEAR AI OAuth + Gemini OAuth (PKCE, S256) + hosted extension/MCP OAuth broker; external auth-proxy rollout still pending |
| DM pairing verification | ✅ | ✅ | ironclaw pairing approve, host APIs |
| Allowlist/blocklist | ✅ | 🚧 | allow_from + pairing store |
| Per-group tool policies | ✅ | ❌ | |
@@ -523,6 +528,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
## Implementation Priorities
### P0 - Core (Already Done)
- ✅ TUI channel with approval overlays
- ✅ HTTP webhook channel
- ✅ DM pairing (ironclaw pairing list/approve, host APIs)
@@ -550,6 +556,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
- ✅ OpenAI-compatible / OpenRouter provider support
### P1 - High Priority
- ❌ Slack channel (real implementation)
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
- ❌ WhatsApp channel
@@ -557,6 +564,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
- ✅ Hooks system (core lifecycle hooks + bundled/plugin/workspace hooks + outbound webhooks)
### P2 - Medium Priority
- ❌ Media handling (images, PDFs)
- ✅ Ollama/local model support (via rig::providers::ollama)
- ❌ Configuration hot-reload
@@ -565,6 +573,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
- ❌ Partial output preservation on abort
### P3 - Lower Priority
- ❌ Discord channel
- ❌ Matrix channel
- ❌ Other messaging platforms
+48 -3
View File
@@ -1,8 +1,8 @@
# LLM Provider Configuration
IronClaw defaults to NEAR AI for model access, but supports any OpenAI-compatible
endpoint as well as Anthropic and Ollama directly. This guide covers the most common
configurations.
endpoint as well as Anthropic, Ollama, and Google Gemini directly. This guide covers
the most common configurations.
## Provider Overview
@@ -11,7 +11,7 @@ configurations.
| NEAR AI | `nearai` | OAuth (browser) | Default; multi-model |
| Anthropic | `anthropic` | `ANTHROPIC_API_KEY` | Claude models |
| OpenAI | `openai` | `OPENAI_API_KEY` | GPT models |
| Google Gemini | `gemini` | `GEMINI_API_KEY` | Gemini models |
| Google Gemini | `gemini_oauth` | OAuth (browser) | Gemini models; function calling |
| io.net | `ionet` | `IONET_API_KEY` | Intelligence API |
| Mistral | `mistral` | `MISTRAL_API_KEY` | Mistral models |
| Yandex AI Studio | `yandex` | `YANDEX_API_KEY` | YandexGPT models |
@@ -62,6 +62,51 @@ 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 exposes chat endpoint at
-158
View File
@@ -1,158 +0,0 @@
"""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)
-148
View File
@@ -1,148 +0,0 @@
"""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
View File
@@ -729,13 +729,13 @@ impl AppBuilder {
self.init_database().await?;
self.init_secrets().await?;
// Post-init validation: if a non-nearai backend was selected but
// credentials were never resolved (deferred resolution found no keys),
// fail early with a clear error instead of a confusing runtime failure.
if self.config.llm.backend != "nearai"
&& self.config.llm.backend != "bedrock"
&& self.config.llm.backend != "openai_codex"
&& self.config.llm.provider.is_none()
// Post-init validation: backends with dedicated config (nearai, gemini_oauth,
// bedrock, openai_codex) handle their own credential resolution. For registry-based
// backends, fail early if no provider config was resolved.
if !matches!(
self.config.llm.backend.as_str(),
"nearai" | "gemini_oauth" | "bedrock" | "openai_codex"
) && self.config.llm.provider.is_none()
{
let backend = &self.config.llm.backend;
anyhow::bail!(
+7 -3
View File
@@ -2343,7 +2343,7 @@ async fn extensions_setup_handler(
"Extension manager not available (secrets store required)".to_string(),
))?;
let secrets = ext_mgr
let setup = ext_mgr
.get_setup_schema(&name)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
@@ -2359,7 +2359,8 @@ async fn extensions_setup_handler(
Ok(Json(ExtensionSetupResponse {
name,
kind,
secrets,
secrets: setup.secrets,
fields: setup.fields,
}))
}
@@ -2377,7 +2378,7 @@ async fn extensions_setup_submit_handler(
// through to the LLM instead of being intercepted as a token.
clear_auth_mode(&state).await;
match ext_mgr.configure(&name, &req.secrets).await {
match ext_mgr.configure(&name, &req.secrets, &req.fields).await {
Ok(result) => {
let mut resp = if result.verification.is_some() || result.activated {
ActionResponse::ok(result.message)
@@ -2385,6 +2386,9 @@ async fn extensions_setup_submit_handler(
ActionResponse::fail(result.message)
};
resp.activated = Some(result.activated);
if result.restart_required || !result.activated {
resp.needs_restart = Some(true);
}
resp.auth_url = result.auth_url.clone();
resp.verification = result.verification.clone();
resp.instructions = result.verification.as_ref().map(|v| v.instructions.clone());
+58 -8
View File
@@ -2791,16 +2791,18 @@ function removeExtension(name) {
function showConfigureModal(name) {
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup')
.then((setup) => {
if (!setup.secrets || setup.secrets.length === 0) {
const secrets = Array.isArray(setup.secrets) ? setup.secrets : [];
const setupFields = Array.isArray(setup.fields) ? setup.fields : [];
if (secrets.length === 0 && setupFields.length === 0) {
showToast('No configuration needed for ' + name, 'info');
return;
}
renderConfigureModal(name, setup.secrets);
renderConfigureModal(name, secrets, setupFields);
})
.catch((err) => showToast('Failed to load setup: ' + err.message, 'error'));
}
function renderConfigureModal(name, secrets) {
function renderConfigureModal(name, secrets, setupFields) {
closeConfigureModal();
const overlay = document.createElement('div');
overlay.className = 'configure-overlay';
@@ -2873,7 +2875,46 @@ function renderConfigureModal(name, secrets) {
field.appendChild(inputRow);
form.appendChild(field);
fields.push({ name: secret.name, input: input });
fields.push({ kind: 'secret', 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);
@@ -3015,9 +3056,16 @@ function startTelegramAutoVerify(name, fields) {
function submitConfigureModal(name, fields, options) {
options = options || {};
const secrets = {};
const setupFields = {};
for (const f of fields) {
if (f.input.value.trim()) {
secrets[f.name] = f.input.value.trim();
const value = f.input.value.trim();
if (!value) {
continue;
}
if (f.kind === 'secret') {
secrets[f.name] = value;
} else {
setupFields[f.name] = value;
}
}
@@ -3034,7 +3082,7 @@ function submitConfigureModal(name, fields, options) {
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup', {
method: 'POST',
body: { secrets },
body: { secrets, fields: setupFields },
})
.then((res) => {
if (res.success) {
@@ -3064,6 +3112,8 @@ function submitConfigureModal(name, fields, options) {
showToast('Opening OAuth authorization for ' + name, 'info');
openOAuthUrl(res.auth_url);
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,
// which will show the toast and refresh extensions — no need to do it here too.
@@ -4012,7 +4062,7 @@ function formatRelativeTime(isoString) {
const absDiff = Math.abs(diffMs);
const future = diffMs < 0;
if (absDiff < 60000)
if (absDiff < 60000)
return future ? I18n.t('time.lessThan1MinuteFromNow') : I18n.t('time.lessThan1MinuteAgo');
if (absDiff < 3600000) {
const m = Math.floor(absDiff / 60000);
+54
View File
@@ -525,6 +525,7 @@ pub struct ExtensionSetupResponse {
pub name: String,
pub kind: String,
pub secrets: Vec<SecretFieldInfo>,
pub fields: Vec<SetupFieldInfo>,
}
#[derive(Debug, Serialize)]
@@ -538,9 +539,23 @@ pub struct SecretFieldInfo {
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)]
pub struct ExtensionSetupRequest {
#[serde(default)]
pub secrets: std::collections::HashMap<String, String>,
#[serde(default)]
pub fields: std::collections::HashMap<String, String>,
}
#[derive(Debug, Serialize)]
@@ -559,6 +574,9 @@ pub struct ActionResponse {
/// Whether the channel was successfully activated after setup.
#[serde(skip_serializing_if = "Option::is_none")]
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.).
#[serde(skip_serializing_if = "Option::is_none")]
pub verification: Option<crate::extensions::VerificationChallenge>,
@@ -573,6 +591,7 @@ impl ActionResponse {
awaiting_token: None,
instructions: None,
activated: None,
needs_restart: None,
verification: None,
}
}
@@ -585,6 +604,7 @@ impl ActionResponse {
awaiting_token: None,
instructions: None,
activated: None,
needs_restart: None,
verification: None,
}
}
@@ -1246,6 +1266,40 @@ mod tests {
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 ----
#[test]
+83 -18
View File
@@ -579,23 +579,27 @@ 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
/// legacy `instance:nonce`/`nonce` forms.
pub fn decode_hosted_oauth_state(state: &str) -> Result<DecodedHostedOAuthState, String> {
if let Some(rest) = state.strip_prefix(&format!("{HOSTED_STATE_PREFIX}."))
&& let Some((payload_b64, checksum)) = rest.rsplit_once('.')
&& let Ok(payload_json) = URL_SAFE_NO_PAD.decode(payload_b64)
{
if let Some(rest) = state.strip_prefix(&format!("{HOSTED_STATE_PREFIX}.")) {
let (payload_b64, checksum) = rest
.rsplit_once('.')
.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);
if checksum != expected_checksum {
return Err("Hosted OAuth state checksum mismatch".to_string());
}
if let Ok(payload) = serde_json::from_slice::<HostedOAuthStatePayload>(&payload_json)
&& !payload.flow_id.trim().is_empty()
{
return Ok(DecodedHostedOAuthState {
flow_id: payload.flow_id,
instance_name: payload.instance_name.filter(|v| !v.is_empty()),
is_legacy: false,
});
let payload: HostedOAuthStatePayload = serde_json::from_slice(&payload_json)
.map_err(|e| format!("Hosted OAuth versioned state JSON parse failed: {e}"))?;
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,
});
}
if let Some((instance_name, flow_id)) = state.split_once(':') {
@@ -1187,14 +1191,14 @@ mod tests {
}
#[test]
fn test_decode_hosted_oauth_state_falls_back_for_non_envelope_ic2_prefix() {
fn test_decode_hosted_oauth_state_rejects_non_envelope_ic2_prefix() {
use crate::cli::oauth_defaults::decode_hosted_oauth_state;
let decoded =
decode_hosted_oauth_state("ic2.provider-owned-state").expect("prefixed fallback");
assert_eq!(decoded.flow_id, "ic2.provider-owned-state");
assert_eq!(decoded.instance_name, None);
assert!(decoded.is_legacy);
// "ic2." prefix must parse as a valid versioned envelope — never fall
// through to legacy handling, which would use the full malformed
// envelope as the flow_id and break OAuth callback lookup (#1441).
decode_hosted_oauth_state("ic2.provider-owned-state")
.expect_err("ic2-prefixed non-envelope state should fail");
}
#[test]
@@ -1244,4 +1248,65 @@ mod tests {
assert!(result.url.contains("code_challenge="));
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);
}
}
+26 -3
View File
@@ -9,6 +9,7 @@ use crate::llm::config::*;
use crate::llm::registry::{ProviderProtocol, ProviderRegistry};
use crate::llm::session::SessionConfig;
use crate::settings::Settings;
impl LlmConfig {
/// Create a test-friendly config without reading env vars.
#[cfg(feature = "libsql")]
@@ -37,6 +38,7 @@ impl LlmConfig {
},
provider: None,
bedrock: None,
gemini_oauth: None,
openai_codex: None,
request_timeout_secs: 120,
cheap_model: None,
@@ -73,11 +75,16 @@ impl LlmConfig {
backend_lower == "nearai" || backend_lower == "near_ai" || backend_lower == "near";
let is_bedrock =
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"
|| backend_lower == "openai-codex"
|| backend_lower == "codex";
if !is_nearai && !is_bedrock && !is_openai_codex && registry.find(&backend_lower).is_none()
if !is_nearai
&& !is_bedrock
&& !is_gemini_oauth
&& !is_openai_codex
&& registry.find(&backend_lower).is_none()
{
tracing::warn!(
"Unknown LLM backend '{}'. Will attempt as openai_compatible fallback.",
@@ -131,8 +138,8 @@ impl LlmConfig {
smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?,
};
// Resolve registry provider config (for non-NearAI, non-Bedrock, non-Codex backends)
let provider = if is_nearai || is_bedrock || is_openai_codex {
// Resolve registry provider config (for non-NearAI, non-Bedrock, non-Gemini, non-Codex backends)
let provider = if is_nearai || is_bedrock || is_gemini_oauth || is_openai_codex {
None
} else {
Some(Self::resolve_registry_provider(
@@ -213,6 +220,19 @@ impl LlmConfig {
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).
// Falls back to NearAI-specific cheap_model in provider chain logic.
let cheap_model = optional_env("LLM_CHEAP_MODEL")?;
@@ -226,6 +246,8 @@ impl LlmConfig {
"nearai".to_string()
} else if is_bedrock {
"bedrock".to_string()
} else if is_gemini_oauth {
"gemini_oauth".to_string()
} else if is_openai_codex {
"openai_codex".to_string()
} else if let Some(ref p) = provider {
@@ -237,6 +259,7 @@ impl LlmConfig {
nearai,
provider,
bedrock,
gemini_oauth,
openai_codex,
request_timeout_secs,
cheap_model,
+2 -2
View File
@@ -56,8 +56,8 @@ pub use self::tunnel::TunnelConfig;
pub use self::wasm::WasmConfig;
pub use self::workspace::WorkspaceConfig;
pub use crate::llm::config::{
BedrockConfig, CacheRetention, LlmConfig, NearAiConfig, OAUTH_PLACEHOLDER, OpenAiCodexConfig,
RegistryProviderConfig,
BedrockConfig, CacheRetention, GeminiOauthConfig, LlmConfig, NearAiConfig, OAUTH_PLACEHOLDER,
OpenAiCodexConfig, RegistryProviderConfig,
};
pub use crate::llm::session::SessionConfig;
-7
View File
@@ -196,12 +196,6 @@ pub struct JobContext {
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".
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 {
@@ -243,7 +237,6 @@ impl JobContext {
metadata: serde_json::Value::Null,
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
user_timezone: "UTC".to_string(),
tool_nesting_depth: 0,
}
}
-1
View File
@@ -134,7 +134,6 @@ impl JobStore for LibSqlBackend {
// TODO(#661): persist user_timezone in agent_jobs table so
// background/routine jobs retain the session's timezone context.
user_timezone: "UTC".to_string(),
tool_nesting_depth: 0,
}))
}
None => Ok(None),
+483 -55
View File
@@ -107,6 +107,21 @@ struct ChannelRuntimeState {
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)]
type TestWasmChannelLoader =
Arc<dyn Fn(&str) -> Result<LoadedChannel, ExtensionError> + Send + Sync>;
@@ -3341,6 +3356,46 @@ impl ExtensionManager {
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
// authoritative signal — setup secrets (client_id/secret) are
// intermediate and may be auto-resolved via builtins.
@@ -3363,31 +3418,13 @@ impl ExtensionManager {
};
}
// No auth section — fall back to checking setup.required_secrets.
let Some(setup) = &cap_file.setup else {
return ToolAuthState::NoAuth;
};
if setup.required_secrets.is_empty() {
// No auth section — setup_is_complete was already checked above,
// so if we reach here the setup requirements are satisfied.
if cap_file.setup.is_none() {
return ToolAuthState::NoAuth;
}
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
}
ToolAuthState::Ready
}
/// Check auth status for a WASM channel (read-only).
@@ -4273,6 +4310,102 @@ impl ExtensionManager {
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) {
let mut pending = self.pending_auth.write().await;
pending.retain(|_, auth| {
@@ -4287,11 +4420,12 @@ impl ExtensionManager {
});
}
/// Get the setup schema for an extension (secret fields and their status).
/// Get the setup schema for an extension (secret/text fields and their status).
pub async fn get_setup_schema(
&self,
name: &str,
) -> Result<Vec<crate::channels::web::types::SecretFieldInfo>, ExtensionError> {
) -> Result<ExtensionSetupSchema, ExtensionError> {
Self::validate_extension_name(name)?;
let kind = self.determine_installed_kind(name).await?;
match kind {
ExtensionKind::WasmChannel => {
@@ -4299,7 +4433,10 @@ impl ExtensionManager {
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
if !cap_path.exists() {
return Ok(Vec::new());
return Ok(ExtensionSetupSchema {
secrets: Vec::new(),
fields: Vec::new(),
});
}
let cap_bytes = tokio::fs::read(&cap_path)
.await
@@ -4308,14 +4445,14 @@ impl ExtensionManager {
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
.map_err(|e| ExtensionError::Other(e.to_string()))?;
let mut fields = Vec::new();
let mut secrets = Vec::new();
for secret in &cap_file.setup.required_secrets {
let provided = self
.secrets
.exists(&self.user_id, &secret.name)
.await
.unwrap_or(false);
fields.push(crate::channels::web::types::SecretFieldInfo {
secrets.push(crate::channels::web::types::SecretFieldInfo {
name: secret.name.clone(),
prompt: secret.prompt.clone(),
optional: secret.optional,
@@ -4323,17 +4460,27 @@ impl ExtensionManager {
auto_generate: secret.auto_generate.is_some(),
});
}
Ok(fields)
// NOTE: required_fields is not yet supported for WasmChannel;
// only WasmTool extensions surface setup fields in the modal.
Ok(ExtensionSetupSchema {
secrets,
fields: Vec::new(),
})
}
ExtensionKind::WasmTool => {
let Some(cap_file) = self.load_tool_capabilities(name).await else {
return Ok(Vec::new());
return Ok(ExtensionSetupSchema {
secrets: Vec::new(),
fields: Vec::new(),
});
};
let mut secrets = Vec::new();
let mut fields = Vec::new();
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 {
// Skip OAuth client_id/secret fields that resolve automatically
if Self::is_auto_resolved_oauth_field(&secret.name, &cap_file) {
continue;
}
@@ -4342,7 +4489,7 @@ impl ExtensionManager {
.exists(&self.user_id, &secret.name)
.await
.unwrap_or(false);
fields.push(crate::channels::web::types::SecretFieldInfo {
secrets.push(crate::channels::web::types::SecretFieldInfo {
name: secret.name.clone(),
prompt: secret.prompt.clone(),
optional: secret.optional,
@@ -4350,10 +4497,26 @@ impl ExtensionManager {
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(fields)
Ok(ExtensionSetupSchema { secrets, fields })
}
_ => Ok(Vec::new()),
_ => Ok(ExtensionSetupSchema {
secrets: Vec::new(),
fields: Vec::new(),
}),
}
}
@@ -4671,29 +4834,31 @@ impl ExtensionManager {
}
}
/// Save setup secrets for an extension, validating names against the capabilities schema.
/// Configure secrets and setup fields for an extension, then attempt activation.
///
/// Configure secrets for an extension: validate, store, auto-generate, and activate.
///
/// This is the single entrypoint for providing secrets to any extension.
/// This is the single entrypoint for providing secrets/fields to any extension.
/// Both the chat auth flow and the Extensions tab setup form call this method.
///
/// - Validates tokens against `validation_endpoint` (if declared in capabilities)
/// - 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)
/// - Activates the extension after configuration
pub async fn configure(
&self,
name: &str,
secrets: &std::collections::HashMap<String, String>,
fields: &std::collections::HashMap<String, String>,
) -> Result<ConfigureResult, ExtensionError> {
Self::validate_extension_name(name)?;
let kind = self.determine_installed_kind(name).await?;
// 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.
// Load allowed secret names and tool setup field definitions from capabilities.
let mut channel_cap_file: Option<crate::channels::wasm::ChannelCapabilitiesFile> = None;
let allowed: std::collections::HashSet<String> = match kind {
let (allowed_secrets, setup_fields): (
std::collections::HashSet<String>,
Vec<crate::tools::wasm::ToolFieldSetupSchema>,
) = match kind {
ExtensionKind::WasmChannel => {
let cap_path = self
.wasm_channels_dir
@@ -4717,27 +4882,28 @@ impl ExtensionManager {
.map(|s| s.name.clone())
.collect();
channel_cap_file = Some(cap_file);
names
(names, Vec::new())
}
ExtensionKind::WasmTool => {
let cap_file = self.load_tool_capabilities(name).await.ok_or_else(|| {
ExtensionError::Other(format!("Capabilities file not found for '{}'", name))
})?;
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 {
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 {
names.insert(auth.secret_name.clone());
}
if names.is_empty() {
if names.is_empty() && required_fields.is_empty() {
return Err(ExtensionError::Other(format!(
"Tool '{}' has no setup or auth schema — no secrets to configure",
"Tool '{}' has no setup or auth schema — nothing to configure",
name
)));
}
names
(names, required_fields)
}
ExtensionKind::McpServer => {
let server = self
@@ -4746,15 +4912,25 @@ impl ExtensionManager {
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
let mut names = std::collections::HashSet::new();
names.insert(server.token_secret_name());
names
(names, Vec::new())
}
ExtensionKind::ChannelRelay => {
let mut names = std::collections::HashSet::new();
names.insert(format!("relay:{}:stream_token", name));
names
(names, Vec::new())
}
};
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.
// The endpoint URL template uses {secret_name} placeholders that are
// substituted with the provided secret value before making the request.
@@ -4804,7 +4980,7 @@ impl ExtensionManager {
// Validate and store each submitted secret
for (secret_name, secret_value) in secrets {
if !allowed.contains(secret_name.as_str()) {
if !allowed_secrets.contains(secret_name.as_str()) {
return Err(ExtensionError::Other(format!(
"Unknown secret '{}' for extension '{}'",
secret_name, name
@@ -4822,6 +4998,70 @@ impl ExtensionManager {
.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)
if let Some(ref cap_file) = channel_cap_file {
for secret_def in &cap_file.setup.required_secrets {
@@ -4869,6 +5109,7 @@ impl ExtensionManager {
name, verification.instructions
),
activated: false,
restart_required,
auth_url: None,
verification: Some(verification),
});
@@ -4926,6 +5167,7 @@ impl ExtensionManager {
return Ok(ConfigureResult {
message,
activated: true,
restart_required,
auth_url,
verification: None,
});
@@ -4939,6 +5181,7 @@ impl ExtensionManager {
return Ok(ConfigureResult {
message: format!("Configuration saved for '{}'.", name),
activated: false,
restart_required,
auth_url: None,
verification: None,
});
@@ -4953,10 +5196,10 @@ impl ExtensionManager {
ExtensionKind::McpServer => self.activate_mcp(name).await,
ExtensionKind::ChannelRelay => self.activate_channel_relay(name).await,
ExtensionKind::WasmTool => {
// WasmTool is handled above and returns early; this branch is unreachable.
return Ok(ConfigureResult {
message: format!("Configuration saved for '{}'.", name),
activated: false,
restart_required,
auth_url: None,
verification: None,
});
@@ -4985,6 +5228,7 @@ impl ExtensionManager {
Ok(ConfigureResult {
message,
activated: true,
restart_required,
auth_url: None,
verification: None,
})
@@ -5008,6 +5252,7 @@ impl ExtensionManager {
name, e
),
activated: false,
restart_required,
auth_url: None,
verification: None,
})
@@ -5124,7 +5369,8 @@ impl ExtensionManager {
let mut secrets = std::collections::HashMap::new();
secrets.insert(secret_name, token.to_string());
self.configure(name, &secrets).await
self.configure(name, &secrets, &std::collections::HashMap::new())
.await
}
/// Read a capabilities.json file and revoke its credential mappings from
@@ -5650,11 +5896,16 @@ mod tests {
// after startup (e.g. via the web UI) would fail with "WASM runtime not
// 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.
fn make_test_manager_with_dirs(
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
tools_dir: std::path::PathBuf,
channels_dir: std::path::PathBuf,
store: Option<Arc<dyn crate::db::Database>>,
) -> crate::extensions::manager::ExtensionManager {
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
use crate::tools::mcp::process::McpProcessManager;
@@ -5681,7 +5932,7 @@ mod tests {
channels_dir,
None, // tunnel_url
"test".to_string(),
None, // db
store,
vec![],
)
}
@@ -5690,7 +5941,180 @@ mod tests {
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
tools_dir: std::path::PathBuf,
) -> crate::extensions::manager::ExtensionManager {
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir)
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir, None)
}
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]
@@ -6077,6 +6501,7 @@ mod tests {
"telegram_bot_token".to_string(),
"123456789:ABCdefGhI".to_string(),
)]),
&std::collections::HashMap::new(),
)
.await
.map_err(|err| format!("configure succeeds: {err}"))?;
@@ -6204,6 +6629,7 @@ mod tests {
"telegram_bot_token".to_string(),
"123456789:ABCdefGhI".to_string(),
)]),
&std::collections::HashMap::new(),
)
.await
.map_err(|err| format!("configure returned challenge: {err}"))?;
@@ -6720,7 +7146,7 @@ mod tests {
let dir = tempfile::tempdir().expect("temp dir");
let tools_dir = dir.path().join("tools");
let channels_dir = dir.path().join("channels");
let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone());
let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone(), None);
let wasm_path = channels_dir.join("telegram.wasm");
let cap_path = channels_dir.join("telegram.capabilities.json");
@@ -7369,7 +7795,9 @@ mod tests {
"tok".to_string(),
);
let result = mgr.configure("test-relay", &secrets).await;
let result = mgr
.configure("test-relay", &secrets, &std::collections::HashMap::new())
.await;
assert!(
result.is_ok(),
"configure should return Ok: {:?}",
+3 -1
View File
@@ -470,6 +470,8 @@ pub struct ConfigureResult {
pub message: String,
/// Whether the extension was successfully activated after configuration.
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).
pub auth_url: Option<String>,
/// Pending manual verification challenge (for Telegram owner binding, etc.).
@@ -498,7 +500,7 @@ pub struct InstalledExtension {
/// Tool names if active.
#[serde(default)]
pub tools: Vec<String>,
/// Whether this extension has a setup schema (required_secrets) that can be configured.
/// Whether this extension has a setup schema (required_secrets/required_fields) that can be configured.
#[serde(default)]
pub needs_setup: bool,
/// Whether this extension has an auth configuration (OAuth or manual token).
-1
View File
@@ -258,7 +258,6 @@ impl Store {
// TODO(#661): persist user_timezone in agent_jobs table so
// background/routine jobs retain the session's timezone context.
user_timezone: "UTC".to_string(),
tool_nesting_depth: 0,
}))
}
None => Ok(None),
+2
View File
@@ -1,5 +1,7 @@
//! Shared test helpers for OpenAI Codex provider tests.
#![cfg(test)]
use crate::config::OpenAiCodexConfig;
/// Build a minimal JWT for testing (header.payload.signature).
+33
View File
@@ -165,6 +165,8 @@ pub struct LlmConfig {
pub provider: Option<RegistryProviderConfig>,
/// AWS Bedrock config (populated when backend=bedrock, requires --features bedrock).
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).
pub openai_codex: Option<OpenAiCodexConfig>,
/// HTTP request timeout in seconds for LLM API calls.
@@ -267,3 +269,34 @@ 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.01.0)
/// - `GEMINI_TOP_K` — top-k sampling (integer)
/// - `GEMINI_SEED` — deterministic generation seed
/// - `GEMINI_PRESENCE_PENALTY` — presence penalty (-2.02.0)
/// - `GEMINI_FREQUENCY_PENALTY` — frequency penalty (-2.02.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
+55
View File
@@ -18,6 +18,7 @@ pub mod config;
pub mod costs;
pub mod error;
pub mod failover;
pub mod gemini_oauth;
mod github_copilot;
pub(crate) mod github_copilot_auth;
mod nearai_chat;
@@ -50,6 +51,7 @@ pub use config::{
};
pub use error::LlmError;
pub use failover::{CooldownConfig, FailoverProvider};
pub use gemini_oauth::GeminiOauthProvider;
pub use nearai_chat::{DEFAULT_MODEL, ModelInfo, NearAiChatProvider, default_models};
pub use openai_codex_provider::OpenAiCodexProvider;
pub use openai_codex_session::{OpenAiCodexSession, OpenAiCodexSessionManager};
@@ -93,6 +95,10 @@ pub async fn create_llm_provider(
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
if config.backend == "bedrock" {
#[cfg(feature = "bedrock")]
@@ -490,6 +496,19 @@ 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
let reg_config = config.provider.as_ref().ok_or_else(|| LlmError::RequestFailed {
provider: config.backend.clone(),
@@ -674,6 +693,17 @@ pub async fn build_provider_chain(
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)]
mod tests {
use super::*;
@@ -705,6 +735,7 @@ mod tests {
nearai: test_nearai_config(),
provider: None,
bedrock: None,
gemini_oauth: None,
request_timeout_secs: 120,
cheap_model: None,
smart_routing_cascade: true,
@@ -786,6 +817,30 @@ 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]
fn test_cheap_model_name_resolution() {
// Generic takes priority
+1
View File
@@ -344,6 +344,7 @@ pub(crate) fn build_nearai_model_fetch_config() -> crate::config::LlmConfig {
nearai: crate::config::NearAiConfig::for_model_discovery(),
provider: None,
bedrock: None,
gemini_oauth: None,
request_timeout_secs: 120,
cheap_model: None,
smart_routing_cascade: false,
-2
View File
@@ -306,8 +306,6 @@ async fn async_main() -> anyhow::Result<()> {
&components.llm,
components.db.as_ref(),
components.secrets_store.as_ref(),
&components.tools,
&components.safety,
)
.await;
let container_job_manager = orch.container_job_manager;
+14 -552
View File
@@ -15,18 +15,15 @@ use tokio::sync::{Mutex, broadcast};
use uuid::Uuid;
use crate::channels::web::types::SseEvent;
use crate::context::JobContext;
use crate::db::Database;
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
use crate::orchestrator::job_manager::ContainerJobManager;
use crate::secrets::SecretsStore;
use crate::tools::ToolExecutor;
use crate::worker::api::JobEventPayload;
use crate::worker::api::{
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
ToolCallRequest, ToolCallResponse,
};
/// A follow-up prompt queued for a Claude Code bridge.
@@ -52,8 +49,6 @@ pub struct OrchestratorState {
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
/// User ID for secret lookups (single-tenant, typically "default").
pub user_id: String,
/// Tool executor for programmatic tool calling (PTC).
pub tool_executor: Option<Arc<ToolExecutor>>,
}
/// The orchestrator's internal API server.
@@ -75,7 +70,6 @@ impl OrchestratorApi {
.route("/worker/{job_id}/event", post(job_event_handler))
.route("/worker/{job_id}/prompt", get(get_prompt_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(
state.token_store.clone(),
worker_auth_middleware,
@@ -297,26 +291,20 @@ async fn job_event_handler(
.unwrap_or("")
.to_string(),
},
"tool_use" => {
// Redact raw parameters from worker-reported tool_use events
// before broadcasting via SSE. Workers are untrusted and may
// include sensitive data (API keys, passwords, PII) in the
// input payload. We replace it with a placeholder to prevent
// leaking secrets to the web UI.
let redacted_input = serde_json::json!({
"_note": "parameters redacted for security"
});
SseEvent::JobToolUse {
job_id: job_id_str,
tool_name: payload
.data
.get("tool_name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string(),
input: redacted_input,
}
}
"tool_use" => SseEvent::JobToolUse {
job_id: job_id_str,
tool_name: payload
.data
.get("tool_name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string(),
input: payload
.data
.get("input")
.cloned()
.unwrap_or(serde_json::Value::Null),
},
"tool_result" => SseEvent::JobToolResult {
job_id: job_id_str,
tool_name: payload
@@ -455,106 +443,6 @@ 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 {
match reason {
crate::llm::FinishReason::Stop => "stop".to_string(),
@@ -592,7 +480,6 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
}
}
@@ -822,7 +709,6 @@ mod tests {
store: None,
secrets_store: Some(secrets_store),
user_id: "default".to_string(),
tool_executor: None,
};
let router = OrchestratorApi::router(state);
@@ -858,7 +744,6 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
};
let job_id = Uuid::new_v4();
@@ -914,7 +799,6 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
};
let job_id = Uuid::new_v4();
@@ -963,7 +847,6 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
};
let job_id = Uuid::new_v4();
@@ -1043,425 +926,4 @@ mod tests {
assert_eq!(handle.worker_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),
}
}
}
-16
View File
@@ -49,9 +49,7 @@ use uuid::Uuid;
use crate::channels::web::types::SseEvent;
use crate::db::Database;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
use crate::secrets::SecretsStore;
use crate::tools::{ToolExecutor, ToolRegistry};
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
/// variable, falling back to 50051.
@@ -77,8 +75,6 @@ pub async fn setup_orchestrator(
llm: &Arc<dyn LlmProvider>,
db: Option<&Arc<dyn Database>>,
secrets_store: Option<&Arc<dyn SecretsStore + Send + Sync>>,
tools: &Arc<ToolRegistry>,
safety: &Arc<SafetyLayer>,
) -> OrchestratorSetup {
let prompt_queue = Arc::new(Mutex::new(
HashMap::<Uuid, VecDeque<api::PendingPrompt>>::new(),
@@ -129,17 +125,6 @@ pub async fn setup_orchestrator(
};
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 {
llm: Arc::clone(llm),
job_manager: Arc::clone(&jm),
@@ -149,7 +134,6 @@ pub async fn setup_orchestrator(
store: db.cloned(),
secrets_store: secrets_store.cloned(),
user_id: "default".to_string(),
tool_executor: Some(tool_executor),
};
tokio::spawn(async move {
+206 -103
View File
@@ -1078,23 +1078,40 @@ impl SetupWizard {
.map(|s| s.display_name().to_string())
.unwrap_or_else(|| def.id.clone())
} else {
current.clone()
match current.as_str() {
"nearai" => "NEAR AI".to_string(),
"gemini_oauth" | "gemini-oauth" => "Gemini API (OAuth)".to_string(),
_ => {
if let Some(def) = registry.find(&current) {
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));
println!();
let is_known = current == "nearai"
|| current == "bedrock"
|| current == "gemini_oauth"
|| current == "gemini-oauth"
|| current == "openai_codex"
|| registry.is_known(&current);
if is_known && confirm("Keep current provider?", true).map_err(SetupError::Io)? {
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.");
return Ok(());
}
if current == "gemini_oauth" || current == "gemini-oauth" {
print_info("Keeping existing Gemini CLI OAuth configuration.");
return Ok(());
}
if current == "openai_codex" {
print_info("Keeping existing OpenAI Codex configuration.");
return Ok(());
@@ -1113,13 +1130,15 @@ impl SetupWizard {
print_info("Select your inference provider:");
println!();
// Build menu: NearAI first, then OpenAI Codex, then registry providers, then Bedrock
// Build menu: NearAI first, then Gemini OAuth, then OpenAI Codex, then registry providers, then Bedrock
let selectable = registry.selectable();
let mut options: Vec<String> = Vec::with_capacity(2 + selectable.len());
let mut provider_ids: Vec<String> = Vec::with_capacity(2 + selectable.len());
let mut options: Vec<String> = Vec::with_capacity(3 + selectable.len());
let mut provider_ids: Vec<String> = Vec::with_capacity(3 + selectable.len());
options.push("NEAR AI - multi-model access via NEAR account".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());
provider_ids.push("openai_codex".to_string());
@@ -1147,6 +1166,8 @@ impl SetupWizard {
if selected_id == "bedrock" {
self.setup_bedrock().await?;
} else if selected_id == "gemini_oauth" {
self.setup_gemini_oauth().await?;
} else {
self.run_provider_setup(selected_id, &registry).await?;
}
@@ -1795,6 +1816,40 @@ impl SetupWizard {
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.
///
/// Branches on the selected LLM backend and fetches models from the
@@ -1818,109 +1873,157 @@ impl SetupWizard {
let backend = self.settings.llm_backend.as_deref().unwrap_or("nearai");
let registry = crate::llm::ProviderRegistry::load();
if backend == "nearai" {
// NEAR AI: use existing provider list_models()
let fetched = self.fetch_nearai_models().await;
let models = if fetched.is_empty() {
crate::llm::default_models()
} else {
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
}
};
// Apply models_filter from setup hint (e.g., Groq "chat" filters non-chat models)
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));
match backend {
"nearai" => {
// NEAR AI: use existing provider list_models()
let fetched = self.fetch_nearai_models().await;
let models = if fetched.is_empty() {
crate::llm::default_models()
} else {
self.select_from_model_list(&models)?;
}
} else {
// Manual model entry
let default = &def.default_model;
fetched.iter().map(|m| (m.clone(), m.clone())).collect()
};
self.select_from_model_list(&models)?;
}
"gemini_oauth" | "gemini-oauth" => {
let default_models: Vec<(String, String)> = vec![
(
"gemini-3.1-pro-preview".into(),
"Gemini 3.1 Pro (Latest, strongest reasoning)".into(),
),
(
"gemini-3.1-pro-preview-customtools".into(),
"Gemini 3.1 Pro Custom Tools (Enhanced tool use)".into(),
),
(
"gemini-3-pro-preview".into(),
"Gemini 3 Pro (Preview)".into(),
),
(
"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(&format!("Model name (default: {default})")).map_err(SetupError::Io)?;
let model_id = if model_id.is_empty() {
default.clone()
} else {
model_id
};
input("Bedrock model ID (e.g., anthropic.claude-v3-sonnet-20240229-v1:0)")
.map_err(SetupError::Io)?;
if model_id.is_empty() {
return Err(SetupError::Config("Model ID is required".to_string()));
}
self.settings.selected_model = Some(model_id.clone());
print_success(&format!("Selected {}", model_id));
}
} else if backend == "bedrock" {
let model_id = input("Bedrock model ID (e.g., anthropic.claude-opus-4-6-v1)")
.map_err(SetupError::Io)?;
if model_id.is_empty() {
return Err(SetupError::Config("Model ID is required".to_string()));
_ => {
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
}
};
// 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(())
-2
View File
@@ -9,7 +9,6 @@ mod json;
mod memory;
mod message;
pub mod path_utils;
pub mod ptc_script;
mod restart;
pub mod routine;
pub mod secrets_tools;
@@ -32,7 +31,6 @@ pub use job::{
pub use json::JsonTool;
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
pub use message::MessageTool;
pub use ptc_script::PtcScriptTool;
pub use restart::RestartTool;
pub use routine::{
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
-371
View File
@@ -1,371 +0,0 @@
//! 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(&params, "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());
}
}
+17 -5
View File
@@ -45,11 +45,23 @@ impl ToolInfoDetail {
}
fn schema_param_names(schema: &serde_json::Value) -> Vec<String> {
schema
.get("properties")
.and_then(|p| p.as_object())
.map(|props| props.keys().cloned().collect())
.unwrap_or_default()
let mut names = std::collections::BTreeSet::new();
if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
names.extend(props.keys().cloned());
}
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 {
+701 -12
View File
@@ -1,4 +1,4 @@
pub(crate) fn prepare_tool_params(
pub fn prepare_tool_params(
tool: &dyn crate::tools::tool::Tool,
params: &serde_json::Value,
) -> serde_json::Value {
@@ -9,14 +9,87 @@ pub(crate) fn prepare_params_for_schema(
params: &serde_json::Value,
schema: &serde_json::Value,
) -> serde_json::Value {
coerce_value(params, schema)
let resolved = resolve_refs(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 {
// This coercer intentionally handles the concrete schema shapes we expose in
// discovery today. It does not resolve combinators like anyOf/oneOf/allOf or
// references via $ref; those schemas pass through unchanged unless they also
// advertise a directly coercible type/property shape.
// This coercer handles concrete schema shapes including discriminated unions
// (oneOf/anyOf with const or single-element enum discriminators), allOf
// merges, and $ref references (resolved in a pre-pass).
if value.is_null() {
return value.clone();
}
@@ -47,12 +120,35 @@ fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_
return value.clone();
}
let properties = schema.get("properties").and_then(|p| p.as_object());
let additional_schema = schema.get("additionalProperties").filter(|v| v.is_object());
let resolved = resolve_effective_properties(schema, obj);
let properties = resolved
.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();
for (key, current) in &mut coerced {
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);
continue;
}
@@ -68,11 +164,179 @@ fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_
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> {
// 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") {
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")
&& let Ok(v) = s.parse::<i64>()
{
@@ -114,10 +378,15 @@ fn schema_allows_type(schema: &serde_json::Value, expected: &str) -> bool {
Some(serde_json::Value::String(t)) => t == expected,
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
_ => match expected {
"object" => schema
.get("properties")
.and_then(|p| p.as_object())
.is_some(),
"object" => {
schema
.get("properties")
.and_then(|p| p.as_object())
.is_some()
|| schema.get("oneOf").is_some()
|| schema.get("anyOf").is_some()
|| schema.get("allOf").is_some()
}
"array" => schema.get("items").is_some(),
_ => false,
},
@@ -325,6 +594,91 @@ mod tests {
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(&params, &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(&params, &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(&params, &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(&params, &schema);
// Required string-only field keeps empty string
assert_eq!(result["name"], serde_json::json!(""));
}
#[test]
fn permissive_schema_is_noop() {
let schema = serde_json::json!({
@@ -339,6 +693,341 @@ mod tests {
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(&params, &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(&params, &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(&params, &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(&params, &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(&params, &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(&params, &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(&params, &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(&params, &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(&params, &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(&params, &schema);
assert_eq!(result["root"]["value"], serde_json::json!(42));
}
#[test]
fn prepare_tool_params_uses_discovery_schema() {
let tool = StubTool {
-490
View File
@@ -1,490 +0,0 @@
//! 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
);
}
}
}
-2
View File
@@ -18,7 +18,6 @@ pub mod redaction;
pub mod schema_validator;
pub mod wasm;
mod executor;
mod registry;
mod tool;
@@ -32,7 +31,6 @@ pub use builder::{
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
};
pub(crate) use coercion::prepare_tool_params;
pub use executor::{PtcError, PtcToolResult, ToolExecutor};
pub use rate_limiter::RateLimiter;
pub use registry::ToolRegistry;
pub use tool::{
+5 -73
View File
@@ -19,12 +19,11 @@ use crate::tools::builder::{
use crate::tools::builtin::{
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, PtcScriptTool,
ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool,
TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool,
ToolSearchTool, ToolUpgradeTool, WriteFileTool,
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool,
ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool,
ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool,
ToolUpgradeTool, WriteFileTool,
};
use crate::tools::executor::ToolExecutor;
use crate::tools::rate_limiter::RateLimiter;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain};
use crate::tools::wasm::{
@@ -79,7 +78,6 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[
"image_edit",
"image_analyze",
"tool_info",
"ptc_script",
];
/// Registry of available tools.
@@ -95,14 +93,6 @@ pub struct ToolRegistry {
rate_limiter: RateLimiter,
/// Reference to the message tool for setting context per-turn.
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 {
@@ -124,7 +114,6 @@ impl ToolRegistry {
secrets_store: None,
rate_limiter: RateLimiter::new(),
message_tool: RwLock::new(None),
tool_executor_slot: Arc::new(std::sync::RwLock::new(None)),
}
}
@@ -149,27 +138,6 @@ impl ToolRegistry {
&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.
pub async fn register(&self, tool: Arc<dyn Tool>) {
let name = tool.name().to_string();
@@ -362,9 +330,8 @@ impl ToolRegistry {
self.register_sync(Arc::new(WriteFileTool::new()));
self.register_sync(Arc::new(ListDirTool::new()));
self.register_sync(Arc::new(ApplyPatchTool::new()));
self.register_sync(Arc::new(PtcScriptTool::new()));
tracing::debug!("Registered 6 development tools");
tracing::debug!("Registered 5 development tools");
}
/// Register memory tools with a workspace.
@@ -692,11 +659,6 @@ impl ToolRegistry {
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
self.register(Arc::new(wrapper)).await;
@@ -927,36 +889,6 @@ mod tests {
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]
async fn test_builtin_tool_cannot_be_shadowed() {
let registry = ToolRegistry::new();
+83 -5
View File
@@ -42,11 +42,38 @@ 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.
fn check_object_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
let mut errors = Vec::new();
// Rule 1: must have "type": "object"
// Report non-array combinator values as errors.
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()) {
Some("object") => {}
Some(other) => {
@@ -54,16 +81,67 @@ fn check_object_schema(schema: &serde_json::Value, path: &str) -> Vec<String> {
return errors;
}
None => {
errors.push(format!("{path}: missing \"type\": \"object\""));
return errors;
if !has_combinators {
errors.push(format!("{path}: missing \"type\": \"object\""));
return errors;
}
}
}
// Rule 2: must have "properties" as an object
// Validate combinator variants recursively
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()) {
Some(p) => p,
None => {
errors.push(format!("{path}: missing or non-object \"properties\""));
if !has_combinators {
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;
}
};
+87 -5
View File
@@ -462,6 +462,22 @@ pub fn redact_params(params: &serde_json::Value, sensitive: &[&str]) -> serde_js
/// on maliciously crafted schemas.
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> {
validate_tool_schema_inner(schema, path, 0)
}
@@ -476,7 +492,18 @@ fn validate_tool_schema_inner(schema: &serde_json::Value, path: &str, depth: usi
return errors;
}
// Rule 1: must have "type": "object" at this level
// Report non-array combinator values as errors.
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()) {
Some("object") => {}
Some(other) => {
@@ -484,16 +511,71 @@ fn validate_tool_schema_inner(schema: &serde_json::Value, path: &str, depth: usi
return errors; // Can't check further
}
None => {
errors.push(format!("{path}: missing \"type\": \"object\""));
return errors;
if !has_combinators {
errors.push(format!("{path}: missing \"type\": \"object\""));
return errors;
}
}
}
// Rule 2: must have "properties" as an object
// Validate combinator variants recursively
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()) {
Some(p) => p,
None => {
errors.push(format!("{path}: missing or non-object \"properties\""));
if !has_combinators {
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;
}
};
+99
View File
@@ -708,6 +708,9 @@ pub struct ToolSetupSchema {
/// Secrets the user must provide before the tool can be used.
#[serde(default)]
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.
@@ -722,6 +725,46 @@ pub struct ToolSecretSetupSchema {
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)]
mod tests {
use crate::tools::wasm::capabilities_schema::{CapabilitiesFile, CredentialLocationSchema};
@@ -1218,6 +1261,20 @@ mod tests {
"prompt": "Google OAuth Client Secret",
"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"
}
]
}
}"#;
@@ -1230,6 +1287,48 @@ mod tests {
assert!(!setup.required_secrets[0].optional);
assert_eq!(setup.required_secrets[1].name, "google_oauth_client_secret");
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]
+1 -1
View File
@@ -139,5 +139,5 @@ pub use loader::{
// Capabilities schema (for parsing *.capabilities.json files)
pub use capabilities_schema::{
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, RateLimitSchema,
ValidationEndpointSchema,
ToolFieldSetupSchema, ToolSetupFieldInputType, ToolSetupSchema, ValidationEndpointSchema,
};
+193 -402
View File
@@ -17,9 +17,9 @@ use wasmtime::component::Linker;
use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
use crate::context::JobContext;
use crate::llm::recording::{HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor};
use crate::safety::LeakDetector;
use crate::secrets::SecretsStore;
use crate::tools::ToolExecutor;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::tools::wasm::capabilities::Capabilities;
use crate::tools::wasm::credential_injector::{
@@ -30,26 +30,6 @@ use crate::tools::wasm::host::{HostState, LogLevel};
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
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.
//
// This creates:
@@ -120,11 +100,9 @@ struct StoreData {
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
/// Reused across multiple `http_request` calls within one execution.
http_runtime: Option<tokio::runtime::Runtime>,
/// Optional tool resolver for programmatic tool calling (PTC).
/// When set, WASM tools can invoke other tools via the `tool_invoke` host function.
tool_resolver: Option<ToolResolver>,
/// Current nesting depth for tool_invoke calls. Prevents infinite recursion.
tool_nesting_depth: u32,
/// Optional HTTP interceptor for testing — returns canned responses
/// instead of making real requests when set.
http_interceptor: Option<Arc<dyn HttpInterceptor>>,
}
impl StoreData {
@@ -133,7 +111,6 @@ impl StoreData {
capabilities: Capabilities,
credentials: HashMap<String, String>,
host_credentials: Vec<ResolvedHostCredential>,
tool_resolver: Option<ToolResolver>,
) -> Self {
// Minimal WASI context: no filesystem, no env vars (security)
let wasi = WasiCtxBuilder::new().build();
@@ -146,8 +123,7 @@ impl StoreData {
credentials,
host_credentials,
http_runtime: None,
tool_resolver,
tool_nesting_depth: 0,
http_interceptor: None,
}
}
@@ -373,6 +349,59 @@ impl near::agent::host::Host for StoreData {
);
}
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 client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(10))
@@ -463,43 +492,63 @@ 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
result.map_err(|e| self.redact_credentials(&e))
}
fn tool_invoke(&mut self, alias: String, params_json: String) -> Result<String, String> {
use crate::tools::executor::MAX_NESTING_DEPTH;
fn tool_invoke(&mut self, alias: String, _params_json: String) -> Result<String, String> {
// 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()?;
// Check nesting depth
if self.tool_nesting_depth >= MAX_NESTING_DEPTH {
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(&params_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)
// Tool invocation requires async context and access to the tool registry,
// which aren't available inside a synchronous WASM callback.
Err("Tool invocation from WASM tools is not yet supported".to_string())
}
fn secret_exists(&mut self, name: String) -> bool {
@@ -530,11 +579,9 @@ pub struct WasmToolWrapper {
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
/// OAuth refresh configuration for auto-refreshing expired tokens.
oauth_refresh: Option<OAuthRefreshConfig>,
/// Direct tool executor reference (for tests that wire it explicitly).
tool_executor: Option<Arc<ToolExecutor>>,
/// 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>>>>>,
/// Optional HTTP interceptor for testing — returns canned responses
/// instead of making real requests when set.
http_interceptor: Option<Arc<dyn HttpInterceptor>>,
}
#[derive(Debug, Clone)]
@@ -561,23 +608,51 @@ impl WasmToolSchemas {
}
fn is_permissive_schema(schema: &serde_json::Value) -> bool {
schema
if schema
.get("properties")
.and_then(|p| p.as_object())
.is_none_or(|p| p.is_empty())
.is_some_and(|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 {
schema
.get("properties")
.and_then(|p| p.as_object())
.map(|props| {
props
.values()
.filter(|prop| schema_is_typed_property(prop))
.count()
})
.unwrap_or(0)
let mut all_props = serde_json::Map::new();
if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
all_props.extend(props.iter().map(|(k, v)| (k.clone(), v.clone())));
}
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()) {
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 {
@@ -623,11 +698,20 @@ impl WasmToolWrapper {
credentials: HashMap::new(),
secrets_store: None,
oauth_refresh: None,
tool_executor: None,
tool_executor_slot: None,
http_interceptor: 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.
pub fn with_description(mut self, description: impl Into<String>) -> Self {
self.description = description.into();
@@ -679,28 +763,6 @@ impl WasmToolWrapper {
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.
pub fn limits(&self) -> &ResourceLimits {
&self.prepared.limits
@@ -729,19 +791,18 @@ impl WasmToolWrapper {
params: serde_json::Value,
context_json: Option<String>,
host_credentials: Vec<ResolvedHostCredential>,
tool_resolver: Option<ToolResolver>,
) -> Result<(String, Vec<crate::tools::wasm::host::LogEntry>), WasmError> {
let engine = self.runtime.engine();
let limits = &self.prepared.limits;
// Create store with fresh state (NEAR pattern: fresh instance per call)
let store_data = StoreData::new(
let mut store_data = StoreData::new(
limits.memory_bytes,
self.capabilities.clone(),
self.credentials.clone(),
host_credentials,
tool_resolver,
);
store_data.http_interceptor = self.http_interceptor.clone();
let mut store = Store::new(engine, store_data);
// Configure fuel if enabled
@@ -839,7 +900,6 @@ pub(super) fn extract_wasm_metadata(
Capabilities::default(),
HashMap::new(),
vec![],
None,
);
let mut store = Store::new(engine, store_data);
@@ -939,48 +999,6 @@ impl Tool for WasmToolWrapper {
// Serialize context for WASM
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
let runtime = Arc::clone(&self.runtime);
let prepared = Arc::clone(&self.prepared);
@@ -998,14 +1016,13 @@ impl Tool for WasmToolWrapper {
description,
schemas,
credentials,
secrets_store: None, // Not needed in blocking task
oauth_refresh: None, // Already used above for pre-refresh
tool_executor: None, // Resolver closure captures the executor
tool_executor_slot: None, // Resolver closure captures the executor
secrets_store: None, // Not needed in blocking task
oauth_refresh: None, // Already used above for pre-refresh
http_interceptor: self.http_interceptor.clone(),
};
tokio::task::spawn_blocking(move || {
wrapper.execute_sync(params, context_json, host_credentials, tool_resolver)
wrapper.execute_sync(params, context_json, host_credentials)
})
.await
.map_err(|e| WasmError::ExecutionPanicked(e.to_string()))?
@@ -1450,15 +1467,33 @@ fn is_private_ip(ip: std::net::IpAddr) -> bool {
}
fn schema_contains_container_properties(schema: &serde_json::Value) -> bool {
schema
let has_container = |props: &serde_json::Map<String, serde_json::Value>| {
props
.values()
.any(|prop| schema_declares_type(prop, "array") || schema_declares_type(prop, "object"))
};
if schema
.get("properties")
.and_then(|p| p.as_object())
.map(|props| {
props.values().any(|prop| {
schema_declares_type(prop, "array") || schema_declares_type(prop, "object")
.is_some_and(has_container)
{
return true;
}
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)
})
})
.unwrap_or(false)
{
return true;
}
}
false
}
fn schema_declares_type(schema: &serde_json::Value, expected: &str) -> bool {
@@ -1516,7 +1551,6 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
@@ -1533,12 +1567,10 @@ mod tests {
TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET,
test_secrets_store,
};
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::tools::tool::Tool;
use crate::tools::wasm::capabilities::Capabilities;
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
use super::WasmToolWrapper;
struct RecordingSecretsStore {
inner: InMemorySecretsStore,
get_decrypted_lookups: Mutex<Vec<(String, String)>>,
@@ -1766,7 +1798,6 @@ mod tests {
Capabilities::default(),
HashMap::new(),
host_credentials,
None,
);
// Should inject for matching host
@@ -1806,7 +1837,6 @@ mod tests {
Capabilities::default(),
HashMap::new(),
host_credentials,
None,
);
let mut headers = HashMap::new();
@@ -1833,7 +1863,6 @@ mod tests {
Capabilities::default(),
HashMap::new(),
host_credentials,
None,
);
let text = "Error: request to https://api.example.com?key=super-secret-token failed";
@@ -2320,244 +2349,6 @@ mod tests {
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(&params, &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(&params, &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(&params, &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
/// injection), not after. If it ran post-injection, the host-injected
/// Slack bot token (`xoxb-...`) would trigger a Block and reject the
-38
View File
@@ -114,36 +114,6 @@ pub struct CredentialResponse {
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 {
/// Create a new client from environment.
///
@@ -429,14 +399,6 @@ 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.
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
let _: serde_json::Value = self
+408
View File
@@ -343,4 +343,412 @@ mod tests {
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();
}
}
+277
View File
@@ -0,0 +1,277 @@
//! 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();
}
}
+99
View File
@@ -0,0 +1,99 @@
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");
}
+112 -15
View File
@@ -23,7 +23,7 @@ use crate::support::metrics::{ToolInvocation, TraceMetrics};
use crate::support::test_channel::{TestChannel, TestChannelHandle};
use crate::support::trace_llm::{LlmTrace, TraceLlm};
use ironclaw::llm::recording::{HttpExchange, ReplayingHttpInterceptor};
use ironclaw::llm::recording::{HttpExchange, HttpInterceptor, ReplayingHttpInterceptor};
// ---------------------------------------------------------------------------
// TestRig
@@ -343,6 +343,13 @@ impl Drop for TestRig {
// 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`.
pub struct TestRigBuilder {
trace: Option<LlmTrace>,
@@ -354,6 +361,7 @@ pub struct TestRigBuilder {
enable_routines: bool,
http_exchanges: Vec<HttpExchange>,
extra_tools: Vec<Arc<dyn Tool>>,
wasm_tools: Vec<WasmToolSpec>,
keep_bootstrap: bool,
}
@@ -370,10 +378,34 @@ impl TestRigBuilder {
enable_routines: false,
http_exchanges: Vec::new(),
extra_tools: Vec::new(),
wasm_tools: Vec::new(),
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.
pub fn with_trace(mut self, trace: LlmTrace) -> Self {
self.trace = Some(trace);
@@ -465,6 +497,7 @@ impl TestRigBuilder {
enable_routines,
http_exchanges: explicit_http_exchanges,
extra_tools,
wasm_tools,
keep_bootstrap,
} = self;
@@ -560,6 +593,20 @@ impl TestRigBuilder {
let scheduler_slot: ironclaw::tools::builtin::SchedulerSlot =
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.
{
// Ensure filesystem/shell dev tools are always available in the
@@ -620,6 +667,69 @@ impl TestRigBuilder {
for tool in extra_tools {
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.
@@ -643,20 +753,7 @@ impl TestRigBuilder {
hooks: components.hooks,
cost_guard: components.cost_guard,
sse_tx: None,
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>)
}
},
http_interceptor,
transcription: None,
document_extraction: None,
sandbox_readiness: ironclaw::agent::SandboxReadiness::Available, // tests don't use real Docker
-1
View File
@@ -65,7 +65,6 @@ async fn core_registration_covers_expected_tools() {
"http",
"json",
"list_dir",
"ptc_script",
"read_file",
"shell",
"time",
-20
View File
@@ -1,20 +0,0 @@
[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]
-52
View File
@@ -1,52 +0,0 @@
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);