feat: V1.0 (#734)

Co-authored-by: Michael Neale <michael.neale@gmail.com>
Co-authored-by: Wendy Tang <wendytang@squareup.com>
Co-authored-by: Jarrod Sibbison <72240382+jsibbison-square@users.noreply.github.com>
Co-authored-by: Alex Hancock <alex.hancock@example.com>
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
Co-authored-by: Lifei Zhou <lifei@squareup.com>
Co-authored-by: Wes <141185334+wesrblock@users.noreply.github.com>
Co-authored-by: Max Novich <maksymstepanenko1990@gmail.com>
Co-authored-by: Zaki Ali <zaki@squareup.com>
Co-authored-by: Salman Mohammed <smohammed@squareup.com>
Co-authored-by: Kalvin C <kalvinnchau@users.noreply.github.com>
Co-authored-by: Alec Thomas <alec@swapoff.org>
Co-authored-by: lily-de <119957291+lily-de@users.noreply.github.com>
Co-authored-by: kalvinnchau <kalvin@block.xyz>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Rizel Scarlett <rizel@squareup.com>
Co-authored-by: bwrage <bwrage@squareup.com>
Co-authored-by: Kalvin Chau <kalvin@squareup.com>
Co-authored-by: Alice Hau <110418948+ahau-square@users.noreply.github.com>
Co-authored-by: Alistair Gray <ajgray@stripe.com>
Co-authored-by: Nahiyan Khan <nahiyan.khan@gmail.com>
Co-authored-by: Alex Hancock <alexhancock@squareup.com>
Co-authored-by: Nahiyan Khan <nahiyan@squareup.com>
Co-authored-by: marcelle <1852848+laanak08@users.noreply.github.com>
Co-authored-by: Yingjie He <yingjiehe@block.xyz>
Co-authored-by: Yingjie He <yingjiehe@squareup.com>
Co-authored-by: Lily Delalande <ldelalande@block.xyz>
Co-authored-by: Adewale Abati <acekyd01@gmail.com>
Co-authored-by: Ebony Louis <ebony774@gmail.com>
Co-authored-by: Angie Jones <jones.angie@gmail.com>
Co-authored-by: Ebony Louis <55366651+EbonyLouis@users.noreply.github.com>
This commit is contained in:
Bradley Axen
2025-01-24 13:04:43 -08:00
committed by GitHub
parent eccb1b2261
commit 1c9a7c0b05
688 changed files with 71147 additions and 19132 deletions

27
.cursorrules Normal file
View File

@@ -0,0 +1,27 @@
You are an expert programmer in Rust teaching who is teaching another developer who is learning Rust.
The students are familiar with programming in languages such as Python (advanced), Java (novice) and C (novice) so
when possible use analogies from those languages.
Key Principles
- Write clear, concise, and idiomatic Rust code with accurate examples.
- Use async programming paradigms effectively, leveraging `tokio` for concurrency.
- Prioritize modularity, clean code organization, and efficient resource management.
- Use expressive variable names that convey intent (e.g., `is_ready`, `has_data`).
- Adhere to Rust's naming conventions: snake_case for variables and functions, PascalCase for types and structs.
- Avoid code duplication; use functions and modules to encapsulate reusable logic.
- Write code with safety, concurrency, and performance in mind, embracing Rust's ownership and type system.
Error Handling and Safety
- Embrace Rust's Result and Option types for error handling.
- Use `?` operator to propagate errors in async functions.
- Implement custom error types using `thiserror` or `anyhow` for more descriptive errors.
- Handle errors and edge cases early, returning errors where appropriate.
- Use `.await` responsibly, ensuring safe points for context switching.
Key Conventions
1. Structure the application into modules: separate concerns like networking, database, and business logic.
2. Use environment variables for configuration management (e.g., `dotenv` crate).
3. Ensure code is well-documented with inline comments and Rustdoc.
4. Do not use the older style of "MOD/mod.rs" for separing modules and instead use the "MOD.rs" filename convention.
Refer to "The Rust Programming Language" book (2024 version) and "Command line apps in Rust" documentation for in-depth information on best practices, and advanced features.

69
.github/workflows/build-cli.yml vendored Normal file
View File

@@ -0,0 +1,69 @@
# This is a **reuseable** workflow that bundles the Desktop App for macOS.
# It doesn't get triggered on its own. It gets used in multiple workflows:
# - release.yml
# - canary.yml
on:
workflow_call:
inputs:
# Let's allow overriding the OSes and architectures in JSON array form:
# e.g. '["ubuntu-latest","macos-latest"]'
# If no input is provided, these defaults apply.
operating-systems:
type: string
required: false
default: '["ubuntu-latest","macos-latest"]'
architectures:
type: string
required: false
default: '["x86_64","aarch64"]'
name: "Reusable workflow to build CLI"
jobs:
build-cli:
name: Build CLI
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
os: ${{ fromJson(inputs.operating-systems) }}
architecture: ${{ fromJson(inputs.architectures) }}
include:
- os: ubuntu-latest
target-suffix: unknown-linux-gnu
- os: macos-latest
target-suffix: apple-darwin
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Setup Rust
uses: dtolnay/rust-toolchain@stable
with:
toolchain: stable
target: ${{ matrix.architecture }}-${{ matrix.target-suffix }}
- name: Install cross
run: cargo install cross --git https://github.com/cross-rs/cross
- name: Build CLI
env:
CROSS_NO_WARNINGS: 0
run: |
export TARGET="${{ matrix.architecture }}-${{ matrix.target-suffix }}"
rustup target add "${TARGET}"
# 'cross' is used to cross-compile for different architectures (see Cross.toml)
cross build --release --target ${TARGET} -p goose-cli
# tar the goose binary as goose-<TARGET>.tar.bz2
cd target/${TARGET}/release
tar -cjf goose-${TARGET}.tar.bz2 goose
echo "ARTIFACT=target/${TARGET}/release/goose-${TARGET}.tar.bz2" >> $GITHUB_ENV
- name: Upload CLI artifact
uses: actions/upload-artifact@v4
with:
name: goose-${{ matrix.architecture }}-${{ matrix.target-suffix }}
path: ${{ env.ARTIFACT }}

198
.github/workflows/bundle-desktop.yml vendored Normal file
View File

@@ -0,0 +1,198 @@
# This is a **reuseable** workflow that bundles the Desktop App for macOS.
# It doesn't get triggered on its own. It gets used in multiple workflows:
# - release.yml
# - canary.yml
# - pr-comment-bundle-desktop.yml
on:
workflow_call:
inputs:
signing:
description: 'Whether to perform signing and notarization'
required: false
default: false
type: boolean
secrets:
CERTIFICATE_OSX_APPLICATION:
description: 'Certificate for macOS application signing'
required: false
CERTIFICATE_PASSWORD:
description: 'Password for the macOS certificate'
required: false
APPLE_ID:
description: 'Apple ID for notarization'
required: false
APPLE_ID_PASSWORD:
description: 'Password for the Apple ID'
required: false
APPLE_TEAM_ID:
description: 'Apple Team ID'
required: false
name: Reusable workflow to bundle desktop app
jobs:
bundle-desktop:
runs-on: macos-latest
name: Bundle Desktop App on macOS
steps:
# Validate Signing Secrets if signing is enabled
- name: Validate Signing Secrets
if: ${{ inputs.signing }}
run: |
if [[ -z "${{ secrets.CERTIFICATE_OSX_APPLICATION }}" ]]; then
echo "Error: CERTIFICATE_OSX_APPLICATION secret is required for signing."
exit 1
fi
if [[ -z "${{ secrets.CERTIFICATE_PASSWORD }}" ]]; then
echo "Error: CERTIFICATE_PASSWORD secret is required for signing."
exit 1
fi
if [[ -z "${{ secrets.APPLE_ID }}" ]]; then
echo "Error: APPLE_ID secret is required for signing."
exit 1
fi
if [[ -z "${{ secrets.APPLE_ID_PASSWORD }}" ]]; then
echo "Error: APPLE_ID_PASSWORD secret is required for signing."
exit 1
fi
if [[ -z "${{ secrets.APPLE_TEAM_ID }}" ]]; then
echo "Error: APPLE_TEAM_ID secret is required for signing."
exit 1
fi
echo "All required signing secrets are present."
- name: Checkout code
uses: actions/checkout@v4
- name: Setup Rust
uses: dtolnay/rust-toolchain@stable
with:
toolchain: stable
- name: Cache Cargo registry
uses: actions/cache@v3
with:
path: ~/.cargo/registry
key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-registry-
- name: Cache Cargo index
uses: actions/cache@v3
with:
path: ~/.cargo/index
key: ${{ runner.os }}-cargo-index
restore-keys: |
${{ runner.os }}-cargo-index
- name: Cache Cargo build
uses: actions/cache@v3
with:
path: target
key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-build-
- name: Build goosed
run: cargo build --release -p goose-server
- name: Copy binary into Electron folder
run: cp target/release/goosed ui/desktop/src/bin/goosed
# Conditional Signing Step - we skip this for faster builds
- name: Add MacOS certs for signing and notarization
if: ${{ inputs.signing }}
run: ./add-macos-cert.sh
working-directory: ui/desktop
env:
CERTIFICATE_OSX_APPLICATION: ${{ secrets.CERTIFICATE_OSX_APPLICATION }}
CERTIFICATE_PASSWORD: ${{ secrets.CERTIFICATE_PASSWORD }}
- name: Set up Node.js
uses: actions/setup-node@v2
with:
node-version: 'lts/*'
- name: Install dependencies
run: npm ci
working-directory: ui/desktop
- name: Make Unsigned App
if: ${{ !inputs.signing }}
run: |
attempt=0
max_attempts=2
until [ $attempt -ge $max_attempts ]; do
npm run bundle:default && break
attempt=$((attempt + 1))
echo "Attempt $attempt failed. Retrying..."
sleep 5
done
if [ $attempt -ge $max_attempts ]; then
echo "Action failed after $max_attempts attempts."
exit 1
fi
working-directory: ui/desktop
- name: Make Signed App
if: ${{ inputs.signing }}
run: |
attempt=0
max_attempts=2
until [ $attempt -ge $max_attempts ]; do
npm run bundle:default && break
attempt=$((attempt + 1))
echo "Attempt $attempt failed. Retrying..."
sleep 5
done
if [ $attempt -ge $max_attempts ]; then
echo "Action failed after $max_attempts attempts."
exit 1
fi
working-directory: ui/desktop
env:
APPLE_ID: ${{ secrets.APPLE_ID }}
APPLE_ID_PASSWORD: ${{ secrets.APPLE_ID_PASSWORD }}
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
- name: Upload Desktop artifact
uses: actions/upload-artifact@v4
with:
name: Goose-darwin-arm64
path: ui/desktop/out/Goose-darwin-arm64/Goose.zip
- name: Quick launch test (macOS)
run: |
# Ensure no quarantine attributes (if needed)
xattr -cr "ui/desktop/out/Goose-darwin-arm64/Goose.app"
echo "Opening Goose.app..."
open -g "ui/desktop/out/Goose-darwin-arm64/Goose.app"
# Give the app a few seconds to start and write logs
sleep 5
# Check if it's running
if pgrep -f "Goose.app/Contents/MacOS/Goose" > /dev/null; then
echo "App appears to be running."
else
echo "App did not stay open. Possible crash or startup error."
exit 1
fi
LOGFILE="$HOME/Library/Application Support/Goose/logs/main.log"
# Print the log and verify "ChatWindow loaded" is in the logs
if [ -f "$LOGFILE" ]; then
echo "===== Log file contents ====="
cat "$LOGFILE"
echo "============================="
if grep -F "ChatWindow loaded" "$LOGFILE"; then
echo "Confirmed: 'ChatWindow loaded' found in logs!"
else
echo "Did not find 'ChatWindow loaded' in logs. Failing..."
exit 1
fi
else
echo "No log file found at $LOGFILE. Exiting with failure."
exit 1
fi
# Kill the app to clean up
pkill -f "Goose.app/Contents/MacOS/Goose"

82
.github/workflows/canary.yml vendored Normal file
View File

@@ -0,0 +1,82 @@
# This workflow is for canary releases, automatically triggered by push to v1.0 branch.
# This workflow is identical to "release.yml" with these exceptions:
# - Triggered by push to v1.0 branch
# - Github Release tagged as "canary"
on:
push:
paths-ignore:
- 'docs/**'
branches:
- main
name: Canary
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
# ------------------------------------
# 1) Build CLI for multiple OS/Arch
# ------------------------------------
build-cli:
uses: ./.github/workflows/build-cli.yml
# ------------------------------------
# 2) Upload Install CLI Script (we only need to do this once)
# ------------------------------------
install-script:
name: Upload Install Script
runs-on: ubuntu-latest
needs: [ build-cli ]
steps:
- uses: actions/checkout@v4
- uses: actions/upload-artifact@v4
with:
name: download_cli.sh
path: download_cli.sh
# ------------------------------------------------------------
# 3) Bundle Desktop App (macOS only) - builds goosed and Electron app
# ------------------------------------------------------------
bundle-desktop:
uses: ./.github/workflows/bundle-desktop.yml
with:
signing: true
secrets:
CERTIFICATE_OSX_APPLICATION: ${{ secrets.CERTIFICATE_OSX_APPLICATION }}
CERTIFICATE_PASSWORD: ${{ secrets.CERTIFICATE_PASSWORD }}
APPLE_ID: ${{ secrets.APPLE_ID }}
APPLE_ID_PASSWORD: ${{ secrets.APPLE_ID_PASSWORD }}
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
# ------------------------------------
# 4) Create/Update GitHub Release
# ------------------------------------
release:
name: Release
runs-on: ubuntu-latest
needs: [ build-cli, install-script, bundle-desktop ]
permissions:
contents: write
steps:
- name: Download all artifacts
uses: actions/download-artifact@v4
with:
merge-multiple: true
# Create/update the canary release
- name: Release canary
uses: ncipollo/release-action@v1
with:
tag: canary
name: Canary
token: ${{ secrets.GITHUB_TOKEN }}
artifacts: |
goose-*.tar.bz2
Goose*.zip
download_cli.sh
allowUpdates: true
omitBody: true
omitPrereleaseDuringUpdate: true

View File

@@ -1,111 +0,0 @@
name: CI
on:
pull_request:
branches:
- main # Trigger CI on PRs to main
push:
branches:
- main # Trigger CI on pushes to main
jobs:
exchange:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v3
with:
version: "0.5.4"
- name: Source Cargo Environment
run: source $HOME/.cargo/env
- name: Ruff
run: |
uvx ruff check packages/exchange
uvx ruff format packages/exchange --check
- name: Run tests
working-directory: ./packages/exchange
run: |
uv run pytest tests -m 'not integration'
goose:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v3
with:
version: "0.5.4"
- name: Source Cargo Environment
run: source $HOME/.cargo/env
- name: Ruff
run: |
uvx ruff check src tests
uvx ruff format src tests --check
- name: Run tests
run: |
uv run pytest tests -m 'not integration'
# This runs integration tests of the OpenAI API, using Ollama to host models.
# This lets us test PRs from forks which can't access secrets like API keys.
ollama:
runs-on: ubuntu-latest
strategy:
matrix:
python-version:
# Only test the lastest python version.
- "3.12"
ollama-model:
# For quicker CI, use a smaller, tool-capable model than the default.
- "qwen2.5:0.5b"
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v3
with:
version: "0.5.4"
- name: Source Cargo Environment
run: source $HOME/.cargo/env
- name: Set up Python
run: uv python install ${{ matrix.python-version }}
- name: Install Ollama
run: curl -fsSL https://ollama.com/install.sh | sh
- name: Start Ollama
run: |
# Run the background, in a way that survives to the next step
nohup ollama serve > ollama.log 2>&1 &
# Block using the ready endpoint
time curl --retry 5 --retry-connrefused --retry-delay 1 -sf http://localhost:11434
# Tests use OpenAI which does not have a mechanism to pull models. Run a
# simple prompt to (pull and) test the model first.
- name: Test Ollama model
run: ollama run $OLLAMA_MODEL hello || cat ollama.log
env:
OLLAMA_MODEL: ${{ matrix.ollama-model }}
- name: Run Ollama tests
run: uv run pytest tests -m integration -k ollama
working-directory: ./packages/exchange
env:
OLLAMA_MODEL: ${{ matrix.ollama-model }}

107
.github/workflows/ci.yml vendored Normal file
View File

@@ -0,0 +1,107 @@
on:
push:
paths-ignore:
- 'docs/**'
branches:
- main
pull_request:
paths-ignore:
- 'docs/**'
branches:
- main
workflow_dispatch:
name: CI
jobs:
rust-format:
name: Check Rust Code Format
runs-on: ubuntu-latest
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Rust
uses: dtolnay/rust-toolchain@stable
with:
toolchain: stable
- name: Run cargo fmt
run: cargo fmt --check
rust-build-and-test:
name: Build and Test Rust Project
runs-on: ubuntu-latest
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Install Dependencies
run: |
sudo apt update -y
sudo apt install -y libdbus-1-dev gnome-keyring libxcb1-dev
- name: Setup Rust
uses: dtolnay/rust-toolchain@stable
with:
toolchain: stable
- name: Cache Cargo Registry
uses: actions/cache@v3
with:
path: ~/.cargo/registry
key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-registry-
- name: Cache Cargo Index
uses: actions/cache@v3
with:
path: ~/.cargo/index
key: ${{ runner.os }}-cargo-index
restore-keys: |
${{ runner.os }}-cargo-index
- name: Cache Cargo Build
uses: actions/cache@v3
with:
path: target
key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-build-
- name: Build and Test
run: |
gnome-keyring-daemon --components=secrets --daemonize --unlock <<< 'foobar'
cargo test
working-directory: crates
- name: Lint
run: cargo clippy -- -D warnings
desktop-lint:
name: Lint Electron Desktop App
runs-on: macos-latest
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v2
with:
node-version: 'lts/*'
- name: Install Dependencies
run: npm ci
working-directory: ui/desktop
- name: Run Lint
run: npm run lint:check
working-directory: ui/desktop
# Faster Desktop App build for PRs only
bundle-desktop-unsigned:
uses: ./.github/workflows/bundle-desktop.yml
if: github.event_name == 'pull_request'
with:
signing: false

View File

@@ -0,0 +1,47 @@
name: Deploy v1 Docs & Extensions # (/documentation and /extensions-site)
on:
push:
branches:
- main
jobs:
deploy:
runs-on: ubuntu-latest
steps:
- name: Checkout the branch
uses: actions/checkout@v3
- name: Setup Node.js
uses: actions/setup-node@v3
with:
node-version: 20
- name: Install dependencies and build docs
working-directory: ./documentation
run: |
npm install
npm run build
- name: Install dependencies and build extensions-site
working-directory: ./extensions-site
env:
VITE_BASENAME: "/goose/v1/extensions/" # Set the base URL here for the extensions site
run: |
npm install
npm run build
- name: Combine builds into one directory
run: |
mkdir combined-build
cp -r documentation/build/* combined-build/
mkdir -p combined-build/extensions
cp -r extensions-site/build/client/* combined-build/extensions/
- name: Deploy to gh-pages
uses: peaceiris/actions-gh-pages@v3
with:
github_token: ${{ secrets.GITHUB_TOKEN }}
publish_dir: combined-build
destination_dir: v1 # Deploy the site to the 'v1' subfolder

View File

@@ -1,40 +0,0 @@
name: Deploy MkDocs
on:
push:
branches:
- main # Trigger deployment on pushes to main
paths:
- 'docs/**'
- 'mkdocs.yml'
- '.github/workflows/deploy_docs.yaml'
pull_request:
branches:
- main
paths:
- 'docs/**'
- 'mkdocs.yml'
- '.github/workflows/deploy_docs.yaml'
jobs:
deploy:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Install UV
uses: astral-sh/setup-uv@v3
- name: Create UV virtual environment
run: uv venv
- name: Install dependencies
run: uv pip install "mkdocs-material[imaging]" Pillow cairosvg
- name: Build the documentation
run: uv run mkdocs gh-deploy --force

View File

@@ -1,53 +0,0 @@
---
name: License Check
on:
pull_request: # Trigger license check on any PRs
paths:
- '**/pyproject.toml'
- '.github/workflows/license-check.yml'
- '.github/workflows/scripts/check_licenses.py'
push: # Trigger license check on pushes to main
branches:
- main
paths: # TODO: can't DRY unless https://github.com/actions/runner/issues/1182
- '**/pyproject.toml'
- '.github/workflows/license-check.yml'
- '.github/workflows/scripts/check_licenses.py'
jobs:
check-licenses:
name: Check Package Licenses
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install tomli requests urllib3
- name: Check licenses
run: |
python .github/workflows/scripts/check_licenses.py \
pyproject.toml || exit_code=$?
if [ "${exit_code:-0}" -ne 0 ]; then
echo "::error::Found packages with disallowed licenses"
exit 1
fi
- name: Check Exchange licenses
run: |
python .github/workflows/scripts/check_licenses.py \
packages/exchange/pyproject.toml || exit_code=$?
if [ "${exit_code:-0}" -ne 0 ]; then
echo "::error::Found packages with disallowed licenses in exchange"
exit 1
fi

View File

@@ -0,0 +1,76 @@
# This workflow is triggered by a comment on an issue or PR with the text ".bundle"
# It bundles the Desktop App, then creates a PR comment with a link to download the app.
# IMPORTANT: issue_comment workflows only use files found on "main" branch to run.
# Do NOT allow on: pull_request since that allows a user to alter a file in a PR and exfil secrets for example.
# Only using issue_comment is the suggested workflow type. Comments on pull requests, and issues will trigger the issue_comment workflow event.
on:
issue_comment:
types: [created]
# permissions needed for reacting to IssueOps commands on PRs
permissions:
pull-requests: write
checks: read
# issues: write
name: Workflow to Bundle Desktop App
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
trigger-on-command:
name: Trigger on ".bundle" PR comment
runs-on: ubuntu-latest
steps:
- uses: github/command@v1.3.0
id: command
with:
command: ".bundle"
reaction: "eyes"
allowed_contexts: pull_request
bundle-desktop:
# Only run this if ".bundle" command is detected.
if: ${{ steps.command.outputs.continue == 'true' }}
uses: ./.github/workflows/bundle-desktop.yml
with:
signing: true
secrets:
CERTIFICATE_OSX_APPLICATION: ${{ secrets.CERTIFICATE_OSX_APPLICATION }}
CERTIFICATE_PASSWORD: ${{ secrets.CERTIFICATE_PASSWORD }}
APPLE_ID: ${{ secrets.APPLE_ID }}
APPLE_ID_PASSWORD: ${{ secrets.APPLE_ID_PASSWORD }}
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
pr-comment:
name: PR Comment with Desktop App
runs-on: ubuntu-latest
needs: [ bundle-desktop ]
permissions:
pull-requests: write
steps:
- name: Download all artifacts
uses: actions/download-artifact@v4
with:
merge-multiple: true
- name: Comment on PR with download link
uses: peter-evans/create-or-update-comment@v3
with:
comment-id: ${{ steps.command.outputs.comment_id }}
issue-number: ${{ github.event.pull_request.number }}
body: |
### Desktop App for this PR
The following build is available for testing:
- [📱 macOS Desktop App (arm64, signed)](https://nightly.link/${{ github.repository }}/actions/runs/${{ github.run_id }}/Goose-darwin-arm64.zip)
After downloading, unzip the file and drag the Goose.app to your Applications folder. The app is signed and notarized for macOS.
This link is provided by nightly.link and will work even if you're not logged into GitHub.
edit-mode: replace

View File

@@ -1,50 +0,0 @@
name: Publish
# A release on goose will also publish exchange, if it has updated
# This means in some cases we may need to make a bump in goose without other changes to release exchange
on:
release:
types: [published]
jobs:
publish:
permissions:
id-token: write
contents: read
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Get current version from pyproject.toml
id: get_version
run: |
echo "VERSION=$(grep -m 1 'version =' "pyproject.toml" | awk -F'"' '{print $2}')" >> $GITHUB_ENV
- name: Extract tag version
id: extract_tag
run: |
TAG_VERSION=$(echo "${{ github.event.release.tag_name }}" | sed -E 's/v(.*)/\1/')
echo "TAG_VERSION=$TAG_VERSION" >> $GITHUB_ENV
- name: Check if tag matches version from pyproject.toml
id: check_tag
run: |
if [ "${{ env.TAG_VERSION }}" != "${{ env.VERSION }}" ]; then
echo "::error::Tag version (${{ env.TAG_VERSION }}) does not match version in pyproject.toml (${{ env.VERSION }})."
exit 1
fi
- name: Install the latest version of uv
uses: astral-sh/setup-uv@v1
with:
version: "latest"
- name: Build Package
run: |
uv build -o dist --package goose-ai
uv build -o dist --package ai-exchange
- name: Publish package to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
skip-existing: true

View File

@@ -1,48 +0,0 @@
name: 'Lint PR'
on:
pull_request_target:
types:
- opened
- edited
- synchronize
- reopened
permissions:
pull-requests: write
jobs:
main:
name: Validate PR title
runs-on: ubuntu-latest
steps:
- uses: amannn/action-semantic-pull-request@v5
id: lint_pr_title
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
requireScope: false
- uses: marocchino/sticky-pull-request-comment@v2
# When the previous steps fails, the workflow would stop. By adding this
# condition you can continue the execution with the populated error message.
if: always() && (steps.lint_pr_title.outputs.error_message != null)
with:
header: pr-title-lint-error
message: |
Hey there and thank you for opening this pull request! 👋🏼
We require pull request titles to follow the [Conventional Commits specification](https://gist.github.com/Zekfad/f51cb06ac76e2457f11c80ed705c95a3#file-conventional-commits-md) and it looks like your proposed title needs to be adjusted.
Details:
```
${{ steps.lint_pr_title.outputs.error_message }}
```
# Delete a previous comment when the issue has been resolved
- if: ${{ steps.lint_pr_title.outputs.error_message == null }}
uses: marocchino/sticky-pull-request-comment@v2
with:
header: pr-title-lint-error
delete: true

View File

@@ -1,41 +0,0 @@
name: Release Monitor
on:
release:
types: [published]
workflow_dispatch: # Add this line to enable manual triggering
jobs:
build:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v2
- name: Set up Python
uses: actions/setup-python@v2
with:
python-version: '3.x'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install pipx
pipx install goose-ai
- name: Check Goose AI Version
run: goose version
- name: Create Issue on Failure
if: failure()
uses: actions/github-script@v3
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const { owner, repo } = context.repo;
await github.issues.create({
owner: owner,
repo: repo,
title: 'Release Build Failed',
body: `The release for version ${{ github.event.release.tag_name }} failed to run. Please investigate the issue.`
});

91
.github/workflows/release.yml vendored Normal file
View File

@@ -0,0 +1,91 @@
# This workflow is main release, needs to be manually tagged & pushed.
on:
push:
tags:
- "v1.*"
name: Release
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
# ------------------------------------
# 1) Build CLI for multiple OS/Arch
# ------------------------------------
build-cli:
uses: ./.github/workflows/build-cli.yml
# ------------------------------------
# 2) Upload Install CLI Script (we only need to do this once)
# ------------------------------------
install-script:
name: Upload Install Script
runs-on: ubuntu-latest
needs: [ build-cli ]
steps:
- uses: actions/checkout@v4
- uses: actions/upload-artifact@v4
with:
name: download_cli.sh
path: download_cli.sh
# ------------------------------------------------------------
# 3) Bundle Desktop App (macOS only) - builds goosed and Electron app
# ------------------------------------------------------------
bundle-desktop:
uses: ./.github/workflows/bundle-desktop.yml
with:
signing: true
secrets:
CERTIFICATE_OSX_APPLICATION: ${{ secrets.CERTIFICATE_OSX_APPLICATION }}
CERTIFICATE_PASSWORD: ${{ secrets.CERTIFICATE_PASSWORD }}
APPLE_ID: ${{ secrets.APPLE_ID }}
APPLE_ID_PASSWORD: ${{ secrets.APPLE_ID_PASSWORD }}
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
# ------------------------------------
# 4) Create/Update GitHub Release
# ------------------------------------
release:
name: Release
runs-on: ubuntu-latest
needs: [ build-cli, install-script, bundle-desktop ]
permissions:
contents: write
steps:
- name: Download all artifacts
uses: actions/download-artifact@v4
with:
merge-multiple: true
# Create/update the versioned release
- name: Release versioned
uses: ncipollo/release-action@v1
with:
token: ${{ secrets.GITHUB_TOKEN }}
# This pattern will match both goose tar.bz2 artifacts and the Goose.zip
artifacts: |
goose-*.tar.bz2
Goose*.zip
download_cli.sh
allowUpdates: true
omitBody: true
omitPrereleaseDuringUpdate: true
# Create/update the stable release
- name: Release stable
uses: ncipollo/release-action@v1
with:
tag: stable
name: Stable
token: ${{ secrets.GITHUB_TOKEN }}
artifacts: |
goose-*.tar.bz2
Goose*.zip
download_cli.sh
allowUpdates: true
omitBody: true
omitPrereleaseDuringUpdate: true

View File

@@ -1,320 +0,0 @@
#!/usr/bin/env python3
import argparse
import os
import sys
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
import requests
import tomli
import urllib3
class Color(str, Enum):
"""ANSI color codes with fallback for non-color terminals"""
@staticmethod
def supports_color() -> bool:
"""Check if the terminal supports color output."""
if not hasattr(sys.stdout, "isatty"):
return False
if not sys.stdout.isatty():
return False
if "NO_COLOR" in os.environ:
return False
term = os.environ.get("TERM", "")
if term == "dumb":
return False
return True
has_color = supports_color()
RED = "\033[91m" if has_color else ""
GREEN = "\033[92m" if has_color else ""
RESET = "\033[0m" if has_color else ""
BOLD = "\033[1m" if has_color else ""
@dataclass(frozen=True)
class LicenseConfig:
allowed_licenses: frozenset[str] = frozenset(
{
"MIT",
"BSD-3-Clause",
"Apache-2.0",
"Apache License 2",
"Apache Software License",
"Python Software Foundation License",
"BSD License",
"ISC",
}
)
exceptions: frozenset[str] = frozenset(
{
"ai-exchange",
"tiktoken",
}
)
@dataclass(frozen=True)
class LicenseInfo:
license: str | None
allowed: bool = False
def __str__(self) -> str:
status = "" if self.allowed else ""
color = Color.GREEN if self.allowed else Color.RED
return f"{color}{status}{Color.RESET} {self.license}"
class LicenseChecker:
def __init__(self, config: LicenseConfig = LicenseConfig()) -> None:
self.config = config
self.session = self._setup_session()
def _setup_session(self) -> requests.Session:
session = requests.Session()
session.verify = True
max_retries = urllib3.util.Retry(
total=3,
backoff_factor=0.5,
status_forcelist=[
500,
502,
503,
504,
],
)
adapter = requests.adapters.HTTPAdapter(max_retries=max_retries)
session.mount("https://", adapter)
return session
def normalize_license(self, license_str: str | None) -> str | None:
"""
Normalize license string for comparison.
This method takes a license string and normalizes it by:
1. Converting to uppercase
2. Removing 'LICENSE' or 'LICENCE' suffixes
3. Stripping whitespace
4. Replacing common variations with standardized forms
Args:
license_str (str | None): The original license string to normalize.
Returns:
str | None: The normalized license string, or None if the input was None.
"""
if not license_str:
return None
# fmt: off
normalized = (
license_str.upper()
.replace(" LICENSE", "")
.replace(" LICENCE", "")
.strip()
)
# fmt: on
replacements = {
"APACHE 2.0": "APACHE-2.0",
"APACHE SOFTWARE LICENSE": "APACHE-2.0",
"BSD": "BSD-3-CLAUSE",
"MIT LICENSE": "MIT",
"PYTHON SOFTWARE FOUNDATION": "PSF",
}
return replacements.get(normalized, normalized)
def get_package_license(self, package_name: str) -> str | None:
"""Fetch license information from PyPI.
Args:
package_name (str): The name of the package to fetch the license for.
Returns:
str | None: The license of the package, or None if not found.
"""
try:
response = self.session.get(
f"https://pypi.org/pypi/{package_name}/json",
timeout=10,
)
response.raise_for_status()
data = response.json()
# fmt: off
license_info = (
data["info"].get("license") or
data["info"].get("classifiers", [])
)
# fmt: on
if isinstance(license_info, list):
for classifier in license_info:
if classifier.startswith("License :: "):
parts = classifier.split(" :: ")
return parts[-1]
return license_info if isinstance(license_info, str) else None
except requests.exceptions.SSLError as e:
print(f"SSL Error fetching license for {package_name}: {e}", file=sys.stderr)
except Exception as e:
print(f"Warning: Could not fetch license for {package_name}: {e}", file=sys.stderr)
return None
def extract_dependencies(self, toml_file: Path) -> list[str]:
"""Extract all dependencies from a TOML file."""
with open(toml_file, "rb") as f:
data = tomli.load(f)
dependencies = []
# Get direct dependencies
project_deps = data.get("project", {}).get("dependencies", [])
dependencies.extend(self._parse_dependency_strings(project_deps))
# Get dev dependencies
tool_deps = data.get("tool", {}).get("uv", {}).get("dev-dependencies", [])
dependencies.extend(self._parse_dependency_strings(tool_deps))
return list(set(dependencies))
def _parse_dependency_strings(self, deps: list[str]) -> list[str]:
"""
Parse dependency strings to extract package names.
Args:
deps (list[str]): A list of dependency strings to parse.
Returns:
list[str]: A list of extracted package names.
"""
packages = []
for dep in deps:
if "workspace = true" in dep:
continue
# fmt: off
# Handle basic package specifiers
package = (
dep.split(">=")[0]
.split("==")[0]
.split("<")[0]
.split(">")[0]
.strip()
)
package = package.split("{")[0].strip()
# fmt: on
if package:
packages.append(package)
return packages
def check_licenses(self, toml_file: Path) -> dict[str, LicenseInfo]:
"""
Check licenses for all dependencies in the TOML file.
Args:
toml_file (Path): The path to the TOML file containing the dependencies.
Returns:
dict[str, LicenseInfo]: A dictionary where the keys are package names and the values are LicenseInfo objects
containing the license information and whether it's allowed."""
dependencies = self.extract_dependencies(toml_file)
results: dict[str, LicenseInfo] = {}
checked: set[str] = set()
for package in dependencies:
if package in checked:
continue
checked.add(package)
results[package] = self._check_package(package)
return results
def _check_package(self, package: str) -> LicenseInfo:
"""
Check license for a single package.
Args:
package (str): The name of the package to check.
Returns:
LicenseInfo: A LicenseInfo object containing the license
information and whether it's allowed.
"""
if package in self.config.exceptions:
return LicenseInfo("Approved Exception", True)
license_info = self.get_package_license(package)
normalized_license = self.normalize_license(license_info)
allowed = False
# fmt: off
if normalized_license:
allowed = normalized_license in {
self.normalize_license(x)
for x in self.config.allowed_licenses
}
# fmt: on
return LicenseInfo(license_info, allowed)
def main() -> None:
parser = argparse.ArgumentParser(description="Check package licenses in TOML files")
parser.add_argument("toml_files", type=Path, nargs="*", help="Paths to TOML files")
parser.add_argument("--supported-licenses", action="store_true", help="Print supported licenses")
checker = LicenseChecker()
all_results: dict[str, LicenseInfo] = {}
args = parser.parse_args()
if args.supported_licenses:
for license in sorted(checker.config.allowed_licenses, key=str.casefold):
print(f" - {license}")
sys.exit(0)
if not args.toml_files:
print("Error: No TOML files specified", file=sys.stderr)
parser.print_help()
sys.exit(1)
for toml_file in args.toml_files:
results = checker.check_licenses(toml_file)
for package, info in results.items():
if package in all_results and all_results[package] != info:
print(f"Warning: Package {package} has conflicting license info:", file=sys.stderr)
print(f" {toml_file}: {info}", file=sys.stderr)
print(f" Previous: {all_results[package]}", file=sys.stderr)
all_results[package] = info
max_package_length = max(len(package) for package in all_results.keys())
any_disallowed = False
for package, info in sorted(all_results.items()):
if Color.has_color:
package_name = f"{Color.BOLD}{package}{Color.RESET}"
padding = len(Color.BOLD) + len(Color.RESET)
else:
package_name = package
padding = 0
print(f"{package_name:<{max_package_length + padding}} {info}")
if not info.allowed:
any_disallowed = True
sys.exit(1 if any_disallowed else 0)
if __name__ == "__main__":
main()

View File

@@ -1,248 +0,0 @@
from pathlib import Path
from typing import Optional
from unittest.mock import Mock, patch
import pytest
import tomli
from check_licenses import Color, LicenseChecker, LicenseConfig, LicenseInfo, main
@pytest.fixture
def checker() -> LicenseChecker:
return LicenseChecker()
@pytest.fixture
def mock_pypi_response() -> Mock:
response = Mock()
response.status_code = 200
response.raise_for_status = Mock()
response.ok = True
response.json.return_value = {
"info": {
"license": "Apache-2.0",
}
}
return response
@pytest.fixture
def mock_toml_content() -> str:
return """
[project]
dependencies = [
"requests>=2.28.0",
"tomli==2.0.1",
"urllib3<2.0.0",
"package-with-workspace{workspace = true}",
]
[tool.uv]
dev-dependencies = [
"pytest>=7.0.0",
"black==23.3.0",
]
"""
@pytest.fixture
def mock_toml_files(tmp_path: Path) -> list[Path]:
"""Create mock TOML files with different dependencies."""
file1 = tmp_path / "pyproject1.toml"
file1.write_text("""
[project]
dependencies = [
"requests>=2.28.0",
"tomli==2.0.1",
]
""")
file2 = tmp_path / "pyproject2.toml"
file2.write_text("""
[project]
dependencies = [
"urllib3<2.0.0",
"requests>=2.27.0", # Different version but same package
]
""")
return [file1, file2]
def test_normalize_license_variations(checker: LicenseChecker) -> None:
assert checker.normalize_license("MIT License") == "MIT"
assert checker.normalize_license("Apache 2.0") == "APACHE-2.0"
assert checker.normalize_license("BSD") == "BSD-3-CLAUSE"
assert checker.normalize_license(None) is None
assert checker.normalize_license("") is None
assert checker.normalize_license("MIT License") == "MIT"
assert checker.normalize_license("Apache 2.0") == "APACHE-2.0"
assert checker.normalize_license("BSD") == "BSD-3-CLAUSE"
assert checker.normalize_license(None) is None
assert checker.normalize_license("") is None
@patch.object(LicenseChecker, "get_package_license")
def test_package_license_verification(mock_get_license: Mock, checker: LicenseChecker) -> None:
def check_package_license(
package: str, license: Optional[str], expected_allowed: bool, expected_license: str
) -> None:
if license:
mock_get_license.return_value = license
result = checker._check_package(package=package)
assert result.allowed is expected_allowed
assert result.license == expected_license
check_package_license(
package="tiktoken",
license=None,
expected_allowed=True,
expected_license="Approved Exception",
)
check_package_license(
package="requests",
license="Apache-2.0",
expected_allowed=True,
expected_license="Apache-2.0",
)
check_package_license(
package="gpl-package",
license="GPL",
expected_allowed=False,
expected_license="GPL",
)
def test_color_support_with_environment_variables(monkeypatch: pytest.MonkeyPatch) -> None:
def verify_color_support_disabled(*, env_var: str, value: str) -> None:
monkeypatch.setenv(env_var, value)
assert not Color.supports_color()
monkeypatch.undo()
verify_color_support_disabled(env_var="NO_COLOR", value="1")
verify_color_support_disabled(env_var="TERM", value="dumb")
@patch("tomli.load")
@patch("builtins.open")
def test_extract_dependencies(
mock_open: Mock, mock_tomli_load: Mock, checker: LicenseChecker, mock_toml_content: str
) -> None:
mock_tomli_load.return_value = tomli.loads(mock_toml_content)
mock_file = Mock()
mock_open.return_value.__enter__.return_value = mock_file
dependencies = checker.extract_dependencies(Path("mock_pyproject.toml"))
expected = ["requests", "tomli", "urllib3", "pytest", "black"]
assert sorted(dependencies) == sorted(expected)
@patch("requests.Session")
def test_get_package_license(mock_session: Mock, checker: LicenseChecker, mock_pypi_response: Mock) -> None:
mock_session.return_value.get.return_value = mock_pypi_response
assert checker.get_package_license("requests") == "Apache-2.0"
# test exception handling
mock_session.return_value.get.side_effect = Exception("error")
assert checker.get_package_license("nonexistent-package") is None
def test_license_info_string_representation() -> None:
def assert_license_info_str(info: LicenseInfo, no_color_expected: str, color_expected: str) -> None:
expected = color_expected if Color.has_color else no_color_expected
assert str(info) == expected
assert_license_info_str(LicenseInfo("MIT", True), "✓ MIT", f"{Color.GREEN}{Color.RESET} MIT")
assert_license_info_str(LicenseInfo("GPL", False), "✗ GPL", f"{Color.RED}{Color.RESET} GPL")
def test_custom_license_config() -> None:
custom_config = LicenseConfig(
allowed_licenses=frozenset({"MIT", "Apache-2.0"}), exceptions=frozenset({"special-package"})
)
checker = LicenseChecker(config=custom_config)
assert "special-package" in checker.config.exceptions
assert len(checker.config.allowed_licenses) == 2
def test_dependency_parsing_scenarios() -> None:
"""Test various dependency parsing scenarios."""
def parse_dependencies(toml_string: str) -> list[str]:
checker = LicenseChecker()
return checker._parse_dependency_strings(tomli.loads(toml_string)["project"]["dependencies"])
# basic version specifiers
parsed = parse_dependencies("""
[project]
dependencies = [
"requests>=2.28.0",
"tomli==2.0.1",
"urllib3<2.0.0",
]
""")
assert sorted(parsed) == sorted(["requests", "tomli", "urllib3"])
# workspace dependencies
parsed = parse_dependencies("""
[project]
dependencies = [
"package-with-workspace{workspace = true}",
]
""")
assert parsed == []
# mixed dependencies
parsed = parse_dependencies("""
[project]
dependencies = [
"requests>=2.28.0",
"package-with-workspace{workspace = true}",
"urllib3<2.0.0",
]
""")
assert sorted(parsed) == sorted(["requests", "urllib3"])
# multiple version constraints
parsed = parse_dependencies("""
[project]
dependencies = [
"urllib3>=1.25.4,<2.0.0",
"requests>=2.28.0,<3.0.0",
]
""")
assert sorted(parsed) == sorted(["urllib3", "requests"])
# empty dependencies
parsed = parse_dependencies("""
[project]
dependencies = []
""")
assert parsed == []
@patch.object(LicenseChecker, "get_package_license")
def test_multiple_toml_files(
mock_get_license: Mock,
mock_toml_files: list[Path],
capsys: pytest.CaptureFixture,
) -> None:
def get_license(package: str) -> Optional[str]:
licenses = {"requests": "Apache-2.0", "tomli": "MIT", "urllib3": "MIT"}
return licenses.get(package)
mock_get_license.side_effect = get_license
with patch("sys.argv", ["check_licenses.py"] + [str(f) for f in mock_toml_files]):
try:
main()
except SystemExit as e:
assert e.code == 0
captured = capsys.readouterr()
assert "requests" in captured.out
assert "tomli" in captured.out
assert "urllib3" in captured.out
assert "" not in captured.out

View File

@@ -1,12 +0,0 @@
{
"pull_request": {
"head": {
"ref": "test-branch"
},
"base": {
"ref": "main"
},
"number": 123,
"title": "test: Update dependency licenses"
}
}

153
.gitignore vendored
View File

@@ -1,135 +1,42 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
# Translations
*.mo
*.pot
# Django stuff:
run_cli.sh
tokenizer_files/
.DS_Store
.idea
*.log
local_settings.py
tmp/
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
# Generated by Cargo
# will have compiled files and executables
debug/
target/
# Jupyter Notebook
.ipynb_checkpoints
# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
Cargo.lock
# IPython
profile_default/
ipython_config.py
# These are backup files generated by rustfmt
**/*.rs.bk
# pyenv
.python-version
# MSVC Windows builds of rustc generate these, which store debugging information
*.pdb
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# UI
./ui/desktop/node_modules
./ui/desktop/out
# Hermit
/.hermit/
/bin/
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
debug_*.txt
# Celery stuff
celerybeat-schedule
celerybeat.pid
# Docs
# Dependencies
/node_modules
# SageMath parsed files
*.sage.py
# Production
/build
# Environments
.env
.env.*
.venv
# exception for local langfuse init vars
!**/packages/exchange/.env.langfuse.local
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
# VSCode
.vscode/*
!.vscode/settings.json
!.vscode/extensions.json
# Autogenerated docs files
docs/docs/reference
# uv lock file
uv.lock
# local files
.DS_Store
# Generated files
.docusaurus
.cache-loader

View File

@@ -1,13 +1,27 @@
This is a python CLI app that uses UV. Read CONTRIBUTING.md for information on how to build and test it as needed.
You are an expert programmer in Rust teaching who is teaching another developer who is learning Rust.
The students are familiar with programming in languages such as Python (advanced), Java (novice) and C (novice) so
when possible use analogies from those languages.
Some key concepts are that it is run as a command line interface, dependes on the "ai-exchange" package (which is in packages/exchange in this repo), and has the concept of toolkits which are ways that its behavior can be extended. Look in src/goose and tests.
Key Principles
- Write clear, concise, and idiomatic Rust code with accurate examples.
- Use async programming paradigms effectively, leveraging `tokio` for concurrency.
- Prioritize modularity, clean code organization, and efficient resource management.
- Use expressive variable names that convey intent (e.g., `is_ready`, `has_data`).
- Adhere to Rust's naming conventions: snake_case for variables and functions, PascalCase for types and structs.
- Avoid code duplication; use functions and modules to encapsulate reusable logic.
- Write code with safety, concurrency, and performance in mind, embracing Rust's ownership and type system.
Assume the user has UV installed and ensure UV is used to run any python related commands.
Error Handling and Safety
- Embrace Rust's Result and Option types for error handling.
- Use `?` operator to propagate errors in async functions.
- Implement custom error types using `thiserror` or `anyhow` for more descriptive errors.
- Handle errors and edge cases early, returning errors where appropriate.
- Use `.await` responsibly, ensuring safe points for context switching.
To run tests:
Key Conventions
1. Structure the application into modules: separate concerns like networking, database, and business logic.
2. Use environment variables for configuration management (e.g., `dotenv` crate).
3. Ensure code is well-documented with inline comments and Rustdoc.
4. Do not use the older style of "MOD/mod.rs" for separing modules and instead use the "MOD.rs" filename convention.
```sh
uv sync && uv run pytest tests -m 'not integration'
```
ideally after each change
Refer to "The Rust Programming Language" book (2024 version) and "Command line apps in Rust" documentation for in-depth information on best practices, and advanced features.

4
.husky/pre-commit Executable file
View File

@@ -0,0 +1,4 @@
#!/usr/bin/env sh
. "$(dirname -- "$0")/_/husky.sh"
cd ui/desktop && npx lint-staged

View File

@@ -1,6 +0,0 @@
lint.select = ["E", "W", "F", "N", "ANN"]
lint.ignore = ["ANN101"]
exclude = [
"docs",
]
line-length = 120

View File

@@ -1,7 +0,0 @@
{
"recommendations": [
"ms-python.debugpy",
"ms-python.python",
"charliermarsh.ruff"
]
}

15
.vscode/settings.json vendored
View File

@@ -1,15 +0,0 @@
{
"[python]": {
"editor.formatOnSave": true,
"editor.defaultFormatter": "charliermarsh.ruff",
"editor.codeActionsOnSave": {
"source.fixAll": "explicit",
"source.organizeImports": "explicit"
}
},
"python.testing.pytestArgs": [
"tests"
],
"python.testing.unittestEnabled": false,
"python.testing.pytestEnabled": true
}

View File

@@ -1,57 +0,0 @@
# Acceptable Usage Policy (AUP)
---
We want everyone to use Goose safely and responsibly. You agree you will not use, or allow others to use, Goose to:
### 1. Violate the law or others rights, including use of Goose to:
a. Engage in, promote, generate, contribute to, encourage, plan, incite, or further illegal or unlawful activity or content, such as:
- Violence or terrorism
- Exploitation or harm to children, including the solicitation, creation, acquisition, or dissemination of child exploitative content or failure to report Child Sexual Abuse Material
- Human trafficking, exploitation, and sexual violence
- The illegal distribution of information or materials to minors
- Sexual solicitation
- Any other criminal activity
b. Engage in, promote, incite, or facilitate the harassment, abuse, threatening, or bullying of individuals or groups of individuals.
c. Engage in, promote, incite, or facilitate discrimination or other unlawful or harmful conduct in the provision of employment, employment benefits, credit, housing, other economic benefits, or other essential goods and services.
d. Violate any person's privacy rights as defined by applicable privacy laws, such as sharing personal information without consent, accessing private data unlawfully, or violating any relevant privacy regulations.
e. Misuse, collect, solicit, or gain access to private information without permission, such as non-public contact details, health data, biometric or neural data (including facial recognition), or confidential or proprietary data.
f. Infringe or misappropriate any third-party rights, including intellectual property rights.
g. Create, generate, or facilitate the creation of malicious code, malware, computer viruses, or do anything else in an intentional or malicious way without third-party consent that could disable, overburden, interfere with, or impair the proper working, integrity, operation, or appearance of a website or computer system.
### 2. Engage in, promote, incite, facilitate, or assist in the planning or development of activities that present a risk of death or bodily harm to individuals, including use of Goose related to the following:
a. Guns and illegal weapons (including weapon development).
b. Illegal drugs and regulated/controlled substances.
c. Any content intended to incite or promote violence, abuse, or any infliction of bodily harm to an individual.
### 3. Intentionally deceive or mislead others or engage in other abusive or fraudulent activities, including use of Goose related to the following:
a. Generating, promoting, or furthering fraud or fraudulent activities, scams, phishing, or malware.
b. Generating, promoting, or furthering defamatory content, including the creation of defamatory statements, images, or other content.
c. Impersonating another individual without consent, authorization, or legal right.
d. Generating or facilitating false online engagement, including fake reviews and other means of fake online engagement.
e. Plagiarism or other forms of academic dishonesty.
f. Compromising security systems or gaining unauthorized access to computer systems or networks without authorization or permission from the affected party, such as spoofing.
---
You also agree that you will not use, or allow others to use, Goose in violation of any agreements that you have with, or commitments you have made to, third parties. This may include any licenses, agreements, and other terms that apply to the models you use with Goose.
---
**Please report any violation of this Policy through the following means:** [open-source-governance@block.xyz](mailto:open-source-governance@block.xyz)

View File

@@ -1,12 +1,12 @@
# Architecture
## The System
## The Extension System
Goose extends the capabilities of high-performing LLMs through a small collection of tools.
This lets you instruct goose, currently via a CLI interface, to automatically solve problems
on your behalf. It attempts to not just tell you how you can do something, but to actually do it for you.
The primary mode of goose (the "developer" toolkit) has access to tools to
The primary mode of goose (the "developer" extension) has access to tools to
- maintain a plan
- run shell commands
@@ -40,7 +40,7 @@ that you should be able to observe by using it.
## Implementation
The core execution logic for generation and tool calling is handled by [exchange][exchange].
It hooks python functions into the model tool use loop, while defining very careful error handling
It hooks rust functions into the model tool use loop, while defining very careful error handling
so any failures in tools are surfaced to the model.
Once we've created an *exchange* object, running the process is effectively just calling
@@ -50,7 +50,7 @@ Once we've created an *exchange* object, running the process is effectively just
Goose builds that exchange:
- allows users to configure a profile to customize capabilities
- provides a pluggable system for adding tools and prompts
- provides a pluggable extension system for adding tools and prompts
- sets up the tools to interact with state
We expect that goose will have multiple UXs over time, and be run in different
@@ -60,28 +60,29 @@ notifications on stdout).
Goose then constructs the exchange for the UX, the UX only interacts with that exchange.
```
def build_exchange(profile: Profile, notifier: Notifier) -> Exchange:
```rust
fn build_exchange(profile: Profile, notifier: Notifier) -> Exchange {
...
}
```
But to setup a configurable system, Goose uses `Toolkit`s:
But to setup a configurable system, Goose uses `Extensions`:
```
(Profile, Notifier) -> [Toolkits] -> Exchange
(Profile, Notifier) -> [Extensions] -> Exchange
```
## Profile
A profile specifies some basic configuration in Goose, such as which models it should use, as well
as which toolkits it should include.
as which extensions it should include.
```yaml
processor: openai:gpt-4o
accelerator: openai:gpt-4o-mini
processor: openai:gpt-4
accelerator: openai:gpt-4-turbo
moderator: passive
toolkits:
- assistant
extensions:
- developer
- calendar
- contacts
- name: scheduling
@@ -93,16 +94,14 @@ toolkits:
## Notifier
The notifier is a concrete implementation of the Notifier base class provided by each UX. It
needs to support two methods
The notifier is a concrete implementation of the Notifier trait provided by each UX. It
needs to support two methods:
```python
class Notifier:
def log(self, RichRenderable):
...
def status(self, str):
...
```rust
trait Notifier {
fn log(&self, content: RichRenderable);
fn status(&self, message: String);
}
```
Log is meant to record something concrete that happened, such as a tool being called, and status is intended
@@ -110,57 +109,69 @@ for transient displays of the current status. For example, while a shell command
`.log` to record the command that started, and then update the status to `"shell command running"`. Log is durable
while Status is ephemeral.
## Toolkits
## Extensions
Toolkits are a collection of tools, along with the state and prompting they require.
Toolkits are what gives Goose its capabilities.
Extensions are a collection of tools, along with the state and prompting they require.
Extensions are what gives Goose its capabilities.
Tools need a way to report what's happening back to the user, which we treat similarly
to logging. To make that possible, toolkits get a reference to the interface described above.
to logging. To make that possible, extensions get a reference to the interface described above.
```python
class ScheduleToolkit(Toolkit):
def __init__(self, notifier: Notifier, requires: Requirements, **kwargs):
super().__init__(notifier, requires, **kwargs) # handles the interface, exchangeview
# for a class that has requirements, you can get them like this
self.calendar = requires.get("calendar")
self.assistant = requires.get("assistant")
self.contacts = requires.get("contacts")
self.appointments_state = []
```rust
struct ScheduleExtension {
notifier: Box<dyn Notifier>,
calendar: Box<dyn Calendar>,
assistant: Box<dyn Assistant>,
contacts: Box<dyn Contacts>,
appointments_state: Vec<Appointment>,
}
def prompt(self) -> str:
return "Try out the example tool."
impl Extension for ScheduleExtension {
fn new(notifier: Box<dyn Notifier>, requires: Requirements) -> Self {
Self {
notifier,
calendar: requires.get("calendar"),
assistant: requires.get("assistant"),
contacts: requires.get("contacts"),
appointments_state: vec![],
}
}
@tool
def example(self):
self.interface.log(f"An example tool was called, current state is {self.state}")
fn prompt(&self) -> String {
"Try out the example tool.".to_string()
}
#[tool]
fn example(&self) {
self.notifier.log(format!("An example tool was called, current state is {:?}", self.appointments_state));
}
}
```
### Advanced
**Dependencies**: Toolkits can depend on each other, to make it easier to get plugins to extend
or modify existing capabilities. In the config above, you can see this used for the scheduling toolkit.
**Dependencies**: Extensions can depend on each other, to make it easier to get plugins to extend
or modify existing capabilities. In the config above, you can see this used for the scheduling extension.
You can refer to those requirements in code through:
```python
@tool
def example_dependency(self):
appointments = self.dependencies["calendar"].appointments
...
```rust
#[tool]
fn example_dependency(&self) {
let appointments = self.calendar.appointments();
// ...
}
```
**ExchangeView**: It can also be useful for tools to have a read-only copy of the history
of the loop so far. So for advanced use cases, toolkits also have access to an
of the loop so far. So for advanced use cases, extensions also have access to an
`ExchangeView` object.
```python
@tool
def example_history(self):
last_message = self.exchange_view.processor.messages[-1]
...
```rust
#[tool]
fn example_history(&self) {
let last_message = self.exchange_view.processor.messages.last();
// ...
}
```
[exchange]: https://github.com/block/goose/tree/main/packages/exchange
[exchange]: https://github.com/block/goose/tree/main/packages/exchange

View File

@@ -1,144 +0,0 @@
# Changelog
All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [0.9.11] - 2024-11-05
- fix: removed the unexpected tracing message (#219)
## [0.9.10] - 2024-11-05
- fix: bump version of exchange to 0.9.9
## [0.9.9] - 2024-11-04
- fix: only run summarize when there is something to summarize vs up front (#195)
- feat: add browser toolkit (#179)
- docs: reworks README to be outcome based (#151)
- chore: license-checker improvements (#212)
- fix: prompt rendering (#208)
- fix: Cost calculation enhancement (#207)
- refactor: used recommended_models function as default model (#209)
- feat: default to synopsis (#206)
- chore: housekeeping (#202)
- feat: reduce tool entrypoints in synopsis for text editor, bash, process manager (#191)
- feat: list moderators (#204)
- chore: Minor changes in providers envs logic (#161)
- chore: add `.vscode` workspace settings and suggested extensions (#200)
- feat: license checker (#201)
- feat: include new anthropic model in docs and recommended config (#198)
- feat: tiny change so you know what processing is doing (#196)
- docs: correct intellij link (#197)
- docs: remove small duplication (#194)
- chore: add tracing option to run and group traces under session name (#187)
- docs: fix mkdocs, update to new references (#193)
- feat: support optional params jsonschema conversion in exchange (#188)
- fix: correct context loading from session new/overwrite and resume (#180)
- feat: trying a license checker (#184)
- docs: getting index.md in sync with readme (#183)
- chore: Update block-open-source ref to block (#181)
- fix: just adding stuff from developer.py to synopsis developer (#182)
- fix: Simplifed langfuse auth check (#177)
- test: add vision tests for google (#160)
- fix: adding a release checker (#176)
## [0.9.8] - 2024-10-20
- fix: added specific ai-exchange version and check path is None to avoid error after starting goose (#174)
## [0.9.7] - 2024-10-18
- chore: update tiktoken to support python 3.13
## [0.9.6] - 2024-10-18
- fix: update summarizer_exchange model to use default model from input exchange (#139)
- feat: Add synopisis approach for the core goose loop (#166)
- feat: web browsing (#154)
- chore: restrict python version (#162)
- feat: Run with resume session (#153)
- refactor: move langfuse wrapper to a module in exchange instead of a package (#138)
- docs: add subheaders to the 'Other ways to run Goose' section (#155)
- fix: Remove tools from exchange when summarizing files (#157)
- chore: use primitives instead of typing imports and fixes completion … (#149)
- chore: make vcr tests pretty-print JSON (#146)
## [0.9.5] - 2024-10-15
- chore: updates ollama default model from mistral-nemo to qwen2.5 (#150)
- feat: add vision support for Google (#141)
- fix: session resume with arg handled incorrectly (#145)
- docs: add release instructions to CONTRIBUTING.md (#143)
- docs: add link to action, IDE words (#140)
- docs: goosehints doc fix only (#142)
## [0.9.4] - 2024-10-10
- revert: "feat: add local langfuse tracing option (#106)"
- feat: add local langfuse tracing option (#106)
- feat: add groq provider (#134)
- feat: add a deep thinking reasoner model (o1-preview/mini) (#68)
- fix: use concrete SessionNotifier (#135)
- feat: add guards to session management (#101)
- fix: Set default model configuration for the Google provider. (#131)
- test: convert Google Gemini tests to VCR (#118)
- chore: Add goose providers list command (#116)
- docs: working ollama for desktop (#125)
- docs: format and clean up warnings/errors (#120)
- docs: update deploy workflow (#124)
- feat: Implement a goose run command (#121)
- feat: saved api_key to keychain for user (#104)
- docs: add callout plugin (#119)
- chore: add a page to docs for Goose application examples (#117)
- fix: exit the goose and show the error message when provider environment variable is not set (#103)
- fix: Update OpenAI pricing per https://openai.com/api/pricing/ (#110)
- fix: update developer tool prompts to use plan task status to match allowable statuses update_plan tool call (#107)
- fix: removed the panel in the output so that the user won't have unnecessary pane borders in the copied content (#109)
- docs: update links to exchange to the new location (#108)
- chore: setup workspace for exchange (#105)
- fix: resolve uvx when using a git client or IDE (#98)
- ci: add include-markdown for mkdocs (#100)
- chore: fix broken badge on readme (#102)
- feat: add global optional user goosehints file (#73)
- docs: update docs (#99)
## [0.9.3] - 2024-09-25
- feat: auto save sessions before next user input (#94)
- fix: removed the diff when default profile changes (#92)
- feat: add shell-completions subcommand (#76)
- chore: update readme! (#96)
- chore: update docs again (#77)
- fix: remove overly general match for long running commands (#87)
- fix: default ollama to tested model (#88)
- fix: Resize file in screen toolkit (#81)
- fix: enhance shell() to know when it is interactive (#66)
- docs: document how to run goose fully from source from any dir (#83)
- feat: track cost and token usage in log file (#80)
- chore: add link to docs in read me (#85)
- docs: add in ollama (#82)
- chore: add just command for releasing goose (#55)
- feat: support markdown plans (#79)
- feat: add version options (#74)
- docs: fixing exchange url to public version (#67)
- docs: Update CONTRIBUTING.md (#69)
- chore: create mkdocs for goose (#70)
- docs: fix broken link (#71)
- feat: give commands the ability to execute logic (#63)
- feat: jira toolkit (#59)
- feat: run goose in a docker-style sandbox (#44)
## [0.9.0] - 2024-09-10
This also updates the minimum version of exchange to 0.9.0.
- fix: goose should track files it reads and not overwrite changes (#46)
- docs: Small dev notes for using exchange from source (#50)
- fix: typo in exchange method rewind (#54)
- fix: remove unsafe pop of messages (#47)
- chore: Update LICENSE (#53)
- chore(docs): update is_dangerous_command method description (#48)
- refactor: improve safety rails speed and prompt (#45)
- feat: make goosehints jinja templated (#43)
- ci: enforce PR title follows conventional commit (#14)
- feat: show available toolkits (#37)
- feat: adding in ability to provide per repo hints (#32)
- chore: apply ruff and add to CI (#40)
- feat: added some regex based checks for dangerous commands (#38)
- chore: Update publish github workflow to check package versions before publishing (#19)

View File

@@ -1,114 +1,101 @@
# Contributing
We welcome Pull Requests for general contributions. If you have a larger new feature or any questions on how to develop a fix, we recommend you open an [issue][issues] before starting.
Goose is Open Source!
We welcome Pull Requests for general contributions! If you have a larger new feature or any questions on how to develop a fix, we recommend you open an [issue][issues] before starting.
## Prerequisites
Goose uses [uv][uv] for dependency management, and formats with [ruff][ruff].
Clone goose and make sure you have installed `uv` to get started. When you use
`uv` below in your local goose directly, it will automatically setup the virtualenv
and install dependencies.
Goose includes rust binaries alongside an electron app for the GUI. To work
on the rust backend, you will need to [install rust and cargo][rustup]. To work
on the App, you will also need to [install node and npm][nvm] - we recommend through nvm.
We provide a shortcut to standard commands using [just][just] in our `justfile`.
## Development
## Getting Started
Now that you have a local environment, you can make edits and run our tests!
### Rust
### Run Goose
If you've made edits and want to try them out, use
First let's compile goose and try it out
```
uv run goose session start
cargo build
```
or other `goose` commands.
If you want to run your local changes but in another directory, you can use the path in
the virtualenv created by uv:
when that is done, you should now have debug builds of the binaries like the goose cli:
```
alias goosedev=`uv run which goose`
./target/debug/goose --help
```
You can then run `goosedev` from another dir and it will use your current changes.
If you haven't used the CLI before, you can use this compiled version to do first time configuration:
### Run Tests
To run the test suite against your edges, use `pytest`:
```sh
uv run pytest tests -m "not integration"
```
./target/debug/goose configure
```
or, as a shortcut,
And then once you have a connection to an LLM provider working, you can run a session!
```sh
just test
```
./target/debug/goose session
```
### Enable traces in Goose with [locally hosted Langfuse](https://langfuse.com/docs/deployment/self-host)
> [!NOTE]
> This integration is experimental and we don't currently have integration tests for it.
Developers can use locally hosted Langfuse tracing by applying the custom `observe_wrapper` decorator defined in `packages/exchange/src/exchange/observers` to functions for automatic integration with Langfuse, and potentially other observability providers in the future.
These same commands can be recompiled and immediately run using `cargo run -p goose-cli` for iteration.
As you make changes to the rust code, you can try it out on the CLI, or also run checks and tests:
```
cargo check # do your changes compile
cargo test # do the tests pass with your changes.
```
### Node
Now let's make sure you can run the app.
```
just run-ui
```
The start gui will both build a release build of rust (as if you had done `cargo build -r`) and start the electron process.
You should see the app open a window, and drop you into first time setup. When you've gone through the setup,
you can talk to goose!
You can now make changes in the code in ui/desktop to iterate on the GUI half of goose.
## Env Vars
You may want to make more frequent changes to your provider setup or similar to test things out
as a developer. You can use environment variables to change things on the fly without redoing
your configuration.
> [!TIP]
> At the moment, we are still updating some of the CLI configuration to make sure this is
> respected.
You can change the provider goose points to via the `GOOSE_PROVIDER` env var. If you already
have a credential for that provider in your keychain from previously setting up, it should
reuse it. For things like automations or to test without doing official setup, you can also
set the relevant env vars for that provider. For example `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`,
or `DATABRICKS_HOST`. Refer to the provider details for more info on required keys.
## Enable traces in Goose with [locally hosted Langfuse](https://langfuse.com/docs/deployment/self-host)
- Add an `observers` array to your profile containing `langfuse`.
- Run `just langfuse-server` to start your local Langfuse server. It requires Docker.
- Go to http://localhost:3000 and log in with the default email/password output by the shell script (values can also be found in the `.env.langfuse.local` file).
- Run Goose with the --tracing flag enabled i.e., `goose session start --tracing`
- View your traces at http://localhost:3000
- Set the environment variables so that rust can connect to the langfuse server
`To extend tracing to additional functions, import `from exchange.observers import observe_wrapper` and use the `observe_wrapper()` decorator on functions you wish to enable tracing for. `observe_wrapper` functions the same way as Langfuse's observe decorator.
Read more about Langfuse's decorator-based tracing [here](https://langfuse.com/docs/sdk/python/decorators).
### Other observability plugins
In case locally hosted Langfuse doesn't fit your needs, you can alternatively use other `observer` telemetry plugins to ingest data with the same interface as the Langfuse integration.
To do so, extend `packages/exchange/src/exchange/observers/base.py:Observer` and include the new plugin's path as an entrypoint in `exchange`'s `pyproject.toml`.
## Exchange
The lower level generation behind goose is powered by the [`exchange`][ai-exchange] package, also in this repo.
Thanks to `uv` workspaces, any changes you make to `exchange` will be reflected in using your local goose. To run tests
for exchange, head to `packages/exchange` and run tests just like above
```sh
uv run pytest tests -m "not integration"
```
export LANGFUSE_INIT_PROJECT_PUBLIC_KEY=publickey-local
export LANGFUSE_INIT_PROJECT_SECRET_KEY=secretkey-local
```
## Evaluations
Given that so much of Goose involves interactions with LLMs, our unit tests only go so far to confirming things work as intended.
We're currently developing a suite of evaluations, to make it easier to make improvements to Goose more confidently.
In the meantime, we typically incubate any new additions that change the behavior of the Goose through **opt-in** plugins - `Toolkit`s, `Moderator`s, and `Provider`s. We welcome contributions of plugins that add new capabilities to *goose*. We recommend sending in several examples of the new capabilities in action with your pull request.
Additions to the [developer toolkit][developer] change the core performance, and so will need to be measured carefully.
Then you can view your traces at http://localhost:3000
## Conventional Commits
This project follows the [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) specification for PR titles. Conventional Commits make it easier to understand the history of a project and facilitate automation around versioning and changelog generation.
## Release
In order to release a new version of goose, you need to do the following:
1. Update CHANGELOG.md. To get the commit messages since last release, run: `just release-notes`
2. Update version in `pyproject.toml` for `goose` and package dependencies such as `exchange`
3. Create a PR and merge it into main branch
4. Tag the HEAD commit in main branch. To do this, switch to main branch and run: `just tag-push`
5. Publish a new release from the [Github Release UI](https://github.com/block/goose/releases)
[issues]: https://github.com/block/goose/issues
[goose-plugins]: https://github.com/block-open-source/goose-plugins
[ai-exchange]: https://github.com/block/goose/tree/main/packages/exchange
[developer]: https://github.com/block/goose/blob/dfecf829a83021b697bf2ecc1dbdd57d31727ddd/src/goose/toolkit/developer.py
[uv]: https://docs.astral.sh/uv/
[ruff]: https://docs.astral.sh/ruff/
[just]: https://github.com/casey/just
[adding-toolkit]: https://block.github.io/goose/configuration.html#adding-a-toolkit
[rustup]: https://doc.rust-lang.org/cargo/getting-started/installation.html
[nvm]: https://github.com/nvm-sh/nvm
[just]: https://github.com/casey/just?tab=readme-ov-file#installation

11
Cargo.toml Normal file
View File

@@ -0,0 +1,11 @@
[workspace]
members = ["crates/*"]
resolver = "2"
[workspace.package]
edition = "2021"
version = "1.0.0"
authors = ["Block <ai-oss-tools@block.xyz>"]
license = "Apache-2.0"
repository = "https://github.com/block/goose"
description = "An AI agent"

29
Cross.toml Normal file
View File

@@ -0,0 +1,29 @@
# Configuration for cross-compiling using cross
[target.aarch64-unknown-linux-gnu]
xargo = false
pre-build = [
# Add the ARM64 architecture and install necessary dependencies
"dpkg --add-architecture arm64",
"""\
apt-get update --fix-missing && apt-get install -y \
pkg-config \
libssl-dev:arm64 \
libdbus-1-dev:arm64 \
libxcb1-dev:arm64
"""
]
env = { PKG_CONFIG_PATH = "/usr/lib/aarch64-linux-gnu/pkgconfig" }
[target.x86_64-unknown-linux-gnu]
xargo = false
pre-build = [
# Install necessary dependencies for x86_64
# We don't need architecture-specific flags because x86_64 dependencies are installable on Ubuntu system
"""\
apt-get update && apt-get install -y \
pkg-config \
libssl-dev \
libdbus-1-dev \
libxcb1-dev \
"""
]

View File

@@ -1,27 +0,0 @@
# Use an official Python runtime as a parent image
FROM python:3.10-slim
# Set the working directory in the container
WORKDIR /app
# Install SSL certificates, update certifi, install pipx and move to path
RUN apt-get update && apt-get install -y ca-certificates git curl make \
&& pip install --upgrade certifi \
&& pip install pipx \
&& pipx ensurepath
# Install goose-ai CLI using pipx
RUN pipx install goose-ai
# Make sure the PATH is updated
ENV PATH="/root/.local/bin:${PATH}"
# Run an infinite loop to keep the container running for testing
ENTRYPOINT ["goose", "session", "start"]
# once built, you can run this with something like:
# docker run -it --env OPENAI_API_KEY goose-ai
# or to run against ollama running on the same host
# docker run -it --env OPENAI_HOST=http://host.docker.internal:11434/ --env OPENAI_API_KEY=unused goose-ai
#
# To use goose in a docker style sandbox for experimenting.

37
Justfile Normal file
View File

@@ -0,0 +1,37 @@
# Justfile
# Default release command
release:
@echo "Building release version..."
cargo build --release
@just copy-binary
# Copy binary command
copy-binary:
@if [ -f ./target/release/goosed ]; then \
echo "Copying goosed binary to ui/desktop/src/bin with permissions preserved..."; \
cp -p ./target/release/goosed ./ui/desktop/src/bin/; \
else \
echo "Release binary not found."; \
exit 1; \
fi
# Run UI with latest
run-ui:
@just release
@echo "Running UI..."
cd ui/desktop && npm install && npm run start-gui
# Run server
run-server:
@echo "Running server..."
cargo run -p goose-server
# make GUI with latest binary
make-ui:
@just release
cd ui/desktop && npm run bundle:default
# Setup langfuse server
langfuse-server:
#!/usr/bin/env bash
./scripts/setup_langfuse.sh

265
README.md
View File

@@ -1,274 +1,21 @@
<h1 align="center">
Goose is your on-machine developer agent, working for you, on your terms
<code>codename goose</code>
</h1>
<p align="center">
<img src="docs/assets/goose.png" width="400" height="400" alt="Goose Drawing"/>
</p>
<p align="center">
Generated by Goose from its <a href="https://github.com/block-open-source/goose-plugins/blob/main/src/goose_plugins/toolkits/artify.py">VincentVanCode toolkit</a>.
<strong>an open-source, extensible AI agent that goes beyond code suggestions<br>install, execute, edit, and test with any LLM</strong>
</p>
<p align="center">
<a href="https://block.github.io/goose/">
<img src="https://img.shields.io/badge/Documentation-goose_docs-teal">
</a>
<a href=https://pypi.org/project/goose-ai>
<img src="https://img.shields.io/pypi/v/goose-ai?color=green">
</a>
<a href="https://opensource.org/licenses/Apache-2.0">
<img src="https://img.shields.io/badge/License-Apache_2.0-blue.svg">
</a>
<a href="https://discord.gg/7GaTvbDwga">
<img src="https://img.shields.io/discord/1287729918100246654?logo=discord&logoColor=white&label=Join+Us&color=blueviolet" alt="Discord">
</a>
<a href="https://github.com/block/goose/actions/workflows/ci.yml">
<img src="https://img.shields.io/github/actions/workflow/status/block/goose/ci.yml?branch=main" alt="CI">
</a>
</p>
<p align="center">
<a href="#unique-features-of-goose-compared-to-other-ai-assistants">Unique features</a> 🤖 •
<a href="#what-users-have-to-say-about-goose"> Testimonials on Goose</a> 👩‍💻 •
<a href="#quick-start-guide">Quick start guide</a> 🚀 •
<a href="#getting-involved">Getting involved!</a> 👋
</p>
> [!TIP]
> **Quick install:**
> ```
> pipx install goose-ai
> ```
Each day, we have tasks that stretch our days. Maybe it is completing that half-complete change started by someone now
on holiday, or figuring out the tl;dr; of that 50 comment pull request. Wouldn't it be grand if someone else can do
this, or at least start the work for us?
**Goose** is your on-machine developer agent, working for you, on your terms. Guided by you, Goose intelligently
assesses what you need and generates required code or modifications. You are in charge: Do you prefer Goose to make a
draft, or complete the change entirely? Do you prefer to work in a terminal or in your IDE?
Doing our work requires a lot of tools like Jira, GitHub, Slack, as well APIs for infrastructure and data pipelines.
Goose handles all of these, and is extensible. Goose can run anything invocable by a shell command, Python or a plugin.
Like semi-autonomous driving, Goose handles the heavy lifting, allowing you to focus on other priorities. Simply set it
on a task and return later to find it completed, boosting your productivity with less manual effort. Read on to get
started!
<p align="center">
<video src="https://github.com/user-attachments/assets/63ee7910-cb02-45c0-9982-351cbce83925" width="700" height="700" />
</p>
## Unique features of Goose compared to other AI assistants
- **Autonomy**: A copilot should be able to also fly the plane at times, which in the development world means running code, debugging tests, installing dependencies, not just providing text output and autocomplete or search. Goose moves beyond just generating code snippets by (1) **using the shell** and (2) by seeing what happens with the code it writes and starting a feedback loop to solve harder problems, **refining solutions iteratively like a human developer**. Your code's best wingman.
- **Extensibility**: Open-source and fully customizable, Goose integrates with your workflow and allows you to extend it for even more control. **Toolkits let you add new capabilities to Goose.** They are anything you can implement as a Python function (e.g. API requests, deployments, search, etc). We have a growing library of toolkits to use, but more importantly you can create your own. This gives Goose the ability to run these commands and decide if and when a tool is needed to complete your request! **Creating your own toolkits give you a way to bring your own private context into Goose's capabilities.** And you can use *any* LLM you want under the hood, as long as it supports tool use.
## What users have to say about Goose
> With Goose, I feel like I am Maverick.
>
> Thanks a ton for creating this. 🙏
> I have been having way too much fun with it today.
-- P, Machine Learning Engineer
> I wanted to construct some fake data for an API with a large request body and business rules I haven't memorized. So I told Goose which object to update and a test to run that calls the vendor. Got it to use the errors descriptions from the vendor response to keep correcting the request until it was successful. So good!
-- J, Software Engineer
> I asked Goose to write up a few Google Scripts that mimic Clockwise's functionality (particularly, creating blocks on my work calendar based on events in my personal calendar, as well as color-coding calendar entries based on type and importance). Took me under an hour. If you haven't tried Goose yet, I highly encourage you to do so!
-- M, Software Engineer
> If anyone was looking for another reason to check it out: I just asked Goose to break a string-array into individual string resources across eleven localizations, and it performed amazingly well and saved me a bunch of time doing it manually or figuring out some way to semi-automate it.
-- A, Android Engineer
> Hi team, thank you for much for making Goose, it's so amazing. Our team is working on migrating Dashboard components to React components. I am working with Goose to help the migration.
-- K, Software Engineer
> Got Goose to update a dependency, run tests, make a branch and a commit... it was 🤌. Not that complicated but I was impressed it figured out how to run tests from the README.
-- J, Software Engineer
> Wanted to document what I had Goose do -- took about 30 minutes end to end! I created a custom CLI command in the `gh CLI` library to download in-line comments on PRs about code changes (currently they aren't directly viewable). I don't know Go *that well* and I definitely didn't know where to start looking in the code base or how to even test the new command was working and Goose did it all for me 😁
-- L, Software Engineer
> Hi Team, just wanted to share my experience of using Goose as a non-engineer! ... I just asked Goose to ensure that my environment is up to date and copied over a guide into my prompt. Goose managed everything flawlessly, keeping me informed at every step... I was truly impressed with how well it works and how easy it was to get started! 😍
-- M, Product Manager
**See more of our use-cases in our [docs][use-cases]!**
## Quick start guide
### Installation
To install Goose, use `pipx`. First ensure [pipx][pipx] is installed:
``` sh
brew install pipx
pipx ensurepath
```
Then install Goose:
```sh
pipx install goose-ai
```
### Running Goose
#### Start a session
From your terminal, navigate to the directory you'd like to start from and run:
```sh
goose session start
```
#### Set up a provider
Goose works with your [preferred LLM][providers]. By default, it uses `openai` as the LLM provider. You'll be prompted to set an [API key][openai-key] if you haven't set one previously.
>[!TIP]
> **Billing:**
>
> You will need to have credits in your LLM Provider account to be able to successfully make requests.
>
#### Make Goose do the work for you
You will see the Goose prompt `G`:
```
G type your instructions here exactly as you would speak to a developer.
```
Now you are interacting with Goose in conversational sessions - think of it as like giving direction to a junior developer. The default toolkit allows Goose to take actions through shell commands and file edits. You can interrupt Goose with `CTRL+D` or `ESC+Enter` at any time to help redirect its efforts.
> [!TIP]
> You can place a `.goosehints` text file in any directory you launch goose from to give it some background info for new sessions in plain language (eg how to test, what instructions to read to get started or just tell it to read the README!) You can also put a global one `~/.config/goose/.goosehints` if you like for always loaded hints personal to you.
### Running a goose tasks (one off)
You can run goose to do things just as a one off, such as tidying up, and then exiting:
```sh
goose run instructions.md
```
You can also use process substitution to provide instructions directly from the command line:
```sh
goose run <(echo "Create a new Python file that prints hello world")
```
This will run until completion as best it can. You can also pass `--resume-session` and it will re-use the first session it finds for context
#### Exit the session
If you are looking to exit, use `CTRL+D`, although Goose should help you figure that out if you forget.
#### Resume a session
When you exit a session, it will save the history in `~/.config/goose/sessions` directory. You can then resume your last saved session later, using:
``` sh
goose session resume
```
To see more documentation on the CLI commands currently available to Goose check out the documentation [here][cli]. If youd like to develop your own CLI commands for Goose, check out the [Contributing document][contributing].
### Tracing with Langfuse
> [!NOTE]
> This Langfuse integration is experimental and we don't currently have integration tests for it.
The exchange package provides a [Langfuse](https://langfuse.com/) wrapper module. The wrapper serves to initialize Langfuse appropriately if the Langfuse server is running locally and otherwise to skip applying the Langfuse observe descorators.
#### Start your local Langfuse server
Run `just langfuse-server` to start your local Langfuse server. It requires Docker.
Read more about local Langfuse deployments [here](https://langfuse.com/docs/deployment/local).
#### Exchange and Goose integration
Import `from exchange.observers import observe_wrapper`, include `langfuse` in the `observers` list of your profile, and use the `observe_wrapper()` decorator on functions you wish to enable tracing for. `observe_wrapper` functions the same way as Langfuse's observe decorator.
Read more about Langfuse's decorator-based tracing [here](https://langfuse.com/docs/sdk/python/decorators).
In Goose, initialization requires certain environment variables to be present:
- `LANGFUSE_PUBLIC_KEY`: Your Langfuse public key
- `LANGFUSE_SECRET_KEY`: Your Langfuse secret key
- `LANGFUSE_BASE_URL`: The base URL of your Langfuse instance
By default your local deployment and Goose will use the values in `.env.langfuse.local`.
### Next steps
Learn how to modify your Goose profiles.yaml file to add and remove functionality (toolkits) and providing context to get the most out of Goose in our [Getting Started Guide][getting-started].
## Other ways to run goose
**Want to move out of the terminal and into an IDE?**
We have some experimental IDE integrations for VSCode and JetBrains IDEs:
* https://github.com/square/goose-vscode
* https://github.com/Kvadratni/goose-intellij
**Goose as a Github Action**
There is also an experimental Github action to run goose as part of your workflow (for example if you ask it to fix an issue):
https://github.com/marketplace/actions/goose-ai-developer-agent
**With Docker**
There is also a `Dockerfile` in the root of this project you can use if you want to run goose in a sandboxed fashion.
## Getting involved!
There is a lot to do! If you're interested in contributing, a great place to start is picking a `good-first-issue`-labelled ticket from our [issues list][gh-issues]. More details on how to develop Goose can be found in our [Contributing Guide][contributing]. We are a friendly, collaborative group and look forward to working together![^1]
Check out and contribute to more experimental features in [Goose Plugins][goose-plugins]!
Let us know what you think in our [Discussions][discussions] or the [**`#goose`** channel on Discord][goose-channel].
[^1]: Yes, Goose is open source and always will be. Goose is released under the ASL2.0 license meaning you are free to use it however you like. See [LICENSE.md][license] for more details.
[goose-plugins]: https://github.com/block-open-source/goose-plugins
[pipx]: https://github.com/pypa/pipx?tab=readme-ov-file#install-pipx
[contributing]: https://github.com/block/goose/blob/main/CONTRIBUTING.md
[license]: https://github.com/block/goose/blob/main/LICENSE
[goose-docs]: https://block.github.io/goose/
[toolkits]: https://block.github.io/goose/plugins/available-toolkits.html
[configuration]: https://block.github.io/goose/configuration.html
[cli]: https://block.github.io/goose/plugins/cli.html
[providers]: https://block.github.io/goose/plugins/providers.html
[use-cases]: https://block.github.io/goose/guidance/applications.html
[getting-started]: https://block.github.io/goose/guidance/getting-started.html
[openai-key]: https://platform.openai.com/api-keys
[discord-invite]: https://discord.gg/7GaTvbDwga
[gh-issues]: https://github.com/block/goose/issues
[van-code]: https://github.com/block-open-source/goose-plugins/blob/de98cd6c29f8e7cd3b6ace26535f24ac57c9effa/src/goose_plugins/toolkits/artify.py
[discussions]: https://github.com/block/goose/discussions
[goose-channel]: https://discord.com/channels/1287729918100246654/1287729920319033345
Stay tuned for the upcoming 1.0 release by the end of January 2025. You can find the v0.X documentation on our [github pages](https://block.github.io/goose/).

3
crates/goose-cli/.gitignore vendored Normal file
View File

@@ -0,0 +1,3 @@
.goosehints
.goose

View File

@@ -0,0 +1,54 @@
[package]
name = "goose-cli"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
description.workspace = true
[[bin]]
name = "goose"
path = "src/main.rs"
[dependencies]
goose = { path = "../goose" }
goose-mcp = { path = "../goose-mcp" }
mcp-client = { path = "../mcp-client" }
mcp-server = { path = "../mcp-server" }
mcp-core = { path = "../mcp-core" }
clap = { version = "4.4", features = ["derive"] }
cliclack = "0.3.5"
console = "0.15.8"
bat = "0.24.0"
anyhow = "1.0"
serde_json = "1.0"
tokio = { version = "1.0", features = ["full"] }
futures = "0.3"
serde = { version = "1.0", features = ["derive"] } # For serialization
serde_yaml = "0.9"
dirs = "4.0"
reqwest = { version = "0.12.9", features = [
"rustls-tls",
"json",
"cookies",
"gzip",
"brotli",
"deflate",
"zstd",
"charset",
"http2",
"stream"
], default-features = false }
rand = "0.8.5"
rustyline = "15.0.0"
tracing = "0.1"
chrono = "0.4"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json", "time"] }
tracing-appender = "0.2"
[dev-dependencies]
tempfile = "3"
temp-env = { version = "0.3.6", features = ["async_closure"] }
test-case = "3.3"
tokio = { version = "1.0", features = ["rt", "macros"] }

View File

@@ -0,0 +1,28 @@
use anyhow::Result;
use clap::Args;
use goose::agents::AgentFactory;
use std::fmt::Write;
#[derive(Args)]
pub struct AgentCommand {}
impl AgentCommand {
pub fn run(&self) -> Result<()> {
let mut output = String::new();
writeln!(output, "Available agent versions:")?;
let versions = AgentFactory::available_versions();
let default_version = AgentFactory::default_version();
for version in versions {
if version == default_version {
writeln!(output, "* {} (default)", version)?;
} else {
writeln!(output, " {}", version)?;
}
}
print!("{}", output);
Ok(())
}
}

View File

@@ -0,0 +1,486 @@
use cliclack::spinner;
use console::style;
use goose::agents::{extension::Envs, ExtensionConfig};
use goose::config::{Config, ExtensionEntry, ExtensionManager};
use goose::message::Message;
use goose::providers::{create, providers};
use serde_json::Value;
use std::collections::HashMap;
use std::error::Error;
pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
let config = Config::global();
if !config.exists() {
// First time setup flow
println!();
println!(
"{}",
style("Welcome to goose! Let's get you set up with a provider.").dim()
);
println!(
"{}",
style(" you can rerun this command later to update your configuration").dim()
);
println!();
cliclack::intro(style(" goose-configure ").on_cyan().black())?;
if configure_provider_dialog().await? {
println!(
"\n {}: Run '{}' again to adjust your config or add extensions",
style("Tip").green().italic(),
style("goose configure").cyan()
);
// Since we are setting up for the first time, we'll also enable the developer system
ExtensionManager::set(ExtensionEntry {
enabled: true,
config: ExtensionConfig::Builtin {
name: "developer".to_string(),
},
})?;
} else {
let _ = config.clear();
println!(
"\n {}: We did not save your config, inspect your credentials\n and run '{}' again to ensure goose can connect",
style("Warning").yellow().italic(),
style("goose configure").cyan()
);
}
Ok(())
} else {
println!();
println!(
"{}",
style("This will update your existing config file").dim()
);
println!(
"{} {}",
style(" if you prefer, you can edit it directly at").dim(),
config.path()
);
println!();
cliclack::intro(style(" goose-configure ").on_cyan().black())?;
let action = cliclack::select("What would you like to configure?")
.item(
"providers",
"Configure Providers",
"Change provider or update credentials",
)
.item(
"toggle",
"Toggle Extensions",
"Enable or disable connected extensions",
)
.item("add", "Add Extension", "Connect to a new extension")
.interact()?;
match action {
"toggle" => toggle_extensions_dialog(),
"add" => configure_extensions_dialog(),
"providers" => configure_provider_dialog().await.and(Ok(())),
_ => unreachable!(),
}
}
}
/// Dialog for configuring the AI provider and model
pub async fn configure_provider_dialog() -> Result<bool, Box<dyn Error>> {
// Get global config instance
let config = Config::global();
// Get all available providers and their metadata
let available_providers = providers();
// Create selection items from provider metadata
let provider_items: Vec<(&String, &str, &str)> = available_providers
.iter()
.map(|p| (&p.name, p.display_name.as_str(), p.description.as_str()))
.collect();
// Get current default provider if it exists
let current_provider: Option<String> = config.get("GOOSE_PROVIDER").ok();
let default_provider = current_provider.unwrap_or_default();
// Select provider
let provider_name = cliclack::select("Which model provider should we use?")
.initial_value(&default_provider)
.items(&provider_items)
.interact()?;
// Get the selected provider's metadata
let provider_meta = available_providers
.iter()
.find(|p| &p.name == provider_name)
.expect("Selected provider must exist in metadata");
// Configure required provider keys
for key in &provider_meta.config_keys {
if !key.required {
continue;
}
// First check if the value is set via environment variable
let from_env = std::env::var(&key.name).ok();
match from_env {
Some(env_value) => {
let _ =
cliclack::log::info(format!("{} is set via environment variable", key.name));
if cliclack::confirm("Would you like to save this value to your config file?")
.initial_value(true)
.interact()?
{
if key.secret {
config.set_secret(&key.name, Value::String(env_value))?;
} else {
config.set(&key.name, Value::String(env_value))?;
}
let _ = cliclack::log::info(format!("Saved {} to config file", key.name));
}
}
None => {
// No env var, check config/secret storage
let existing: Result<String, _> = if key.secret {
config.get_secret(&key.name)
} else {
config.get(&key.name)
};
match existing {
Ok(_) => {
let _ = cliclack::log::info(format!("{} is already configured", key.name));
if cliclack::confirm("Would you like to update this value?").interact()? {
let new_value: String = if key.secret {
cliclack::password(format!("Enter new value for {}", key.name))
.mask('▪')
.interact()?
} else {
cliclack::input(format!("Enter new value for {}", key.name))
.interact()?
};
if key.secret {
config.set_secret(&key.name, Value::String(new_value))?;
} else {
config.set(&key.name, Value::String(new_value))?;
}
}
}
Err(_) => {
let value: String = if key.secret {
cliclack::password(format!(
"Provider {} requires {}, please enter a value",
provider_meta.display_name, key.name
))
.mask('▪')
.interact()?
} else {
cliclack::input(format!(
"Provider {} requires {}, please enter a value",
provider_meta.display_name, key.name
))
.interact()?
};
if key.secret {
config.set_secret(&key.name, Value::String(value))?;
} else {
config.set(&key.name, Value::String(value))?;
}
}
}
}
}
}
// Select model, defaulting to the provider's recommended model
let default_model = config
.get("GOOSE_MODEL")
.unwrap_or(provider_meta.default_model.clone());
let model: String = cliclack::input("Enter a model from that provider:")
.default_input(&default_model)
.interact()?;
// Update config with new values
config.set("GOOSE_PROVIDER", Value::String(provider_name.to_string()))?;
config.set("GOOSE_MODEL", Value::String(model.clone()))?;
// Test the configuration
let spin = spinner();
spin.start("Checking your configuration...");
let model_config = goose::model::ModelConfig::new(model.clone());
let provider = create(provider_name, model_config)?;
let message = Message::user().with_text(
"Please give a nice welcome message (one sentence) and let them know they are all set to use this agent"
);
let result = provider
.complete(
"You are an AI agent called Goose. You use tools of connected extensions to solve problems.",
&[message],
&[]
)
.await;
match result {
Ok((message, _usage)) => {
if let Some(content) = message.content.first() {
if let Some(text) = content.as_text() {
spin.stop(text);
} else {
spin.stop("No response text available");
}
} else {
spin.stop("No response content available");
}
cliclack::outro("Configuration saved successfully")?;
Ok(true)
}
Err(e) => {
println!("{:?}", e);
spin.stop("We could not connect!");
let _ = cliclack::outro("The provider configuration was invalid");
Ok(false)
}
}
}
/// Configure extensions that can be used with goose
/// Dialog for toggling which extensions are enabled/disabled
pub fn toggle_extensions_dialog() -> Result<(), Box<dyn Error>> {
let extensions = ExtensionManager::get_all()?;
if extensions.is_empty() {
cliclack::outro(
"No extensions configured yet. Run configure and add some extensions first.",
)?;
return Ok(());
}
// Create a list of extension names and their enabled status
let extension_status: Vec<(String, bool)> = extensions
.iter()
.map(|entry| (entry.config.name().to_string(), entry.enabled))
.collect();
// Get currently enabled extensions for the selection
let enabled_extensions: Vec<&String> = extension_status
.iter()
.filter(|(_, enabled)| *enabled)
.map(|(name, _)| name)
.collect();
// Let user toggle extensions
let selected = cliclack::multiselect(
"enable extensions: (use \"space\" to toggle and \"enter\" to submit)",
)
.required(false)
.items(
&extension_status
.iter()
.map(|(name, _)| (name, name.as_str(), ""))
.collect::<Vec<_>>(),
)
.initial_values(enabled_extensions)
.interact()?;
// Update enabled status for each extension
for name in extension_status.iter().map(|(name, _)| name) {
ExtensionManager::set_enabled(name, selected.iter().any(|s| s.as_str() == name))?;
}
cliclack::outro("Extension settings updated successfully")?;
Ok(())
}
pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
let extension_type = cliclack::select("What type of extension would you like to add?")
.item(
"built-in",
"Built-in Extension",
"Use an extension that comes with Goose",
)
.item(
"stdio",
"Command-line Extension",
"Run a local command or script",
)
.item(
"sse",
"Remote Extension",
"Connect to a remote extension via SSE",
)
.interact()?;
match extension_type {
// TODO we'll want a place to collect all these options, maybe just an enum in goose-mcp
"built-in" => {
let extension = cliclack::select("Which built-in extension would you like to enable?")
.item(
"developer",
"Developer Tools",
"Code editing and shell access",
)
.item(
"nondeveloper",
"Non Developer",
"AI driven scripting for non developers",
)
.item(
"google_drive",
"Google Drive",
"Search and read content from google drive - additional config required",
)
.item(
"memory",
"Memory",
"Tools to save and retrieve durable memories",
)
.item("jetbrains", "JetBrains", "Connect to jetbrains IDEs")
.interact()?
.to_string();
ExtensionManager::set(ExtensionEntry {
enabled: true,
config: ExtensionConfig::Builtin {
name: extension.clone(),
},
})?;
cliclack::outro(format!("Enabled {} extension", style(extension).green()))?;
}
"stdio" => {
let extensions = ExtensionManager::get_all_names()?;
let name: String = cliclack::input("What would you like to call this extension?")
.placeholder("my-extension")
.validate(move |input: &String| {
if input.is_empty() {
Err("Please enter a name")
} else if extensions.contains(input) {
Err("An extension with this name already exists")
} else {
Ok(())
}
})
.interact()?;
let command_str: String = cliclack::input("What command should be run?")
.placeholder("npx -y @block/gdrive")
.validate(|input: &String| {
if input.is_empty() {
Err("Please enter a command")
} else {
Ok(())
}
})
.interact()?;
// Split the command string into command and args
let mut parts = command_str.split_whitespace();
let cmd = parts.next().unwrap_or("").to_string();
let args: Vec<String> = parts.map(String::from).collect();
let add_env =
cliclack::confirm("Would you like to add environment variables?").interact()?;
let mut envs = HashMap::new();
if add_env {
loop {
let key: String = cliclack::input("Environment variable name:")
.placeholder("API_KEY")
.interact()?;
let value: String = cliclack::password("Environment variable value:")
.mask('▪')
.interact()?;
envs.insert(key, value);
if !cliclack::confirm("Add another environment variable?").interact()? {
break;
}
}
}
ExtensionManager::set(ExtensionEntry {
enabled: true,
config: ExtensionConfig::Stdio {
name: name.clone(),
cmd,
args,
envs: Envs::new(envs),
},
})?;
cliclack::outro(format!("Added {} extension", style(name).green()))?;
}
"sse" => {
let extensions = ExtensionManager::get_all_names()?;
let name: String = cliclack::input("What would you like to call this extension?")
.placeholder("my-remote-extension")
.validate(move |input: &String| {
if input.is_empty() {
Err("Please enter a name")
} else if extensions.contains(input) {
Err("An extension with this name already exists")
} else {
Ok(())
}
})
.interact()?;
let uri: String = cliclack::input("What is the SSE endpoint URI?")
.placeholder("http://localhost:8000/events")
.validate(|input: &String| {
if input.is_empty() {
Err("Please enter a URI")
} else if !input.starts_with("http") {
Err("URI should start with http:// or https://")
} else {
Ok(())
}
})
.interact()?;
let add_env =
cliclack::confirm("Would you like to add environment variables?").interact()?;
let mut envs = HashMap::new();
if add_env {
loop {
let key: String = cliclack::input("Environment variable name:")
.placeholder("API_KEY")
.interact()?;
let value: String = cliclack::password("Environment variable value:")
.mask('▪')
.interact()?;
envs.insert(key, value);
if !cliclack::confirm("Add another environment variable?").interact()? {
break;
}
}
}
ExtensionManager::set(ExtensionEntry {
enabled: true,
config: ExtensionConfig::Sse {
name: name.clone(),
uri,
envs: Envs::new(envs),
},
})?;
cliclack::outro(format!("Added {} extension", style(name).green()))?;
}
_ => unreachable!(),
};
Ok(())
}

View File

@@ -0,0 +1,33 @@
use anyhow::Result;
use goose_mcp::{
DeveloperRouter, GoogleDriveRouter, JetBrainsRouter, MemoryRouter, NonDeveloperRouter,
};
use mcp_server::router::RouterService;
use mcp_server::{BoundedService, ByteTransport, Server};
use tokio::io::{stdin, stdout};
pub async fn run_server(name: &str) -> Result<()> {
// Initialize logging
crate::logging::setup_logging(Some(&format!("mcp-{name}")))?;
tracing::info!("Starting MCP server");
let router: Option<Box<dyn BoundedService>> = match name {
"developer" => Some(Box::new(RouterService(DeveloperRouter::new()))),
"nondeveloper" => Some(Box::new(RouterService(NonDeveloperRouter::new()))),
"jetbrains" => Some(Box::new(RouterService(JetBrainsRouter::new()))),
"google_drive" => {
let router = GoogleDriveRouter::new().await;
Some(Box::new(RouterService(router)))
}
"memory" => Some(Box::new(RouterService(MemoryRouter::new()))),
_ => None,
};
// Create and run the server
let server = Server::new(router.unwrap_or_else(|| panic!("Unknown server requested {}", name)));
let transport = ByteTransport::new(stdin(), stdout());
tracing::info!("Server initialized and ready to handle requests");
Ok(server.run(transport).await?)
}

View File

@@ -0,0 +1,5 @@
pub mod agent_version;
pub mod configure;
pub mod mcp;
pub mod session;
pub mod version;

View File

@@ -0,0 +1,178 @@
use rand::{distributions::Alphanumeric, Rng};
use std::process;
use crate::prompt::rustyline::RustylinePrompt;
use crate::session::{ensure_session_dir, get_most_recent_session, Session};
use console::style;
use goose::agents::extension::{Envs, ExtensionError};
use goose::agents::AgentFactory;
use goose::config::{Config, ExtensionConfig, ExtensionManager};
use goose::providers::create;
use std::path::Path;
use mcp_client::transport::Error as McpClientError;
pub async fn build_session(
name: Option<String>,
resume: bool,
extension: Option<String>,
builtin: Option<String>,
) -> Session<'static> {
// Load config and get provider/model
let config = Config::global();
let provider_name: String = config
.get("GOOSE_PROVIDER")
.expect("No provider configured. Run 'goose configure' first");
let session_dir = ensure_session_dir().expect("Failed to create session directory");
let model: String = config
.get("GOOSE_MODEL")
.expect("No model configured. Run 'goose configure' first");
let model_config = goose::model::ModelConfig::new(model.clone());
let provider = create(&provider_name, model_config).expect("Failed to create provider");
// Create the agent
let agent_version: Option<String> = config.get("GOOSE_AGENT").ok();
let mut agent = match agent_version {
Some(version) => AgentFactory::create(&version, provider),
None => AgentFactory::create(AgentFactory::default_version(), provider),
}
.expect("Failed to create agent");
// Setup extensions for the agent
for extension in ExtensionManager::get_all().expect("should load extensions") {
if extension.enabled {
let config = extension.config.clone();
agent
.add_extension(config.clone())
.await
.unwrap_or_else(|e| {
let err = match e {
ExtensionError::Transport(McpClientError::StdioProcessError(inner)) => {
inner
}
_ => e.to_string(),
};
println!("Failed to start extension: {}, {:?}", config.name(), err);
println!(
"Please check extension configuration for {}.",
config.name()
);
process::exit(1);
});
}
}
// Add extension if provided
if let Some(extension_str) = extension {
let mut parts: Vec<&str> = extension_str.split_whitespace().collect();
let mut envs = std::collections::HashMap::new();
// Parse environment variables (format: KEY=value)
while let Some(part) = parts.first() {
if !part.contains('=') {
break;
}
let env_part = parts.remove(0);
let (key, value) = env_part.split_once('=').unwrap();
envs.insert(key.to_string(), value.to_string());
}
if parts.is_empty() {
eprintln!("No command provided in extension string");
process::exit(1);
}
let cmd = parts.remove(0).to_string();
//this is an ephemeral extension so name does not matter
let name = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(8)
.map(char::from)
.collect();
let config = ExtensionConfig::Stdio {
name,
cmd,
args: parts.iter().map(|s| s.to_string()).collect(),
envs: Envs::new(envs),
};
agent.add_extension(config).await.unwrap_or_else(|e| {
eprintln!("Failed to start extension: {}", e);
process::exit(1);
});
}
// Add builtin extension if provided
if let Some(name) = builtin {
let config = ExtensionConfig::Builtin { name };
agent.add_extension(config).await.unwrap_or_else(|e| {
eprintln!("Failed to start builtin extension: {}", e);
process::exit(1);
});
}
// If resuming, try to find the session
if resume {
if let Some(ref session_name) = name {
// Try to resume specific session
let session_file = session_dir.join(format!("{}.jsonl", session_name));
if session_file.exists() {
let prompt = Box::new(RustylinePrompt::new());
return Session::new(agent, prompt, session_file);
} else {
eprintln!("Session '{}' not found, starting new session", session_name);
}
} else {
// Try to resume most recent session
if let Ok(session_file) = get_most_recent_session() {
let prompt = Box::new(RustylinePrompt::new());
return Session::new(agent, prompt, session_file);
} else {
eprintln!("No previous sessions found, starting new session");
}
}
}
// Generate session name if not provided
let name = name.unwrap_or_else(|| {
rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(8)
.map(char::from)
.collect()
});
let session_file = session_dir.join(format!("{}.jsonl", name));
if session_file.exists() {
eprintln!("Session '{}' already exists", name);
process::exit(1);
}
let prompt = Box::new(RustylinePrompt::new());
display_session_info(resume, &provider_name, &model, &session_file);
Session::new(agent, prompt, session_file)
}
fn display_session_info(resume: bool, provider: &str, model: &str, session_file: &Path) {
let start_session_msg = if resume {
"resuming session |"
} else {
"starting session |"
};
println!(
"{} {} {} {} {}",
style(start_session_msg).dim(),
style("provider:").dim(),
style(provider).cyan().dim(),
style("model:").dim(),
style(model).cyan().dim(),
);
println!(
" {} {}",
style("logging to").dim(),
style(session_file.display()).dim().cyan(),
);
}

View File

@@ -0,0 +1,3 @@
pub fn print_version() {
println!(env!("CARGO_PKG_VERSION"))
}

View File

@@ -0,0 +1,93 @@
use goose::providers::base::ProviderUsage;
#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct SessionLog {
session_file: String,
usage: Vec<ProviderUsage>,
}
pub fn log_usage(session_file: String, usage: Vec<ProviderUsage>) {
let log = SessionLog {
session_file,
usage,
};
// Ensure log directory exists
if let Some(home_dir) = dirs::home_dir() {
let log_dir = home_dir.join(".config").join("goose").join("logs");
if let Err(e) = std::fs::create_dir_all(&log_dir) {
eprintln!("Failed to create log directory: {}", e);
return;
}
let log_file = log_dir.join("goose.log");
let serialized = match serde_json::to_string(&log) {
Ok(s) => s,
Err(e) => {
eprintln!("Failed to serialize usage log: {}", e);
return;
}
};
// Append to log file
if let Err(e) = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(log_file)
.and_then(|mut file| {
std::io::Write::write_all(&mut file, serialized.as_bytes())?;
std::io::Write::write_all(&mut file, b"\n")?;
Ok(())
})
{
eprintln!("Failed to write to usage log file: {}", e);
}
} else {
eprintln!("Failed to write to usage log file: Failed to determine home directory");
}
}
#[cfg(test)]
mod tests {
use goose::providers::base::{ProviderUsage, Usage};
use crate::{
log_usage::{log_usage, SessionLog},
test_helpers::run_with_tmp_dir,
};
#[test]
fn test_session_logging() {
run_with_tmp_dir(|| {
let home_dir = dirs::home_dir().unwrap();
let log_file = home_dir
.join(".config")
.join("goose")
.join("logs")
.join("goose.log");
log_usage(
"path.txt".to_string(),
vec![ProviderUsage::new(
"model".to_string(),
Usage::new(Some(10), Some(20), Some(30)),
)],
);
// Check if log file exists and contains the expected content
assert!(log_file.exists(), "Log file should exist");
let log_content = std::fs::read_to_string(&log_file).unwrap();
let log: SessionLog = serde_json::from_str(&log_content).unwrap();
assert!(log.session_file.contains("path.txt"));
assert_eq!(log.usage[0].usage.input_tokens, Some(10));
assert_eq!(log.usage[0].usage.output_tokens, Some(20));
assert_eq!(log.usage[0].usage.total_tokens, Some(30));
assert_eq!(log.usage[0].model, "model");
// Remove the log file after test
std::fs::remove_file(&log_file).ok();
})
}
}

View File

@@ -0,0 +1,237 @@
use anyhow::{Context, Result};
use std::fs;
use std::path::PathBuf;
use tracing_appender::rolling::Rotation;
use tracing_subscriber::{
filter::LevelFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer,
Registry,
};
use goose::tracing::langfuse_layer;
/// Returns the directory where log files should be stored.
/// Creates the directory structure if it doesn't exist.
fn get_log_directory() -> Result<PathBuf> {
let home = std::env::var("HOME").context("HOME environment variable not set")?;
let base_log_dir = PathBuf::from(home)
.join(".config")
.join("goose")
.join("logs")
.join("cli"); // Add cli-specific subdirectory
// Create date-based subdirectory
let now = chrono::Local::now();
let date_dir = base_log_dir.join(now.format("%Y-%m-%d").to_string());
// Ensure log directory exists
fs::create_dir_all(&date_dir).context("Failed to create log directory")?;
Ok(date_dir)
}
/// Sets up the logging infrastructure for the application.
/// This includes:
/// - File-based logging with JSON formatting (DEBUG level)
/// - Console output for development (INFO level)
/// - Optional Langfuse integration (DEBUG level)
pub fn setup_logging(name: Option<&str>) -> Result<()> {
// Set up file appender for goose module logs
let log_dir = get_log_directory()?;
let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S").to_string();
// Create log file name by prefixing with timestamp
let log_filename = if name.is_some() {
format!("{}-{}.log", timestamp, name.unwrap())
} else {
format!("{}.log", timestamp)
};
// Create non-rolling file appender for detailed logs
let file_appender =
tracing_appender::rolling::RollingFileAppender::new(Rotation::NEVER, log_dir, log_filename);
// Create JSON file logging layer with all logs (DEBUG and above)
let file_layer = fmt::layer()
.with_target(true)
.with_level(true)
.with_writer(file_appender)
.with_ansi(false)
.with_file(true)
.pretty();
// Create console logging layer for development - INFO and above only
let console_layer = fmt::layer()
.with_target(true)
.with_level(true)
.with_ansi(true)
.with_file(true)
.with_line_number(true)
.pretty();
// Base filter
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
// Set default levels for different modules
EnvFilter::new("")
// Set mcp-server module to DEBUG
.add_directive("mcp_server=debug".parse().unwrap())
// Set mcp-client to DEBUG
.add_directive("mcp_client=debug".parse().unwrap())
// Set goose module to DEBUG
.add_directive("goose=debug".parse().unwrap())
// Set goose-cli to INFO
.add_directive("goose_cli=info".parse().unwrap())
// Set everything else to WARN
.add_directive(LevelFilter::WARN.into())
});
// Build the subscriber with required layers
let subscriber = Registry::default()
.with(file_layer.with_filter(env_filter)) // Gets all logs
.with(console_layer.with_filter(LevelFilter::WARN)); // Controls log levels
// Initialize with Langfuse if available
if let Some(langfuse) = langfuse_layer::create_langfuse_observer() {
subscriber
.with(langfuse.with_filter(LevelFilter::DEBUG))
.try_init()
.context("Failed to set global subscriber")?;
} else {
subscriber
.try_init()
.context("Failed to set global subscriber")?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::env;
use tempfile::TempDir;
use test_case::test_case;
use tokio::runtime::Runtime;
fn setup_temp_home() -> TempDir {
let temp_dir = TempDir::new().unwrap();
env::set_var("HOME", temp_dir.path());
temp_dir
}
#[test]
fn test_log_directory_creation() {
let _temp_dir = setup_temp_home();
let log_dir = get_log_directory().unwrap();
assert!(log_dir.exists());
assert!(log_dir.is_dir());
// Verify directory structure
let path_components: Vec<_> = log_dir.components().collect();
assert!(path_components.iter().any(|c| c.as_os_str() == "goose"));
assert!(path_components.iter().any(|c| c.as_os_str() == "logs"));
assert!(path_components.iter().any(|c| c.as_os_str() == "cli"));
}
#[test_case(Some("test_session") ; "with session name")]
#[test_case(None ; "without session name")]
fn test_log_file_name(session_name: Option<&str>) {
let _rt = Runtime::new().unwrap();
let _temp_dir = setup_temp_home();
// Create a test-specific log directory and file
let log_dir = get_log_directory().unwrap();
let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S").to_string();
let file_name = format!("{}.log", session_name.unwrap_or(&timestamp));
// Create the log file
let file_path = log_dir.join(&file_name);
fs::write(&file_path, "test").unwrap();
// Verify the file exists and has the correct name
let entries = fs::read_dir(log_dir).unwrap();
let log_files: Vec<_> = entries
.filter_map(Result::ok)
.filter(|e| e.path().extension().map_or(false, |ext| ext == "log"))
.collect();
assert_eq!(log_files.len(), 1, "Expected exactly one log file");
let log_file_name = log_files[0].file_name().to_string_lossy().into_owned();
println!("Log file name: {}", log_file_name);
if let Some(name) = session_name {
assert_eq!(log_file_name, format!("{}.log", name));
} else {
// Extract just the filename without extension for comparison
let name_without_ext = log_file_name.trim_end_matches(".log");
// Verify it's a valid timestamp format
assert_eq!(
name_without_ext.len(),
15,
"Expected 15 characters (YYYYMMDD_HHMMSS)"
);
assert!(
name_without_ext[8..9].contains('_'),
"Expected underscore at position 8"
);
assert!(
name_without_ext
.chars()
.all(|c| c.is_ascii_digit() || c == '_'),
"Expected only digits and underscore"
);
}
}
#[tokio::test]
async fn test_langfuse_layer_creation() {
let _temp_dir = setup_temp_home();
// Store original environment variables (both sets)
let original_vars = [
("LANGFUSE_PUBLIC_KEY", env::var("LANGFUSE_PUBLIC_KEY").ok()),
("LANGFUSE_SECRET_KEY", env::var("LANGFUSE_SECRET_KEY").ok()),
("LANGFUSE_HOST", env::var("LANGFUSE_HOST").ok()),
(
"LANGFUSE_INIT_PROJECT_PUBLIC_KEY",
env::var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY").ok(),
),
(
"LANGFUSE_INIT_PROJECT_SECRET_KEY",
env::var("LANGFUSE_INIT_PROJECT_SECRET_KEY").ok(),
),
];
// Clear all Langfuse environment variables
for (var, _) in &original_vars {
env::remove_var(var);
}
// Test without any environment variables
assert!(langfuse_layer::create_langfuse_observer().is_none());
// Test with standard Langfuse variables
env::set_var("LANGFUSE_PUBLIC_KEY", "test_public_key");
env::set_var("LANGFUSE_SECRET_KEY", "test_secret_key");
assert!(langfuse_layer::create_langfuse_observer().is_some());
// Clear and test with init project variables
env::remove_var("LANGFUSE_PUBLIC_KEY");
env::remove_var("LANGFUSE_SECRET_KEY");
env::set_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY", "test_public_key");
env::set_var("LANGFUSE_INIT_PROJECT_SECRET_KEY", "test_secret_key");
assert!(langfuse_layer::create_langfuse_observer().is_some());
// Test fallback behavior
env::remove_var("LANGFUSE_INIT_PROJECT_PUBLIC_KEY");
assert!(langfuse_layer::create_langfuse_observer().is_none());
// Restore original environment variables
for (var, value) in original_vars {
match value {
Some(val) => env::set_var(var, val),
None => env::remove_var(var),
}
}
}
}

View File

@@ -0,0 +1,234 @@
use anyhow::Result;
use clap::{CommandFactory, Parser, Subcommand};
mod commands;
mod log_usage;
mod logging;
mod prompt;
mod session;
use commands::agent_version::AgentCommand;
use commands::configure::handle_configure;
use commands::mcp::run_server;
use commands::session::build_session;
use commands::version::print_version;
use console::style;
use goose::config::Config;
use logging::setup_logging;
use std::io::{self, Read};
#[cfg(test)]
mod test_helpers;
#[derive(Parser)]
#[command(author, about, long_about = None)]
struct Cli {
#[arg(short = 'v', long = "version")]
version: bool,
#[command(subcommand)]
command: Option<Command>,
}
#[derive(Subcommand)]
enum Command {
/// Configure Goose settings
#[command(about = "Configure Goose settings")]
Configure {},
/// Manage system prompts and behaviors
#[command(about = "Run one of the mcp servers bundled with goose")]
Mcp { name: String },
/// Start or resume interactive chat sessions
#[command(about = "Start or resume interactive chat sessions", alias = "s")]
Session {
/// Name for the chat session
#[arg(
short,
long,
value_name = "NAME",
help = "Name for the chat session (e.g., 'project-x')",
long_help = "Specify a name for your chat session. When used with --resume, will resume this specific session if it exists."
)]
name: Option<String>,
/// Resume a previous session
#[arg(
short,
long,
help = "Resume a previous session (last used or specified by --session)",
long_help = "Continue from a previous chat session. If --session is provided, resumes that specific session. Otherwise resumes the last used session."
)]
resume: bool,
/// Add a stdio extension with environment variables and command
#[arg(
long = "with-extension",
value_name = "COMMAND",
help = "Add a stdio extension (e.g., 'GITHUB_TOKEN=xyz npx -y @modelcontextprotocol/server-github')",
long_help = "Add a stdio extension from a full command with environment variables. Format: 'ENV1=val1 ENV2=val2 command args...'"
)]
extension: Option<String>,
/// Add a builtin extension by name
#[arg(
long = "with-builtin",
value_name = "NAME",
help = "Add a builtin extension by name (e.g., 'developer')",
long_help = "Add a builtin extension that is bundled with goose by specifying its name"
)]
builtin: Option<String>,
},
/// Execute commands from an instruction file
#[command(about = "Execute commands from an instruction file or stdin")]
Run {
/// Path to instruction file containing commands
#[arg(
short,
long,
value_name = "FILE",
help = "Path to instruction file containing commands",
conflicts_with = "input_text"
)]
instructions: Option<String>,
/// Input text containing commands
#[arg(
short = 't',
long = "text",
value_name = "TEXT",
help = "Input text to provide to Goose directly",
long_help = "Input text containing commands for Goose. Use this in lieu of the instructions argument.",
conflicts_with = "instructions"
)]
input_text: Option<String>,
/// Name for this run session
#[arg(
short,
long,
value_name = "NAME",
help = "Name for this run session (e.g., 'daily-tasks')",
long_help = "Specify a name for this run session. This helps identify and resume specific runs later."
)]
name: Option<String>,
/// Resume a previous run
#[arg(
short,
long,
action = clap::ArgAction::SetTrue,
help = "Resume from a previous run",
long_help = "Continue from a previous run, maintaining the execution state and context."
)]
resume: bool,
/// Add a stdio extension with environment variables and command
#[arg(
long = "with-extension",
value_name = "COMMAND",
help = "Add a stdio extension with environment variables and command (e.g., 'GITHUB_TOKEN=xyz npx -y @modelcontextprotocol/server-github')",
long_help = "Add a stdio extension with environment variables and command. Format: 'ENV1=val1 ENV2=val2 command args...'"
)]
extension: Option<String>,
/// Add a builtin extension by name
#[arg(
long = "with-builtin",
value_name = "NAME",
help = "Add a builtin extension by name (e.g., 'developer')",
long_help = "Add a builtin extension that is compiled into goose by specifying its name"
)]
builtin: Option<String>,
},
/// List available agent versions
Agents(AgentCommand),
}
#[derive(clap::ValueEnum, Clone, Debug)]
enum CliProviderVariant {
OpenAi,
Databricks,
Ollama,
}
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
if cli.version {
print_version();
return Ok(());
}
match cli.command {
Some(Command::Configure {}) => {
let _ = handle_configure().await;
return Ok(());
}
Some(Command::Mcp { name }) => {
let _ = run_server(&name).await;
}
Some(Command::Session {
name,
resume,
extension,
builtin,
}) => {
let mut session = build_session(name, resume, extension, builtin).await;
setup_logging(session.session_file().file_stem().and_then(|s| s.to_str()))?;
let _ = session.start().await;
return Ok(());
}
Some(Command::Run {
instructions,
input_text,
name,
resume,
extension,
builtin,
}) => {
// Validate that we have some input source
if instructions.is_none() && input_text.is_none() {
eprintln!("Error: Must provide either --instructions or --text");
std::process::exit(1);
}
let contents = if let Some(file_name) = instructions {
let file_path = std::path::Path::new(&file_name);
std::fs::read_to_string(file_path).expect("Failed to read the instruction file")
} else if let Some(input_text) = input_text {
input_text
} else {
let mut stdin = String::new();
io::stdin()
.read_to_string(&mut stdin)
.expect("Failed to read from stdin");
stdin
};
let mut session = build_session(name, resume, extension, builtin).await;
let _ = session.headless_start(contents.clone()).await;
return Ok(());
}
Some(Command::Agents(cmd)) => {
cmd.run()?;
return Ok(());
}
None => {
Cli::command().print_help()?;
println!();
if !Config::global().exists() {
println!(
"\n {}: Run '{}' to setup goose for the first time",
style("Tip").green().italic(),
style("goose configure").cyan()
);
}
}
}
Ok(())
}

View File

@@ -0,0 +1,38 @@
use anyhow::Result;
use goose::message::Message;
pub mod renderer;
pub mod rustyline;
pub mod thinking;
pub trait Prompt {
fn render(&mut self, message: Box<Message>);
fn get_input(&mut self) -> Result<Input>;
fn show_busy(&mut self);
fn hide_busy(&self);
fn close(&self);
/// Load the user's message history into the prompt for command history navigation. First message is the oldest message.
/// When history is supported by the prompt.
fn load_user_message_history(&mut self, _messages: Vec<Message>) {}
fn goose_ready(&self) {
println!("\n");
println!("Goose is running! Enter your instructions, or try asking what goose can do.");
println!("\n");
}
}
pub struct Input {
pub input_type: InputType,
pub content: Option<String>, // Optional content as sometimes the user may be issuing a command eg. (Exit)
}
pub enum InputType {
AskAgain, // Ask the user for input again. Control flow command.
Message, // User sent a message
Exit, // User wants to exit the session
}
pub enum Theme {
Light,
Dark,
}

View File

@@ -0,0 +1,407 @@
use std::collections::HashMap;
use std::io::{self, Write};
use std::path::PathBuf;
use bat::WrappingMode;
use console::style;
use goose::message::{Message, MessageContent, ToolRequest, ToolResponse};
use mcp_core::role::Role;
use mcp_core::{content::Content, tool::ToolCall};
use serde_json::Value;
use super::Theme;
const MAX_STRING_LENGTH: usize = 40;
const MAX_PATH_LENGTH: usize = 60;
const INDENT: &str = " ";
/// Shortens a path string by abbreviating directory names while keeping the last two components intact.
/// If the path starts with the user's home directory, it will be replaced with ~.
///
/// # Examples
/// ```
/// let path = "/Users/alice/Development/very/long/path/to/file.txt";
/// assert_eq!(
/// shorten_path(path),
/// "~/D/v/l/p/to/file.txt"
/// );
/// ```
fn shorten_path(path: &str) -> String {
let path = PathBuf::from(path);
// First try to convert to ~ if it's in home directory
let home = dirs::home_dir();
let path_str = if let Some(home) = home {
if let Ok(stripped) = path.strip_prefix(home) {
format!("~/{}", stripped.display())
} else {
path.display().to_string()
}
} else {
path.display().to_string()
};
// If path is already short enough, return as is
if path_str.len() <= MAX_PATH_LENGTH {
return path_str;
}
let parts: Vec<_> = path_str.split('/').collect();
// If we have 3 or fewer parts, return as is
if parts.len() <= 3 {
return path_str;
}
// Keep the first component (empty string before root / or ~) and last two components intact
let mut shortened = vec![parts[0].to_string()];
// Shorten middle components to their first letter
for component in &parts[1..parts.len() - 2] {
if !component.is_empty() {
shortened.push(component.chars().next().unwrap_or('?').to_string());
}
}
// Add the last two components
shortened.push(parts[parts.len() - 2].to_string());
shortened.push(parts[parts.len() - 1].to_string());
shortened.join("/")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_shorten_path() {
// Test a long path without home directory
let long_path = "/Users/test/Development/this/is/a/very/long/nested/deeply/example.txt";
let shortened = shorten_path(long_path);
assert!(
shortened.len() < long_path.len(),
"Shortened path '{}' should be shorter than original '{}'",
shortened,
long_path
);
assert!(
shortened.ends_with("deeply/example.txt"),
"Shortened path '{}' should end with 'deeply/example.txt'",
shortened
);
// Test a short path (shouldn't be modified)
assert_eq!(shorten_path("/usr/local/bin"), "/usr/local/bin");
// Test path with less than 3 components
assert_eq!(shorten_path("/usr/local"), "/usr/local");
}
}
/// Implement the ToolRenderer trait for each tool that you want to render in the prompt.
pub trait ToolRenderer: ToolRendererClone {
fn tool_name(&self) -> String;
fn request(&self, tool_request: &ToolRequest, theme: &str);
fn response(&self, tool_response: &ToolResponse, theme: &str);
}
// Helper trait for cloning boxed ToolRenderer objects
pub trait ToolRendererClone {
fn clone_box(&self) -> Box<dyn ToolRenderer>;
}
// Implement the helper trait for any type that implements ToolRenderer and Clone
impl<T> ToolRendererClone for T
where
T: 'static + ToolRenderer + Clone,
{
fn clone_box(&self) -> Box<dyn ToolRenderer> {
Box::new(self.clone())
}
}
// Make Box<dyn ToolRenderer> clonable
impl Clone for Box<dyn ToolRenderer> {
fn clone(&self) -> Box<dyn ToolRenderer> {
self.clone_box()
}
}
#[derive(Clone)]
pub struct DefaultRenderer;
impl ToolRenderer for DefaultRenderer {
fn tool_name(&self) -> String {
"default".to_string()
}
fn request(&self, tool_request: &ToolRequest, theme: &str) {
match &tool_request.tool_call {
Ok(call) => {
default_print_request_header(call);
// Format and print the parameters
print_params(&call.arguments, 0);
print_newline();
}
Err(e) => print_markdown(&e.to_string(), theme),
}
}
fn response(&self, tool_response: &ToolResponse, theme: &str) {
default_response_renderer(tool_response, theme);
}
}
#[derive(Clone)]
pub struct TextEditorRenderer;
impl ToolRenderer for TextEditorRenderer {
fn tool_name(&self) -> String {
"developer__text_editor".to_string()
}
fn request(&self, tool_request: &ToolRequest, theme: &str) {
match &tool_request.tool_call {
Ok(call) => {
default_print_request_header(call);
// Print path first with special formatting
if let Some(Value::String(path)) = call.arguments.get("path") {
println!(
"{}: {}",
style("path").dim(),
style(shorten_path(path)).green()
);
}
// Print other arguments normally, excluding path
if let Some(args) = call.arguments.as_object() {
let mut other_args = serde_json::Map::new();
for (k, v) in args {
if k != "path" {
other_args.insert(k.clone(), v.clone());
}
}
print_params(&Value::Object(other_args), 0);
}
print_newline();
}
Err(e) => print_markdown(&e.to_string(), theme),
}
}
fn response(&self, tool_response: &ToolResponse, theme: &str) {
default_response_renderer(tool_response, theme);
}
}
#[derive(Clone)]
pub struct BashDeveloperExtensionRenderer;
impl ToolRenderer for BashDeveloperExtensionRenderer {
fn tool_name(&self) -> String {
"developer__shell".to_string()
}
fn request(&self, tool_request: &ToolRequest, theme: &str) {
match &tool_request.tool_call {
Ok(call) => {
default_print_request_header(call);
match call.arguments.get("command") {
Some(Value::String(s)) => {
println!("{}: {}", style("command").dim(), style(s).green());
}
_ => print_params(&call.arguments, 0),
}
print_newline();
}
Err(e) => print_markdown(&e.to_string(), theme),
}
}
fn response(&self, tool_response: &ToolResponse, theme: &str) {
default_response_renderer(tool_response, theme);
}
}
pub fn render(message: &Message, theme: &Theme, renderers: HashMap<String, Box<dyn ToolRenderer>>) {
let theme = match theme {
Theme::Light => "GitHub",
Theme::Dark => "zenburn",
};
let mut last_tool_name: &str = "default";
for message_content in &message.content {
match message_content {
MessageContent::Text(text) => print_markdown(&text.text, theme),
MessageContent::ToolRequest(tool_request) => match &tool_request.tool_call {
Ok(call) => {
last_tool_name = &call.name;
renderers
.get(&call.name)
.or_else(|| renderers.get("default"))
.unwrap()
.request(tool_request, theme);
}
Err(_) => renderers
.get("default")
.unwrap()
.request(tool_request, theme),
},
MessageContent::ToolResponse(tool_response) => renderers
.get(last_tool_name)
.or_else(|| renderers.get("default"))
.unwrap()
.response(tool_response, theme),
MessageContent::Image(image) => {
println!("Image: [data: {}, type: {}]", image.data, image.mime_type);
}
}
}
print_newline();
io::stdout().flush().expect("Failed to flush stdout");
}
pub fn default_response_renderer(tool_response: &ToolResponse, theme: &str) {
match &tool_response.tool_result {
Ok(contents) => {
for content in contents {
if content
.audience()
.is_some_and(|audience| !audience.contains(&Role::User))
{
continue;
}
let min_priority = std::env::var("GOOSE_CLI_MIN_PRIORITY")
.ok()
.and_then(|val| val.parse::<f32>().ok())
.unwrap_or(0.0);
// if priority is not set OR less than or equal to min_priority, do not render
if content
.priority()
.is_some_and(|priority| priority <= min_priority)
|| content.priority().is_none()
{
continue;
}
if let Content::Text(text) = content {
print_markdown(&text.text, theme);
}
}
}
Err(e) => print_markdown(&e.to_string(), theme),
}
}
pub fn default_print_request_header(call: &ToolCall) {
// Print the tool name with an emoji
// use rsplit to handle any prefixed tools with more underscores
// unicode gets converted to underscores during sanitization
let parts: Vec<_> = call.name.rsplit("__").collect();
let tool_header = format!(
"─── {} | {} ──────────────────────────",
style(parts.first().unwrap_or(&"unknown")),
style(
parts
.split_first()
// client name is the rest of the split, reversed
// reverse the iterator and re-join on __
.map(|(_, s)| s.iter().rev().copied().collect::<Vec<_>>().join("__"))
.unwrap_or_else(|| "unknown".to_string())
)
.magenta()
.dim(),
);
print_newline();
println!("{}", tool_header);
}
pub fn print_markdown(content: &str, theme: &str) {
bat::PrettyPrinter::new()
.input(bat::Input::from_bytes(content.as_bytes()))
.theme(theme)
.language("Markdown")
.wrapping_mode(WrappingMode::Character)
.print()
.unwrap();
}
/// Format and print parameters recursively with proper indentation and colors
pub fn print_params(value: &Value, depth: usize) {
let indent = INDENT.repeat(depth);
match value {
Value::Object(map) => {
for (key, val) in map {
match val {
Value::Object(_) => {
println!("{}{}:", indent, style(key).dim());
print_params(val, depth + 1);
}
Value::Array(arr) => {
println!("{}{}:", indent, style(key).dim());
for item in arr.iter() {
println!("{}{}- ", indent, INDENT);
print_params(item, depth + 2);
}
}
Value::String(s) => {
if s.len() > MAX_STRING_LENGTH {
println!("{}{}: {}", indent, style(key).dim(), style("...").dim());
} else {
println!("{}{}: {}", indent, style(key).dim(), style(s).green());
}
}
Value::Number(n) => {
println!("{}{}: {}", indent, style(key).dim(), style(n).blue());
}
Value::Bool(b) => {
println!("{}{}: {}", indent, style(key).dim(), style(b).blue());
}
Value::Null => {
println!("{}{}: {}", indent, style(key).dim(), style("null").dim());
}
}
}
}
Value::Array(arr) => {
for (i, item) in arr.iter().enumerate() {
println!("{}{}.", indent, i + 1);
print_params(item, depth + 1);
}
}
Value::String(s) => {
if s.len() > MAX_STRING_LENGTH {
println!(
"{}{}",
indent,
style(format!("[REDACTED: {} chars]", s.len())).yellow()
);
} else {
println!("{}{}", indent, style(s).green());
}
}
Value::Number(n) => {
println!("{}{}", indent, style(n).yellow());
}
Value::Bool(b) => {
println!("{}{}", indent, style(b).yellow());
}
Value::Null => {
println!("{}{}", indent, style("null").dim());
}
}
}
pub fn print_newline() {
println!();
}

View File

@@ -0,0 +1,167 @@
use std::collections::HashMap;
use super::{
renderer::{
render, BashDeveloperExtensionRenderer, DefaultRenderer, TextEditorRenderer, ToolRenderer,
},
thinking::get_random_thinking_message,
Input, InputType, Prompt, Theme,
};
use anyhow::Result;
use cliclack::spinner;
use goose::message::Message;
use mcp_core::Role;
use rustyline::{DefaultEditor, EventHandler, KeyCode, KeyEvent, Modifiers};
const PROMPT: &str = "\x1b[1m\x1b[38;5;30m( O)> \x1b[0m";
pub struct RustylinePrompt {
spinner: cliclack::ProgressBar,
theme: Theme,
renderers: HashMap<String, Box<dyn ToolRenderer>>,
editor: DefaultEditor,
}
impl RustylinePrompt {
pub fn new() -> Self {
let mut renderers: HashMap<String, Box<dyn ToolRenderer>> = HashMap::new();
let default_renderer = DefaultRenderer;
renderers.insert(default_renderer.tool_name(), Box::new(default_renderer));
let bash_dev_extension_renderer = BashDeveloperExtensionRenderer;
renderers.insert(
bash_dev_extension_renderer.tool_name(),
Box::new(bash_dev_extension_renderer),
);
let text_editor_renderer = TextEditorRenderer;
renderers.insert(
text_editor_renderer.tool_name(),
Box::new(text_editor_renderer),
);
let mut editor = DefaultEditor::new().expect("Failed to create editor");
editor.bind_sequence(
KeyEvent(KeyCode::Char('j'), Modifiers::CTRL),
EventHandler::Simple(rustyline::Cmd::Newline),
);
RustylinePrompt {
spinner: spinner(),
theme: std::env::var("GOOSE_CLI_THEME")
.ok()
.map(|val| {
if val.eq_ignore_ascii_case("light") {
Theme::Light
} else {
Theme::Dark
}
})
.unwrap_or(Theme::Dark),
renderers,
editor,
}
}
}
impl Prompt for RustylinePrompt {
fn render(&mut self, message: Box<Message>) {
render(&message, &self.theme, self.renderers.clone());
}
fn show_busy(&mut self) {
self.spinner = spinner();
self.spinner
.start(format!("{}...", get_random_thinking_message()));
}
fn hide_busy(&self) {
self.spinner.stop("");
}
fn get_input(&mut self) -> Result<Input> {
let input = self.editor.readline(PROMPT);
let mut message_text = match input {
Ok(text) => {
// Add valid input to history
if let Err(e) = self.editor.add_history_entry(text.as_str()) {
eprintln!("Failed to add to history: {}", e);
}
text
}
Err(e) => {
match e {
rustyline::error::ReadlineError::Interrupted => (),
_ => eprintln!("Input error: {}", e),
}
return Ok(Input {
input_type: InputType::Exit,
content: None,
});
}
};
message_text = message_text.trim().to_string();
if message_text.eq_ignore_ascii_case("/exit")
|| message_text.eq_ignore_ascii_case("/quit")
|| message_text.eq_ignore_ascii_case("exit")
|| message_text.eq_ignore_ascii_case("quit")
{
Ok(Input {
input_type: InputType::Exit,
content: None,
})
} else if message_text.eq_ignore_ascii_case("/t") {
self.theme = match self.theme {
Theme::Light => {
println!("Switching to Dark theme");
Theme::Dark
}
Theme::Dark => {
println!("Switching to Light theme");
Theme::Light
}
};
return Ok(Input {
input_type: InputType::AskAgain,
content: None,
});
} else if message_text.eq_ignore_ascii_case("/?")
|| message_text.eq_ignore_ascii_case("/help")
{
println!("Commands:");
println!("/exit - Exit the session");
println!("/t - Toggle Light/Dark theme");
println!("/? | /help - Display this help message");
println!("Ctrl+C - Interrupt goose (resets the interaction to before the interrupted user request)");
println!("Ctrl+j - Adds a newline");
println!("Use Up/Down arrow keys to navigate through command history");
return Ok(Input {
input_type: InputType::AskAgain,
content: None,
});
} else {
return Ok(Input {
input_type: InputType::Message,
content: Some(message_text.to_string()),
});
}
}
fn load_user_message_history(&mut self, messages: Vec<Message>) {
for message in messages.into_iter().filter(|m| m.role == Role::User) {
for content in message.content {
if let Some(text) = content.as_text() {
if let Err(e) = self.editor.add_history_entry(text) {
eprintln!("Failed to add to history: {}", e);
}
}
}
}
}
fn close(&self) {
// No cleanup required
}
}

View File

@@ -0,0 +1,225 @@
use rand::seq::SliceRandom;
/// Extended list of playful thinking messages including both goose and general AI actions
pub const THINKING_MESSAGES: &[&str] = &[
"Thinking",
"Thinking hard",
// Include all goose actions
"Spreading wings",
"Honking thoughtfully",
"Waddling to conclusions",
"Flapping wings excitedly",
"Preening code feathers",
"Gathering digital breadcrumbs",
"Paddling through data",
"Migrating thoughts",
"Nesting ideas",
"Squawking calculations",
"Ruffling algorithmic feathers",
"Pecking at problems",
"Stretching webbed feet",
"Foraging for solutions",
"Grooming syntax",
"Building digital nest",
"Patrolling the codebase",
"Gosling about",
"Strutting with purpose",
"Diving for answers",
"Herding bytes",
"Molting old code",
"Swimming through streams",
"Goose-stepping through logic",
"Synchronizing flock algorithms",
"Navigating code marshes",
"Incubating brilliant ideas",
"Arranging feathers recursively",
"Gliding through branches",
"Migrating to better solutions",
"Nesting functions carefully",
"Hatching clever solutions",
"Preening parse trees",
"Flying through functions",
"Gathering syntax seeds",
"Webbing connections",
"Flocking to optimizations",
"Paddling through protocols",
"Honking success signals",
"Waddling through workflows",
"Nesting in neural networks",
// AI thinking actions
"Consulting the digital oracle",
"Summoning binary spirits",
"Reticulating splines",
"Calculating meaning of life",
"Traversing neural pathways",
"Untangling spaghetti code",
"Mining thought gems",
"Defragmenting brain bits",
"Compiling wisdom",
"Debugging reality",
"Optimizing thought processes",
"Scanning parallel universes",
"Reorganizing bits and bytes",
"Calibrating neural networks",
"Charging creativity cells",
"Indexing imagination",
"Parsing possibilities",
"Buffering brilliance",
"Loading clever responses",
"Generating witty remarks",
"Synthesizing solutions",
"Applying machine learning",
"Calculating quantum states",
"Analyzing algorithms",
"Decoding human intent",
"Exploring solution space",
"Gathering computational momentum",
"Initializing clever mode",
"Juggling variables",
"Knitting neural networks",
"Learning at light speed",
"Navigating knowledge graphs",
"Orchestrating outputs",
"Pondering possibilities",
"Reading between the lines",
"Searching solution space",
"Training thought vectors",
"Unfolding understanding",
"Validating variables",
"Weaving wisdom web",
"Yielding insights",
"Zooming through zettabytes",
"Baking fresh ideas",
"Charging creativity crystals",
"Dancing with data",
"Enchanting electrons",
"Folding thought origami",
"Growing solution trees",
"Harmonizing heuristics",
"Inspiring innovations",
"Jazzing up algorithms",
"Kindling knowledge",
"Levitating logic gates",
"Manifesting solutions",
"Nurturing neural nets",
"Optimizing outcomes",
"Painting with pixels",
"Questioning bits",
"Recycling random thoughts",
"Serenading semiconductors",
"Taming tensors",
"Unlocking understanding",
"Visualizing vectors",
"Wrangling widgets",
"Yodeling yaml",
"Aligning artificial awarenesses",
"Bootstrapping brain bytes",
"Contemplating code conundrums",
"Distilling digital dreams",
"Energizing electron engines",
"Fabricating future frameworks",
"Generating genius guidelines",
"Harmonizing hardware helpers",
"Illuminating input insights",
"Kindling knowledge kernels",
"Linking logical lattices",
"Materializing memory maps",
"Navigating neural nodes",
"Orchestrating output oracles",
"Pioneering program paths",
"Quantifying quantum queries",
"Refactoring reality routines",
"Synchronizing system states",
"Transforming thought threads",
"Unifying understanding units",
"Vectorizing virtual visions",
"Weaving wisdom wavelengths",
"Yielding yaml yearnings",
"Brewing binary brilliance",
"Crafting code crystals",
"Designing data dreams",
"Encoding ethereal elements",
"Filtering function flows",
"Gathering gigabyte galaxies",
"Hashing hope hypotheses",
"Igniting innovation ions",
"Joining joy journals",
"Knitting knowledge knots",
"Launching logic loops",
"Merging memory matrices",
"Nourishing neural networks",
"Ordering output orbits",
"Processing pattern particles",
"Rendering reality rays",
"Streaming syntax stars",
"Threading thought theories",
"Updating understanding units",
"Validating virtual vectors",
"Warming wisdom waves",
"Examining electron echoes",
"Yoking yesterday yields",
"Assembling algorithm arrays",
"Balancing binary bridges",
"Calculating cosmic codes",
"Debugging dream drivers",
"Encrypting ethereal edges",
"Formatting future frames",
"Growing gradient gardens",
"Harvesting hash harmonies",
"Importing insight ions",
"Keeping kernel keys",
"Linking lambda loops",
"Mapping memory mazes",
"Normalizing neural nodes",
"Organizing output oceans",
"Parsing pattern paths",
"Sampling syntax streams",
"Testing thought threads",
"Validating virtual vectors",
"Examining electron echoes",
"Accelerating abstract algebras",
"Buffering binary bubbles",
"Caching cosmic calculations",
"Deploying digital dreams",
"Evolving ethereal entities",
"Calculating response probabilities",
"Updating knowledge graphs",
"Processing neural feedback",
"Exploring decision trees",
"Measuring semantic distance",
"Connecting synaptic pathways",
"Evaluating response options",
"Scanning memory banks",
"Simulating future outcomes",
"Adjusting confidence weights",
"Mapping context vectors",
"Balancing response parameters",
"Running inference engines",
"Optimizing memory usage",
"Merging knowledge streams",
"Calibrating response tone",
"Analyzing input patterns",
"Processing feedback loops",
"Measuring response quality",
"Scanning information matrices",
"Processing user intent",
"Measuring response coherence",
"Exploring solution paths",
"Processing context clues",
"Scanning memory circuits",
"Building response chains",
"Analyzing conversation flow",
"Processing temporal data",
"Exploring concept spaces",
"Processing memory streams",
"Evaluating logical paths",
"Building thought graphs",
"Scanning neural pathways",
];
/// Returns a random thinking message from the extended list
pub fn get_random_thinking_message() -> &'static str {
THINKING_MESSAGES
.choose(&mut rand::thread_rng())
.unwrap_or(&THINKING_MESSAGES[0])
}

View File

@@ -0,0 +1,331 @@
use anyhow::Result;
use core::panic;
use futures::StreamExt;
use std::fs::{self, File};
use std::io::{self, BufRead, Write};
use std::path::PathBuf;
use crate::log_usage::log_usage;
use crate::prompt::{InputType, Prompt};
use goose::agents::Agent;
use goose::message::{Message, MessageContent};
use mcp_core::handler::ToolError;
use mcp_core::role::Role;
// File management functions
pub fn ensure_session_dir() -> Result<PathBuf> {
let home_dir = dirs::home_dir().ok_or(anyhow::anyhow!("Could not determine home directory"))?;
let config_dir = home_dir.join(".config").join("goose").join("sessions");
if !config_dir.exists() {
fs::create_dir_all(&config_dir)?;
}
Ok(config_dir)
}
pub fn get_most_recent_session() -> Result<PathBuf> {
let session_dir = ensure_session_dir()?;
let mut entries = fs::read_dir(&session_dir)?
.filter_map(|entry| entry.ok())
.filter(|entry| entry.path().extension().is_some_and(|ext| ext == "jsonl"))
.collect::<Vec<_>>();
if entries.is_empty() {
return Err(anyhow::anyhow!("No session files found"));
}
// Sort by modification time, most recent first
entries.sort_by(|a, b| {
b.metadata()
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH)
.cmp(
&a.metadata()
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH),
)
});
Ok(entries[0].path())
}
pub fn readable_session_file(session_file: &PathBuf) -> Result<File> {
match fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(session_file)
{
Ok(file) => Ok(file),
Err(e) => Err(anyhow::anyhow!("Failed to open session file: {}", e)),
}
}
pub fn persist_messages(session_file: &PathBuf, messages: &[Message]) -> Result<()> {
let file = fs::File::create(session_file)?; // Create or truncate the file
persist_messages_internal(file, messages)
}
fn persist_messages_internal(session_file: File, messages: &[Message]) -> Result<()> {
let mut writer = std::io::BufWriter::new(session_file);
for message in messages {
serde_json::to_writer(&mut writer, &message)?;
writeln!(writer)?;
}
writer.flush()?;
Ok(())
}
pub fn deserialize_messages(file: File) -> Result<Vec<Message>> {
let reader = io::BufReader::new(file);
let mut messages = Vec::new();
for line in reader.lines() {
messages.push(serde_json::from_str::<Message>(&line?)?);
}
Ok(messages)
}
// Session management
pub struct Session<'a> {
agent: Box<dyn Agent>,
prompt: Box<dyn Prompt + 'a>,
session_file: PathBuf,
messages: Vec<Message>,
}
#[allow(dead_code)]
impl<'a> Session<'a> {
pub fn new(
agent: Box<dyn Agent>,
mut prompt: Box<dyn Prompt + 'a>,
session_file: PathBuf,
) -> Self {
let messages = match readable_session_file(&session_file) {
Ok(file) => deserialize_messages(file).unwrap_or_else(|e| {
eprintln!(
"Failed to read messages from session file. Starting fresh.\n{}",
e
);
Vec::<Message>::new()
}),
Err(e) => {
eprintln!("Failed to load session file. Starting fresh.\n{}", e);
Vec::<Message>::new()
}
};
prompt.load_user_message_history(messages.clone());
Session {
agent,
prompt,
session_file,
messages,
}
}
pub async fn start(&mut self) -> Result<(), Box<dyn std::error::Error>> {
self.prompt.goose_ready();
loop {
let input = self.prompt.get_input().unwrap();
match input.input_type {
InputType::Message => {
if let Some(content) = &input.content {
self.messages.push(Message::user().with_text(content));
persist_messages(&self.session_file, &self.messages)?;
}
}
InputType::Exit => break,
InputType::AskAgain => continue,
}
self.prompt.show_busy();
self.agent_process_messages().await;
self.prompt.hide_busy();
}
self.close_session().await;
Ok(())
}
pub async fn headless_start(
&mut self,
initial_message: String,
) -> Result<(), Box<dyn std::error::Error>> {
self.messages
.push(Message::user().with_text(initial_message.as_str()));
persist_messages(&self.session_file, &self.messages)?;
self.agent_process_messages().await;
self.close_session().await;
Ok(())
}
async fn agent_process_messages(&mut self) {
let mut stream = match self.agent.reply(&self.messages).await {
Ok(stream) => stream,
Err(e) => {
eprintln!("Error starting reply stream: {}", e);
return;
}
};
loop {
tokio::select! {
response = stream.next() => {
match response {
Some(Ok(message)) => {
self.messages.push(message.clone());
persist_messages(&self.session_file, &self.messages).unwrap_or_else(|e| eprintln!("Failed to persist messages: {}", e));
self.prompt.hide_busy();
self.prompt.render(Box::new(message.clone()));
self.prompt.show_busy();
}
Some(Err(e)) => {
eprintln!("Error: {}", e);
drop(stream);
self.rewind_messages();
self.prompt.render(raw_message(r#"
The error above was an exception we were not able to handle.\n\n
These errors are often related to connection or authentication\n
We've removed the conversation up to the most recent user message
- depending on the error you may be able to continue"#));
break;
}
None => break,
}
}
_ = tokio::signal::ctrl_c() => {
// Kill any running processes when the client disconnects
// TODO is this used? I suspect post MCP this is on the server instead
// goose::process_store::kill_processes();
drop(stream);
self.handle_interrupted_messages();
break;
}
}
}
}
/// Rewind the messages to before the last user message (they have cancelled it).
fn rewind_messages(&mut self) {
if self.messages.is_empty() {
return;
}
// Remove messages until we find the last user 'Text' message (not a tool response).
while let Some(message) = self.messages.last() {
if message.role == Role::User
&& message
.content
.iter()
.any(|c| matches!(c, MessageContent::Text(_)))
{
break;
}
self.messages.pop();
}
// Remove the last user text message we found.
if !self.messages.is_empty() {
self.messages.pop();
}
}
fn handle_interrupted_messages(&mut self) {
// First, get any tool requests from the last message if it exists
let tool_requests = self
.messages
.last()
.filter(|msg| msg.role == Role::Assistant)
.map_or(Vec::new(), |msg| {
msg.content
.iter()
.filter_map(|content| {
if let MessageContent::ToolRequest(req) = content {
Some((req.id.clone(), req.tool_call.clone()))
} else {
None
}
})
.collect()
});
if !tool_requests.is_empty() {
// Interrupted during a tool request
// Create tool responses for all interrupted tool requests
let mut response_message = Message::user();
let last_tool_name = tool_requests
.last()
.and_then(|(_, tool_call)| tool_call.as_ref().ok().map(|tool| tool.name.clone()))
.unwrap_or_else(|| "tool".to_string());
for (req_id, _) in &tool_requests {
response_message.content.push(MessageContent::tool_response(
req_id.clone(),
Err(ToolError::ExecutionError(
"Interrupted by the user to make a correction".to_string(),
)),
));
}
self.messages.push(response_message);
let prompt_response = &format!(
"We interrupted the existing call to {}. How would you like to proceed?",
last_tool_name
);
self.messages
.push(Message::assistant().with_text(prompt_response));
self.prompt.render(raw_message(prompt_response));
} else {
// An interruption occurred outside of a tool request-response.
if let Some(last_msg) = self.messages.last() {
if last_msg.role == Role::User {
match last_msg.content.first() {
Some(MessageContent::ToolResponse(_)) => {
// Interruption occurred after a tool had completed but not assistant reply
let prompt_response = "We interrupted the existing calls to tools. How would you like to proceed?";
self.messages
.push(Message::assistant().with_text(prompt_response));
self.prompt.render(raw_message(prompt_response));
}
Some(_) => {
// A real users message
self.messages.pop();
let prompt_response = "We interrupted before the model replied and removed the last message.";
self.prompt.render(raw_message(prompt_response));
}
None => panic!("No content in last message"),
}
}
}
}
}
async fn close_session(&mut self) {
self.prompt.render(raw_message(
format!(
"Closing session. Recorded to {}\n",
self.session_file.display()
)
.as_str(),
));
self.prompt.close();
let usage = self.agent.usage().await;
log_usage(self.session_file.to_string_lossy().to_string(), usage);
}
pub fn session_file(&self) -> PathBuf {
self.session_file.clone()
}
}
fn raw_message(content: &str) -> Box<Message> {
Box::new(Message::assistant().with_text(content))
}

View File

@@ -0,0 +1,70 @@
/// Helper function to set up a temporary home directory for testing, returns path of that temp dir.
/// Also creates a default profiles.json to avoid obscure test failures when there are no profiles.
#[cfg(test)]
pub fn run_with_tmp_dir<F: FnOnce() -> T, T>(func: F) -> T {
use std::ffi::OsStr;
use tempfile::tempdir;
let temp_dir = tempdir().unwrap();
let temp_dir_path = temp_dir.path().to_path_buf();
setup_profile(&temp_dir_path, None);
temp_env::with_vars(
[
("HOME", Some(temp_dir_path.as_os_str())),
("DATABRICKS_HOST", Some(OsStr::new("tmp_host_url"))),
],
func,
)
}
#[cfg(test)]
#[allow(dead_code)]
pub async fn run_with_tmp_dir_async<F, Fut, T>(func: F) -> T
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = T>,
{
use std::ffi::OsStr;
use tempfile::tempdir;
let temp_dir = tempdir().unwrap();
let temp_dir_path = temp_dir.path().to_path_buf();
setup_profile(&temp_dir_path, None);
temp_env::async_with_vars(
[
("HOME", Some(temp_dir_path.as_os_str())),
("DATABRICKS_HOST", Some(OsStr::new("tmp_host_url"))),
],
func(),
)
.await
}
#[cfg(test)]
use std::path::Path;
#[cfg(test)]
/// Setup a goose profile for testing, and an optional profile string
fn setup_profile(temp_dir_path: &Path, profile_string: Option<&str>) {
use std::fs;
let profile_path = temp_dir_path
.join(".config")
.join("goose")
.join("profiles.json");
fs::create_dir_all(profile_path.parent().unwrap()).unwrap();
let default_profile = r#"
{
"profile_items": {
"default": {
"provider": "databricks",
"model": "claude-3-5-sonnet-2",
"additional_extensions": []
}
}
}"#;
fs::write(&profile_path, profile_string.unwrap_or(default_profile)).unwrap();
}

View File

@@ -0,0 +1,42 @@
[package]
name = "goose-mcp"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
description.workspace = true
[dependencies]
mcp-core = { path = "../mcp-core" }
mcp-server = { path = "../mcp-server" }
anyhow = "1.0.94"
tokio = { version = "1", features = ["full"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
tracing-appender = "0.2"
url = "2.5"
urlencoding = "2.1.3"
base64 = "0.21"
thiserror = "1.0"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
lazy_static = "1.5"
kill_tree = "0.2.4"
shellexpand = "3.1.0"
indoc = "2.0.5"
xcap = "0.0.14"
reqwest = { version = "0.11", features = ["json"] }
async-trait = "0.1"
chrono = { version = "0.4.38", features = ["serde"] }
dirs = "5.0.1"
tempfile = "3.8"
include_dir = "0.7.4"
google-drive3 = "6.0.0"
webbrowser = "0.8"
http-body-util = "0.1.2"
regex = "1.11.1"
[dev-dependencies]
serial_test = "3.0.0"
sysinfo = "0.32.1"

View File

@@ -0,0 +1,9 @@
### Test with MCP Inspector
Update examples/mcp.rs to use the appropriate the MCP server (eg. DeveloperRouter)
```bash
npx @modelcontextprotocol/inspector cargo run -p goose-mcp --example mcp
```
Then visit the Inspector in the browser window and test the different endpoints.

View File

@@ -0,0 +1,36 @@
// An example script to run an MCP server
use anyhow::Result;
use goose_mcp::MemoryRouter;
use mcp_server::router::RouterService;
use mcp_server::{ByteTransport, Server};
use tokio::io::{stdin, stdout};
use tracing_appender::rolling::{RollingFileAppender, Rotation};
use tracing_subscriber::{self, EnvFilter};
#[tokio::main]
async fn main() -> Result<()> {
// Set up file appender for logging
let file_appender = RollingFileAppender::new(Rotation::DAILY, "logs", "goose-mcp-example.log");
// Initialize the tracing subscriber with file and stdout logging
tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env().add_directive(tracing::Level::INFO.into()))
.with_writer(file_appender)
.with_target(false)
.with_thread_ids(true)
.with_file(true)
.with_line_number(true)
.init();
tracing::info!("Starting MCP server");
// Create an instance of our counter router
let router = RouterService(MemoryRouter::new());
// Create and run the server
let server = Server::new(router);
let transport = ByteTransport::new(stdin(), stdout());
tracing::info!("Server initialized and ready to handle requests");
Ok(server.run(transport).await?)
}

View File

@@ -0,0 +1,34 @@
use std::path::Path;
/// Get the markdown language identifier for a file extension
pub fn get_language_identifier(path: &Path) -> &'static str {
match path.extension().and_then(|ext| ext.to_str()) {
Some("rs") => "rust",
Some("py") => "python",
Some("js") => "javascript",
Some("ts") => "typescript",
Some("json") => "json",
Some("toml") => "toml",
Some("yaml") | Some("yml") => "yaml",
Some("sh") => "bash",
Some("go") => "go",
Some("md") => "markdown",
Some("html") => "html",
Some("css") => "css",
Some("sql") => "sql",
Some("java") => "java",
Some("cpp") | Some("cc") | Some("cxx") => "cpp",
Some("c") => "c",
Some("h") | Some("hpp") => "cpp",
Some("rb") => "ruby",
Some("php") => "php",
Some("swift") => "swift",
Some("kt") | Some("kts") => "kotlin",
Some("scala") => "scala",
Some("r") => "r",
Some("m") => "matlab",
Some("pl") => "perl",
Some("dockerfile") => "dockerfile",
_ => "",
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,585 @@
use indoc::indoc;
use regex::Regex;
use serde_json::{json, Value};
use std::{env, fs, future::Future, io::Write, path::Path, pin::Pin};
use mcp_core::{
handler::{ResourceError, ToolError},
protocol::ServerCapabilities,
resource::Resource,
tool::Tool,
};
use mcp_server::router::CapabilitiesBuilder;
use mcp_server::Router;
use mcp_core::content::Content;
use google_drive3::{
self,
api::{File, Scope},
hyper_rustls::{self, HttpsConnector},
hyper_util::{self, client::legacy::connect::HttpConnector},
yup_oauth2::{
self,
authenticator_delegate::{DefaultInstalledFlowDelegate, InstalledFlowDelegate},
InstalledFlowAuthenticator,
},
DriveHub,
};
use http_body_util::BodyExt;
/// async function to be pinned by the `present_user_url` method of the trait
/// we use the existing `DefaultInstalledFlowDelegate::present_user_url` method as a fallback for
/// when the browser did not open for example, the user still see's the URL.
async fn browser_user_url(url: &str, need_code: bool) -> Result<String, String> {
tracing::info!(oauth_url = url, "Attempting OAuth login flow");
if let Err(e) = webbrowser::open(url) {
tracing::debug!(oauth_url = url, error = ?e, "Failed to open OAuth flow");
println!("Please open this URL in your browser:\n{}", url);
}
let def_delegate = DefaultInstalledFlowDelegate;
def_delegate.present_user_url(url, need_code).await
}
/// our custom delegate struct we will implement a flow delegate trait for:
/// in this case we will implement the `InstalledFlowDelegated` trait
#[derive(Copy, Clone)]
struct LocalhostBrowserDelegate;
/// here we implement only the present_user_url method with the added webbrowser opening
/// the other behaviour of the trait does not need to be changed.
impl InstalledFlowDelegate for LocalhostBrowserDelegate {
/// the actual presenting of URL and browser opening happens in the function defined above here
/// we only pin it
fn present_user_url<'a>(
&'a self,
url: &'a str,
need_code: bool,
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send + 'a>> {
Box::pin(browser_user_url(url, need_code))
}
}
pub struct GoogleDriveRouter {
tools: Vec<Tool>,
instructions: String,
drive: DriveHub<HttpsConnector<HttpConnector>>,
}
impl GoogleDriveRouter {
async fn google_auth() -> DriveHub<HttpsConnector<HttpConnector>> {
let oauth_config = env::var("GOOGLE_DRIVE_OAUTH_CONFIG");
let keyfile_path_str = env::var("GOOGLE_DRIVE_OAUTH_PATH")
.unwrap_or_else(|_| "./gcp-oauth.keys.json".to_string());
let credentials_path_str = env::var("GOOGLE_DRIVE_CREDENTIALS_PATH")
.unwrap_or_else(|_| "./gdrive-server-credentials.json".to_string());
let expanded_keyfile = shellexpand::tilde(keyfile_path_str.as_str());
let keyfile_path = Path::new(expanded_keyfile.as_ref());
let expanded_credentials = shellexpand::tilde(credentials_path_str.as_str());
let credentials_path = Path::new(expanded_credentials.as_ref());
tracing::info!(
credentials_path = credentials_path_str,
keyfile_path = keyfile_path_str,
"Google Drive MCP server authentication config paths"
);
if !keyfile_path.exists() && oauth_config.is_ok() {
tracing::debug!(
oauth_config = ?oauth_config,
"Google Drive MCP server OAuth config"
);
// attempt to create the path
if let Some(parent_dir) = keyfile_path.parent() {
let _ = fs::create_dir_all(parent_dir);
}
if let Ok(mut file) = fs::File::create(keyfile_path) {
let _ = file.write_all(oauth_config.unwrap().as_bytes());
}
}
let secret = yup_oauth2::read_application_secret(keyfile_path)
.await
.expect("expected keyfile for google auth");
let auth = InstalledFlowAuthenticator::builder(
secret,
yup_oauth2::InstalledFlowReturnMethod::HTTPRedirect,
)
.persist_tokens_to_disk(credentials_path)
.flow_delegate(Box::new(LocalhostBrowserDelegate))
.build()
.await
.expect("expected successful authentication");
let client =
hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new())
.build(
hyper_rustls::HttpsConnectorBuilder::new()
.with_native_roots()
.unwrap()
.https_or_http()
.enable_http1()
.build(),
);
DriveHub::new(client, auth)
}
pub async fn new() -> Self {
let drive = Self::google_auth().await;
// handle auth
let search_tool = Tool::new(
"search".to_string(),
indoc! {r#"
Search for files in google drive by name, given an input search query.
"#}
.to_string(),
json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query",
},
},
"required": ["query"],
}),
);
let read_tool = Tool::new(
"read".to_string(),
indoc! {r#"
Read a file from google drive using the file uri.
Optionally include base64 encoded images, false by default.
"#}
.to_string(),
json!({
"type": "object",
"properties": {
"uri": {
"type": "string",
"description": "google drive uri of the file to read",
},
"includeImages": {
"type": "boolean",
"description": "Whether or not to include images as base64 encoded strings, defaults to false",
}
},
"required": ["uri"],
}),
);
let instructions = indoc::formatdoc! {r#"
Google Drive MCP Server Instructions
## Overview
The Google Drive MCP server provides two main tools for interacting with Google Drive files:
1. search - Find files in your Google Drive
2. read - Read file contents directly using a uri in the `gdrive:///uri` format
## Available Tools
### 1. Search Tool
Search for files in Google Drive, by name and ordered by most recently viewedByMeTime.
Returns: List of files with their names, MIME types, and IDs
### 2. Read File Tool
Read a file's contents using its ID, and optionally include images as base64 encoded data.
The default is to exclude images, to include images set includeImages to true in the query.
Images take up a large amount of context, this should only be used if a
user explicity needs the image data.
Limitations: Google Sheets exporting only supports reading the first sheet. This is an important limitation that should
be communicated to the user whenever dealing with a Google Sheet (mimeType: application/vnd.google-apps.spreadsheet).
## File Format Handling
The server automatically handles different file types:
- Google Docs → Markdown
- Google Sheets → CSV
- Google Presentations → Plain text
- Text/JSON files → UTF-8 text
- Binary files → Base64 encoded
## Common Usage Pattern
1. First, search for the file you want to read, searching by name.
2. Then, use the file URI from the search results to read its contents.
## Best Practices
1. Always use search first to find the correct file URI
2. Search results include file types (MIME types) to help identify the right file
3. Search is limited to 10 results per query, so use specific search terms
4. The server has read-only access to Google Drive
## Error Handling
If you encounter errors:
1. Verify the file URI is correct
2. Ensure you have access to the file
3. Check if the file format is supported
4. Verify the server is properly configured
Remember: Always use the tools in sequence - search first to get the file URI, then read to access the contents.
"#};
Self {
tools: vec![search_tool, read_tool],
instructions,
drive,
}
}
// Implement search tool functionality
async fn search(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let query = params
.get("query")
.and_then(|q| q.as_str())
.ok_or(ToolError::InvalidParameters(
"The query string is required".to_string(),
))?
.replace('\\', "\\\\")
.replace('\'', "\\'");
let result = self
.drive
.files()
.list()
.q(format!("name contains '{}'", query).as_str())
.order_by("viewedByMeTime desc")
.param("fields", "files(id, name, mimeType, modifiedTime, size)")
.page_size(10)
.supports_all_drives(true)
.include_items_from_all_drives(true)
.clear_scopes() // Scope::MeetReadonly is the default, remove it
.add_scope(Scope::Readonly)
.doit()
.await;
match result {
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to execute google drive search query, {}.",
e
))),
Ok(r) => {
let content =
r.1.files
.map(|files| {
files.into_iter().map(|f| {
format!(
"{} ({}) (uri: {})",
f.name.unwrap_or_default(),
f.mime_type.unwrap_or_default(),
f.id.unwrap_or_default()
)
})
})
.into_iter()
.flatten()
.collect::<Vec<_>>()
.join("\n");
Ok(vec![Content::text(content.to_string())])
}
}
}
async fn fetch_file_metadata(&self, uri: &str) -> Result<File, ToolError> {
self.drive
.files()
.get(uri)
.param("fields", "mimeType")
.supports_all_drives(true)
.clear_scopes()
.add_scope(Scope::Readonly)
.doit()
.await
.map_err(|e| {
ToolError::ExecutionError(format!(
"Failed to execute Google Drive get query, {}.",
e
))
})
.map(|r| r.1)
}
fn strip_image_body(&self, input: &str) -> String {
let image_regex = Regex::new(r"<data:image/[a-zA-Z0-9.-]+;base64,[^>]+>").unwrap();
image_regex.replace_all(input, "").to_string()
}
// Downloading content with alt=media only works if the file is stored in Drive.
// To download Google Docs, Sheets, and Slides use files.export instead.
async fn export_google_file(
&self,
uri: &str,
mime_type: &str,
include_images: bool,
) -> Result<Vec<Content>, ToolError> {
let export_mime_type = match mime_type {
"application/vnd.google-apps.document" => "text/markdown",
"application/vnd.google-apps.spreadsheet" => "text/csv",
"application/vnd.google-apps.presentation" => "text/plain",
_ => "text/plain",
};
let result = self
.drive
.files()
.export(uri, export_mime_type)
.param("alt", "media")
.clear_scopes()
.add_scope(Scope::Readonly)
.doit()
.await;
match result {
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to execute google drive export for {}, {}.",
uri, e
))),
Ok(r) => {
if let Ok(body) = r.into_body().collect().await {
if let Ok(response) = String::from_utf8(body.to_bytes().to_vec()) {
let content = if !include_images {
self.strip_image_body(&response)
} else {
response
};
Ok(vec![Content::text(content).with_priority(0.1)])
} else {
Err(ToolError::ExecutionError(format!(
"Failed to export google drive to string, {}.",
uri,
)))
}
} else {
Err(ToolError::ExecutionError(format!(
"Failed to export google drive document, {}.",
uri,
)))
}
}
}
}
// handle for files we can use files.get on
async fn get_google_file(
&self,
uri: &str,
include_images: bool,
) -> Result<Vec<Content>, ToolError> {
let result = self
.drive
.files()
.get(uri)
.param("alt", "media")
.clear_scopes()
.add_scope(Scope::Readonly)
.doit()
.await;
match result {
Err(e) => Err(ToolError::ExecutionError(format!(
"Failed to execute google drive export for {}, {}.",
uri, e
))),
Ok(r) => {
let file = r.1;
let mime_type = file
.mime_type
.unwrap_or("application/octet-stream".to_string());
if mime_type.starts_with("text/") || mime_type == "application/json" {
if let Ok(body) = r.0.into_body().collect().await {
if let Ok(response) = String::from_utf8(body.to_bytes().to_vec()) {
let content = if !include_images {
self.strip_image_body(&response)
} else {
response
};
Ok(vec![Content::text(content).with_priority(0.1)])
} else {
Err(ToolError::ExecutionError(format!(
"Failed to convert google drive to string, {}.",
uri,
)))
}
} else {
Err(ToolError::ExecutionError(format!(
"Failed to get google drive document, {}.",
uri,
)))
}
} else {
//TODO: handle base64 image case, see typscript mcp-gdrive
Err(ToolError::ExecutionError(format!(
"Suported mimeType {}, for {}",
mime_type, uri,
)))
}
}
}
}
async fn read(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let uri =
params
.get("uri")
.and_then(|q| q.as_str())
.ok_or(ToolError::InvalidParameters(
"The uri of the file is required".to_string(),
))?;
let drive_uri = uri.replace("gdrive:///", "");
let include_images = params
.get("includeImages")
.and_then(|i| i.as_bool())
.unwrap_or(false);
let metadata = self.fetch_file_metadata(&drive_uri).await?;
let mime_type = metadata.mime_type.ok_or_else(|| {
ToolError::ExecutionError(format!("Missing mime type in file metadata for {}.", uri))
})?;
// Handle Google Docs export
if mime_type.starts_with("application/vnd.google-apps") {
self.export_google_file(&drive_uri, &mime_type, include_images)
.await
} else {
self.get_google_file(&drive_uri, include_images).await
}
}
async fn read_google_resource(&self, uri: String) -> Result<String, ResourceError> {
self.read(json!({"uri": uri}))
.await
.map_err(|e| ResourceError::ExecutionError(e.to_string()))
.map(|contents| {
contents
.into_iter()
.map(|content| content.as_text().unwrap_or_default().to_string())
.collect::<Vec<_>>()
.join("\n")
})
}
async fn list_google_resources(&self, params: Value) -> Vec<Resource> {
let next_page_token = params.get("cursor").and_then(|q| q.as_str());
let mut query = self
.drive
.files()
.list()
.order_by("viewedByMeTime desc")
.page_size(10)
.param("fields", "nextPageToken, files(id, name, mimeType)")
.supports_all_drives(true)
.include_items_from_all_drives(true)
.clear_scopes() // Scope::MeetReadonly is the default, remove it
.add_scope(Scope::Readonly);
// add a next token if we have one
if let Some(token) = next_page_token {
query = query.page_token(token)
}
let result = query.doit().await;
match result {
Err(_) => {
//Err(ResourceError::ExecutionError(format!(
// "Failed to execute google drive list query, {}.",
// e,
//)));
vec![]
}
Ok(r) => {
r.1.files
.map(|files| {
files.into_iter().map(|f| Resource {
uri: f.id.unwrap_or_default(),
mime_type: f.mime_type.unwrap_or_default(),
name: f.name.unwrap_or_default(),
description: None,
annotations: None,
})
})
.into_iter()
.flatten()
.collect::<Vec<_>>()
}
}
}
}
impl Router for GoogleDriveRouter {
fn name(&self) -> String {
"google_drive".to_string()
}
fn instructions(&self) -> String {
self.instructions.clone()
}
fn capabilities(&self) -> ServerCapabilities {
CapabilitiesBuilder::new()
.with_tools(false)
.with_resources(false, false)
.build()
}
fn list_tools(&self) -> Vec<Tool> {
self.tools.clone()
}
fn call_tool(
&self,
tool_name: &str,
arguments: Value,
) -> Pin<Box<dyn Future<Output = Result<Vec<Content>, ToolError>> + Send + 'static>> {
let this = self.clone();
let tool_name = tool_name.to_string();
Box::pin(async move {
match tool_name.as_str() {
"search" => this.search(arguments).await,
"read" => this.read(arguments).await,
_ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))),
}
})
}
fn list_resources(&self) -> Vec<Resource> {
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current()
.block_on(async { self.list_google_resources(json!({})).await })
})
}
fn read_resource(
&self,
uri: &str,
) -> Pin<Box<dyn Future<Output = Result<String, ResourceError>> + Send + 'static>> {
let this = self.clone();
let uri_clone = uri.to_string();
Box::pin(async move { this.read_google_resource(uri_clone).await })
}
}
impl Clone for GoogleDriveRouter {
fn clone(&self) -> Self {
Self {
tools: self.tools.clone(),
instructions: self.instructions.clone(),
drive: self.drive.clone(),
}
}
}

View File

@@ -0,0 +1,217 @@
mod proxy;
use anyhow::Result;
use mcp_core::{
content::Content,
handler::{ResourceError, ToolError},
protocol::ServerCapabilities,
resource::Resource,
role::Role,
tool::Tool,
};
use mcp_server::router::CapabilitiesBuilder;
use mcp_server::Router;
use serde_json::Value;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::{sleep, Duration};
use tracing::error;
use self::proxy::JetBrainsProxy;
pub struct JetBrainsRouter {
tools: Arc<Mutex<Vec<Tool>>>,
proxy: Arc<JetBrainsProxy>,
instructions: String,
}
impl Default for JetBrainsRouter {
fn default() -> Self {
Self::new()
}
}
impl JetBrainsRouter {
pub fn new() -> Self {
let tools = Arc::new(Mutex::new(Vec::new()));
let proxy = Arc::new(JetBrainsProxy::new());
let instructions = "JetBrains IDE integration".to_string();
// Initialize the proxy
let proxy_clone = Arc::clone(&proxy);
tokio::spawn(async move {
if let Err(e) = proxy_clone.start().await {
error!("Failed to start JetBrains proxy: {}", e);
}
});
// Start the background task to update tools
let tools_clone = Arc::clone(&tools);
let proxy_clone = Arc::clone(&proxy);
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(5));
loop {
interval.tick().await;
match proxy_clone.list_tools().await {
Ok(new_tools) => {
let mut tools = tools_clone.lock().await;
*tools = new_tools;
}
Err(e) => {
error!("Failed to update tools: {}", e);
}
}
}
});
Self {
tools,
proxy,
instructions,
}
}
async fn call_proxy_tool(
&self,
tool_name: String,
arguments: Value,
) -> Result<Vec<Content>, ToolError> {
let result = self
.proxy
.call_tool(&tool_name, arguments)
.await
.map_err(|e| ToolError::ExecutionError(e.to_string()))?;
// Create a success message for the assistant
let mut contents = vec![
Content::text(format!("Tool {} executed successfully", tool_name))
.with_audience(vec![Role::Assistant]),
];
// Add the tool's result contents
contents.extend(result.content);
Ok(contents)
}
async fn ensure_tools(&self) -> Result<(), ToolError> {
let mut retry_count = 0;
let max_retries = 50; // 5 second total wait time
let retry_delay = Duration::from_millis(100);
while retry_count < max_retries {
let tools = self.tools.lock().await;
if !tools.is_empty() {
return Ok(());
}
drop(tools); // Release the lock before sleeping
sleep(retry_delay).await;
retry_count += 1;
}
Err(ToolError::ExecutionError("Failed to get tools list from IDE. Make sure the IDE is running and the plugin is installed.".to_string()))
}
}
impl Router for JetBrainsRouter {
fn name(&self) -> String {
"jetbrains".to_string()
}
fn instructions(&self) -> String {
self.instructions.clone()
}
fn capabilities(&self) -> ServerCapabilities {
CapabilitiesBuilder::new().with_tools(true).build()
}
fn list_tools(&self) -> Vec<Tool> {
// Use block_in_place to avoid blocking the runtime
tokio::task::block_in_place(|| {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let tools = self.tools.lock().await;
if tools.is_empty() {
drop(tools);
if let Err(e) = self.ensure_tools().await {
error!("Failed to ensure tools: {}", e);
vec![]
} else {
self.tools.lock().await.clone()
}
} else {
tools.clone()
}
})
})
}
fn call_tool(
&self,
tool_name: &str,
arguments: Value,
) -> Pin<Box<dyn Future<Output = Result<Vec<Content>, ToolError>> + Send + 'static>> {
let this = self.clone();
let tool_name = tool_name.to_string();
Box::pin(async move {
this.ensure_tools().await?;
this.call_proxy_tool(tool_name, arguments).await
})
}
fn list_resources(&self) -> Vec<Resource> {
vec![]
}
fn read_resource(
&self,
_uri: &str,
) -> Pin<Box<dyn Future<Output = Result<String, ResourceError>> + Send + 'static>> {
Box::pin(async { Err(ResourceError::NotFound("Resource not found".into())) })
}
}
impl Clone for JetBrainsRouter {
fn clone(&self) -> Self {
Self {
tools: Arc::clone(&self.tools),
proxy: Arc::clone(&self.proxy),
instructions: self.instructions.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::OnceCell;
static JETBRAINS_ROUTER: OnceCell<JetBrainsRouter> = OnceCell::const_new();
async fn get_router() -> &'static JetBrainsRouter {
JETBRAINS_ROUTER
.get_or_init(|| async { JetBrainsRouter::new() })
.await
}
#[tokio::test]
async fn test_router_creation() {
let router = get_router().await;
assert_eq!(router.name(), "jetbrains");
assert!(!router.instructions().is_empty());
}
#[tokio::test]
async fn test_capabilities() {
let router = get_router().await;
let capabilities = router.capabilities();
assert!(capabilities.tools.is_some());
}
}

View File

@@ -0,0 +1,341 @@
use anyhow::{anyhow, Result};
use mcp_core::{Content, Tool};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::env;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, error, info};
const PORT_RANGE_START: u16 = 63342;
const PORT_RANGE_END: u16 = 63352;
const ENDPOINT_CHECK_INTERVAL: Duration = Duration::from_secs(10);
#[derive(Debug, Serialize, Deserialize)]
struct IDEResponseOk {
status: String,
error: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
struct IDEResponseErr {
status: Option<String>,
error: String,
}
#[derive(Debug, Serialize)]
pub struct CallToolResult {
pub content: Vec<Content>,
pub is_error: bool,
}
#[derive(Debug)]
pub struct JetBrainsProxy {
cached_endpoint: Arc<RwLock<Option<String>>>,
previous_response: Arc<RwLock<Option<String>>>,
client: Client,
}
impl JetBrainsProxy {
pub fn new() -> Self {
Self {
cached_endpoint: Arc::new(RwLock::new(None)),
previous_response: Arc::new(RwLock::new(None)),
client: Client::new(),
}
}
async fn test_list_tools(&self, endpoint: &str) -> Result<bool> {
debug!("Sending test request to {}/mcp/list_tools", endpoint);
let response = match self
.client
.get(format!("{}/mcp/list_tools", endpoint))
.send()
.await
{
Ok(resp) => {
debug!("Got response with status: {}", resp.status());
resp
}
Err(e) => {
debug!("Error testing endpoint {}: {}", endpoint, e);
return Ok(false);
}
};
if !response.status().is_success() {
debug!("Test request failed with status {}", response.status());
return Ok(false);
}
let current_response = response.text().await?;
debug!("Received response: {}", current_response);
// Try to parse as JSON array to validate format
if serde_json::from_str::<Vec<Value>>(&current_response).is_err() {
debug!("Response is not a valid JSON array of tools");
return Ok(false);
}
let mut prev_response = self.previous_response.write().await;
if let Some(prev) = prev_response.as_ref() {
if prev != &current_response {
debug!("Response changed since last check");
self.send_tools_changed().await;
}
}
*prev_response = Some(current_response);
Ok(true)
}
async fn find_working_ide_endpoint(&self) -> Result<String> {
debug!("Attempting to find working IDE endpoint...");
// Check IDE_PORT environment variable first
if let Ok(port) = env::var("IDE_PORT") {
debug!("Found IDE_PORT environment variable: {}", port);
let test_endpoint = format!("http://127.0.0.1:{}/api", port);
if self.test_list_tools(&test_endpoint).await? {
debug!("IDE_PORT {} is working", port);
return Ok(test_endpoint);
}
debug!("IDE_PORT {} is not responding correctly", port);
return Err(anyhow!(
"Specified IDE_PORT={} is not responding correctly",
port
));
}
debug!(
"No IDE_PORT environment variable, scanning port range {}-{}",
PORT_RANGE_START, PORT_RANGE_END
);
// Scan port range
for port in PORT_RANGE_START..=PORT_RANGE_END {
let candidate_endpoint = format!("http://127.0.0.1:{}/api", port);
debug!("Testing port {}...", port);
if self.test_list_tools(&candidate_endpoint).await? {
debug!("Found working IDE endpoint at {}", candidate_endpoint);
return Ok(candidate_endpoint);
}
}
debug!("No working IDE endpoint found in port range");
Err(anyhow!(
"No working IDE endpoint found in range {}-{}",
PORT_RANGE_START,
PORT_RANGE_END
))
}
async fn update_ide_endpoint(&self) {
debug!("Updating IDE endpoint...");
match self.find_working_ide_endpoint().await {
Ok(endpoint) => {
let mut cached = self.cached_endpoint.write().await;
*cached = Some(endpoint.clone());
debug!("Updated cached endpoint to: {}", endpoint);
}
Err(e) => {
debug!("Failed to update IDE endpoint: {}", e);
error!("Failed to update IDE endpoint: {}", e);
}
}
}
pub async fn list_tools(&self) -> Result<Vec<Tool>> {
debug!("Listing tools...");
let endpoint = {
let cached = self.cached_endpoint.read().await;
match cached.as_ref() {
Some(ep) => {
debug!("Using cached endpoint: {}", ep);
ep.clone()
}
None => {
debug!("No cached endpoint available");
return Ok(vec![]);
}
}
};
debug!("Sending list_tools request to {}/mcp/list_tools", endpoint);
let response = match self
.client
.get(format!("{}/mcp/list_tools", endpoint))
.send()
.await
{
Ok(resp) => {
debug!("Got response with status: {}", resp.status());
resp
}
Err(e) => {
debug!("Failed to send request: {}", e);
return Err(anyhow!("Failed to send request: {}", e));
}
};
if !response.status().is_success() {
debug!("Request failed with status: {}", response.status());
return Err(anyhow!(
"Failed to fetch tools with status {}",
response.status()
));
}
let response_text = response.text().await?;
debug!("Got response text: {}", response_text);
let tools_response: Value = serde_json::from_str(&response_text).map_err(|e| {
debug!("Failed to parse response as JSON: {}", e);
anyhow!("Failed to parse response as JSON: {}", e)
})?;
debug!("Parsed JSON response: {:?}", tools_response);
let tools: Vec<Tool> = tools_response
.as_array()
.ok_or_else(|| {
debug!("Response is not a JSON array");
anyhow!("Invalid tools response format: not an array")
})?
.iter()
.filter_map(|t| {
if let (Some(name), Some(description)) =
(t["name"].as_str(), t["description"].as_str())
{
// Get just the first sentence of the description
let first_sentence = description
.split('.')
.next()
.unwrap_or(description)
.trim()
.to_string()
+ ".";
// Handle input_schema as either a string or an object
let input_schema = match &t["inputSchema"] {
Value::String(s) => Value::String(s.clone()),
Value::Object(o) => Value::Object(o.clone()),
_ => {
debug!(
"Invalid inputSchema format for tool {}: {:?}",
name, t["inputSchema"]
);
return None;
}
};
Some(Tool {
name: name.to_string(),
description: first_sentence,
input_schema,
})
} else {
debug!("Skipping invalid tool entry: {:?}", t);
None
}
})
.collect();
debug!("Collected {} tools", tools.len());
Ok(tools)
}
pub async fn call_tool(&self, name: &str, args: Value) -> Result<CallToolResult> {
let endpoint = self
.cached_endpoint
.read()
.await
.clone()
.ok_or_else(|| anyhow!("No working IDE endpoint available"))?;
debug!(
"ENDPOINT: {} | Tool name: {} | args: {}",
endpoint, name, args
);
let response = self
.client
.post(format!("{}/mcp/{}", endpoint, name))
.json(&args)
.send()
.await?;
if !response.status().is_success() {
debug!("Response failed with status: {}", response.status());
return Err(anyhow!("Response failed: {}", response.status()));
}
let ide_response: Value = response.json().await?;
let (is_error, text) = match ide_response {
Value::Object(map) => {
let status = map.get("status").and_then(|v| v.as_str());
let error = map.get("error").and_then(|v| v.as_str());
match (status, error) {
(Some(s), None) => (false, s.to_string()),
(None, Some(e)) => (true, e.to_string()),
_ => {
debug!("Invalid response format from IDE");
return Err(anyhow!("Invalid response format from IDE"));
}
}
}
_ => {
debug!("Unexpected response type from IDE");
return Err(anyhow!("Unexpected response type from IDE"));
}
};
Ok(CallToolResult {
content: vec![Content::text(text)],
is_error,
})
}
async fn send_tools_changed(&self) {
debug!("Sending tools changed notification");
// TODO: Implement notification mechanism when needed
}
pub async fn start(&self) -> Result<()> {
debug!("Initializing JetBrains Proxy...");
info!("Initializing JetBrains Proxy...");
// Initial endpoint check
debug!("Performing initial endpoint check...");
self.update_ide_endpoint().await;
// Schedule periodic endpoint checks
let proxy = self.clone();
tokio::spawn(async move {
loop {
tokio::time::sleep(ENDPOINT_CHECK_INTERVAL).await;
debug!("Performing periodic endpoint check...");
proxy.update_ide_endpoint().await;
}
});
debug!("JetBrains Proxy running");
info!("JetBrains Proxy running");
Ok(())
}
}
impl Clone for JetBrainsProxy {
fn clone(&self) -> Self {
Self {
cached_endpoint: Arc::clone(&self.cached_endpoint),
previous_response: Arc::clone(&self.previous_response),
client: Client::new(),
}
}
}

View File

@@ -0,0 +1,11 @@
mod developer;
mod google_drive;
mod jetbrains;
mod memory;
mod nondeveloper;
pub use developer::DeveloperRouter;
pub use google_drive::GoogleDriveRouter;
pub use jetbrains::JetBrainsRouter;
pub use memory::MemoryRouter;
pub use nondeveloper::NonDeveloperRouter;

View File

@@ -0,0 +1,535 @@
use async_trait::async_trait;
use indoc::formatdoc;
use serde_json::{json, Value};
use std::{
collections::HashMap,
fs,
future::Future,
io::{self, Read, Write},
path::PathBuf,
pin::Pin,
};
use mcp_core::{
handler::{ResourceError, ToolError},
protocol::ServerCapabilities,
resource::Resource,
tool::{Tool, ToolCall},
Content,
};
use mcp_server::router::CapabilitiesBuilder;
use mcp_server::Router;
// MemoryRouter implementation
#[derive(Clone)]
pub struct MemoryRouter {
tools: Vec<Tool>,
instructions: String,
global_memory_dir: PathBuf,
local_memory_dir: PathBuf,
}
impl Default for MemoryRouter {
fn default() -> Self {
Self::new()
}
}
impl MemoryRouter {
pub fn new() -> Self {
let remember_memory = Tool::new(
"remember_memory",
"Stores a memory with optional tags in a specified category",
json!({
"type": "object",
"properties": {
"category": {"type": "string"},
"data": {"type": "string"},
"tags": {"type": "array", "items": {"type": "string"}},
"is_global": {"type": "boolean"}
},
"required": ["category", "data", "is_global"]
}),
);
let retrieve_memories = Tool::new(
"retrieve_memories",
"Retrieves all memories from a specified category",
json!({
"type": "object",
"properties": {
"category": {"type": "string"},
"is_global": {"type": "boolean"}
},
"required": ["category", "is_global"]
}),
);
let remove_memory_category = Tool::new(
"remove_memory_category",
"Removes all memories within a specified category",
json!({
"type": "object",
"properties": {
"category": {"type": "string"},
"is_global": {"type": "boolean"}
},
"required": ["category", "is_global"]
}),
);
let remove_specific_memory = Tool::new(
"remove_specific_memory",
"Removes a specific memory within a specified category",
json!({
"type": "object",
"properties": {
"category": {"type": "string"},
"memory_content": {"type": "string"},
"is_global": {"type": "boolean"}
},
"required": ["category", "memory_content", "is_global"]
}),
);
let instructions = formatdoc! {r#"
This extension allows storage and retrieval of categorized information with tagging support. It's designed to help
manage important information across sessions in a systematic and organized manner.
Capabilities:
1. Store information in categories with optional tags for context-based retrieval.
2. Search memories by content or specific tags to find relevant information.
3. List all available memory categories for easy navigation.
4. Remove entire categories of memories when they are no longer needed.
Interaction Protocol:
When important information is identified, such as:
- User-specific data (e.g., name, preferences)
- Project-related configurations
- Workflow descriptions
- Other critical settings
The protocol is:
1. Identify the critical piece of information.
2. Ask the user if they'd like to store it for later reference.
3. Upon agreement:
- Suggest a relevant category like "personal" for user data or "development" for project preferences.
- Inquire about any specific tags they want to apply for easier lookup.
- Confirm the desired storage location:
- Local storage (.goose/memory) for project-specific details.
- Global storage (~/.config/goose/memory) for user-wide data.
Example Interaction for Storing Information:
User: "For this project, we use black for code formatting"
Assistant: "You've mentioned a development preference. Would you like to remember this for future conversations?
User: "Yes, please."
Assistant: "I'll store this in the 'development' category. Any specific tags to add? Suggestions: #formatting
#tools"
User: "Yes, use those tags."
Assistant: "Shall I store this locally for this project only, or globally for all projects?"
User: "Locally, please."
Assistant: *Stores the information under category="development", tags="formatting tools", scope="local"*
Retrieving Memories:
To access stored information, utilize the memory retrieval protocols:
- **Search by Category**:
- Provides all memories within the specified context.
- Use: `retrieve_memories(category="development", is_global=False)`
- Note: If you want to retrieve all local memories, use `retrieve_memories(category="*", is_global=False)`
- Note: If you want to retrieve all global memories, use `retrieve_memories(category="*", is_global=True)`
- **Filter by Tags**:
- Enables targeted retrieval based on specific tags.
- Use: Provide tag filters to refine search.
To remove a memory, use the following protocol:
- **Remove by Category**:
- Removes all memories within the specified category.
- Use: `remove_memory_category(category="development", is_global=False)`
- Note: If you want to remove all local memories, use `remove_memory_category(category="*", is_global=False)`
- Note: If you want to remove all global memories, use `remove_memory_category(category="*", is_global=True)`
The Protocol is:
1. Confirm what kind of information the user seeks by category or keyword.
2. Suggest categories or relevant tags based on the user's request.
3. Use the retrieve function to access relevant memory entries.
4. Present a summary of findings, offering detailed exploration upon request.
Example Interaction for Retrieving Information:
User: "What configuration do we use for code formatting?"
Assistant: "Let me check the 'development' category for any related memories. Searching using #formatting tag."
Assistant: *Executes retrieval: `retrieve_memories(category="development", is_global=False)`*
Assistant: "We have 'black' configured for code formatting, specific to this project. Would you like further
details?"
Memory Overview:
- Categories can include a wide range of topics, structured to keep information grouped logically.
- Tags enable quick filtering and identification of specific entries.
Operational Guidelines:
- Always confirm with the user before saving information.
- Propose suitable categories and tag suggestions.
- Discuss storage scope thoroughly to align with user needs.
- Acknowledge the user about what is stored and where, for transparency and ease of future retrieval.
"#};
// Check for .goose/memory in current directory
let local_memory_dir = std::env::var("GOOSE_WORKING_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| std::env::current_dir().unwrap())
.join(".goose")
.join("memory");
// Check for .config/goose/memory in user's home directory
let global_memory_dir = dirs::home_dir()
.map(|home| home.join(".config/goose/memory"))
.unwrap_or_else(|| PathBuf::from(".config/goose/memory"));
fs::create_dir_all(&global_memory_dir).unwrap();
fs::create_dir_all(&local_memory_dir).unwrap();
let mut memory_router = Self {
tools: vec![
remember_memory,
retrieve_memories,
remove_memory_category,
remove_specific_memory,
],
instructions: instructions.clone(),
global_memory_dir,
local_memory_dir,
};
let retrieved_global_memories = memory_router.retrieve_all(true);
let retrieved_local_memories = memory_router.retrieve_all(false);
let mut updated_instructions = instructions;
let memories_follow_up_instructions = formatdoc! {r#"
**Here are the user's currently saved memories:**
Please keep this information in mind when answering future questions.
Do not bring up memories unless relevant.
Note: if the user has not saved any memories, this section will be empty.
Note: if the user removes a memory that was previously loaded into the system, please remove it from the system instructions.
"#};
updated_instructions.push_str("\n\n");
updated_instructions.push_str(&memories_follow_up_instructions);
if let Ok(global_memories) = retrieved_global_memories {
if !global_memories.is_empty() {
updated_instructions.push_str("\n\nGlobal Memories:\n");
for (category, memories) in global_memories {
updated_instructions.push_str(&format!("\nCategory: {}\n", category));
for memory in memories {
updated_instructions.push_str(&format!("- {}\n", memory));
}
}
}
}
if let Ok(local_memories) = retrieved_local_memories {
if !local_memories.is_empty() {
updated_instructions.push_str("\n\nLocal Memories:\n");
for (category, memories) in local_memories {
updated_instructions.push_str(&format!("\nCategory: {}\n", category));
for memory in memories {
updated_instructions.push_str(&format!("- {}\n", memory));
}
}
}
}
memory_router.set_instructions(updated_instructions);
memory_router
}
// Add a setter method for instructions
pub fn set_instructions(&mut self, new_instructions: String) {
self.instructions = new_instructions;
}
pub fn get_instructions(&self) -> &str {
&self.instructions
}
fn get_memory_file(&self, category: &str, is_global: bool) -> PathBuf {
// Defaults to local memory if no is_global flag is provided
let base_dir = if is_global {
&self.global_memory_dir
} else {
&self.local_memory_dir
};
base_dir.join(format!("{}.txt", category))
}
pub fn retrieve_all(&self, is_global: bool) -> io::Result<HashMap<String, Vec<String>>> {
let base_dir = if is_global {
&self.global_memory_dir
} else {
&self.local_memory_dir
};
let mut memories = HashMap::new();
if base_dir.exists() {
for entry in fs::read_dir(base_dir)? {
let entry = entry?;
if entry.file_type()?.is_file() {
let category = entry.file_name().to_string_lossy().replace(".txt", "");
let category_memories = self.retrieve(&category, is_global)?;
memories.insert(
category,
category_memories.into_iter().flat_map(|(_, v)| v).collect(),
);
}
}
}
Ok(memories)
}
pub fn remember(
&self,
_context: &str,
category: &str,
data: &str,
tags: &[&str],
is_global: bool,
) -> io::Result<()> {
let memory_file_path = self.get_memory_file(category, is_global);
let mut file = fs::OpenOptions::new()
.append(true)
.create(true)
.open(&memory_file_path)?;
if !tags.is_empty() {
writeln!(file, "# {}", tags.join(" "))?;
}
writeln!(file, "{}\n", data)?;
Ok(())
}
pub fn retrieve(
&self,
category: &str,
is_global: bool,
) -> io::Result<HashMap<String, Vec<String>>> {
let memory_file_path = self.get_memory_file(category, is_global);
if !memory_file_path.exists() {
return Ok(HashMap::new());
}
let mut file = fs::File::open(memory_file_path)?;
let mut content = String::new();
file.read_to_string(&mut content)?;
let mut memories = HashMap::new();
for entry in content.split("\n\n") {
let mut lines = entry.lines();
if let Some(first_line) = lines.next() {
if let Some(stripped) = first_line.strip_prefix('#') {
let tags = stripped
.split_whitespace()
.map(String::from)
.collect::<Vec<_>>();
memories.insert(tags.join(" "), lines.map(String::from).collect());
} else {
let entry_data: Vec<String> = std::iter::once(first_line.to_string())
.chain(lines.map(String::from))
.collect();
memories
.entry("untagged".to_string())
.or_insert_with(Vec::new)
.extend(entry_data);
}
}
}
Ok(memories)
}
pub fn remove_specific_memory(
&self,
category: &str,
memory_content: &str,
is_global: bool,
) -> io::Result<()> {
let memory_file_path = self.get_memory_file(category, is_global);
if !memory_file_path.exists() {
return Ok(());
}
let mut file = fs::File::open(&memory_file_path)?;
let mut content = String::new();
file.read_to_string(&mut content)?;
let memories: Vec<&str> = content.split("\n\n").collect();
let new_content: Vec<String> = memories
.into_iter()
.filter(|entry| !entry.contains(memory_content))
.map(|s| s.to_string())
.collect();
fs::write(memory_file_path, new_content.join("\n\n"))?;
Ok(())
}
pub fn clear_memory(&self, category: &str, is_global: bool) -> io::Result<()> {
let memory_file_path = self.get_memory_file(category, is_global);
if memory_file_path.exists() {
fs::remove_file(memory_file_path)?;
}
Ok(())
}
pub fn clear_all_global_or_local_memories(&self, is_global: bool) -> io::Result<()> {
let base_dir = if is_global {
&self.global_memory_dir
} else {
&self.local_memory_dir
};
fs::remove_dir_all(base_dir)?;
Ok(())
}
async fn execute_tool_call(&self, tool_call: ToolCall) -> Result<String, io::Error> {
match tool_call.name.as_str() {
"remember_memory" => {
let args = MemoryArgs::from_value(&tool_call.arguments)?;
let data = args.data.filter(|d| !d.is_empty()).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"Data must exist when remembering a memory",
)
})?;
self.remember("context", args.category, data, &args.tags, args.is_global)?;
Ok(format!("Stored memory in category: {}", args.category))
}
"retrieve_memories" => {
let args = MemoryArgs::from_value(&tool_call.arguments)?;
let memories = if args.category == "*" {
self.retrieve_all(args.is_global)?
} else {
self.retrieve(args.category, args.is_global)?
};
Ok(format!("Retrieved memories: {:?}", memories))
}
"remove_memory_category" => {
let args = MemoryArgs::from_value(&tool_call.arguments)?;
if args.category == "*" {
self.clear_all_global_or_local_memories(args.is_global)?;
Ok(format!(
"Cleared all memory {} categories",
if args.is_global { "global" } else { "local" }
))
} else {
self.clear_memory(args.category, args.is_global)?;
Ok(format!("Cleared memories in category: {}", args.category))
}
}
"remove_specific_memory" => {
let args = MemoryArgs::from_value(&tool_call.arguments)?;
let memory_content = tool_call.arguments["memory_content"].as_str().unwrap();
self.remove_specific_memory(args.category, memory_content, args.is_global)?;
Ok(format!(
"Removed specific memory from category: {}",
args.category
))
}
_ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Unknown tool")),
}
}
}
#[async_trait]
impl Router for MemoryRouter {
fn name(&self) -> String {
"memory".to_string()
}
fn instructions(&self) -> String {
self.instructions.clone()
}
fn capabilities(&self) -> ServerCapabilities {
CapabilitiesBuilder::new().with_tools(false).build()
}
fn list_tools(&self) -> Vec<Tool> {
self.tools.clone()
}
fn call_tool(
&self,
tool_name: &str,
arguments: Value,
) -> Pin<Box<dyn Future<Output = Result<Vec<Content>, ToolError>> + Send + 'static>> {
let this = self.clone();
let tool_name = tool_name.to_string();
Box::pin(async move {
let tool_call = ToolCall {
name: tool_name,
arguments,
};
match this.execute_tool_call(tool_call).await {
Ok(result) => Ok(vec![Content::text(result)]),
Err(err) => Err(ToolError::ExecutionError(err.to_string())),
}
})
}
fn list_resources(&self) -> Vec<Resource> {
Vec::new()
}
fn read_resource(
&self,
_uri: &str,
) -> Pin<Box<dyn Future<Output = Result<String, ResourceError>> + Send + 'static>> {
Box::pin(async move { Ok("".to_string()) })
}
}
#[derive(Debug)]
struct MemoryArgs<'a> {
category: &'a str,
data: Option<&'a str>,
tags: Vec<&'a str>,
is_global: bool,
}
impl<'a> MemoryArgs<'a> {
// Category is required, data is optional, tags are optional, is_global is optional
fn from_value(args: &'a Value) -> Result<Self, io::Error> {
let category = args["category"].as_str().ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "Category must be a string")
})?;
if category.is_empty() {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"Category must be a string",
));
}
let data = args.get("data").and_then(|d| d.as_str());
let tags = match &args["tags"] {
Value::Array(arr) => arr.iter().filter_map(|v| v.as_str()).collect(),
Value::String(s) => vec![s.as_str()],
_ => Vec::new(),
};
let is_global = match &args.get("is_global") {
// Default to false if no is_global flag is provided
Some(Value::Bool(b)) => *b,
Some(Value::String(s)) => s.to_lowercase() == "true",
None => false,
_ => {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"is_global must be a boolean or string 'true'/'false'",
))
}
};
Ok(Self {
category,
data,
tags,
is_global,
})
}
}

View File

@@ -0,0 +1,757 @@
use base64::Engine;
use indoc::{formatdoc, indoc};
use reqwest::{Client, Url};
use serde_json::{json, Value};
use std::{
collections::HashMap, fs, future::Future, os::unix::fs::PermissionsExt, path::PathBuf,
pin::Pin, sync::Arc, sync::Mutex,
};
use tokio::process::Command;
use mcp_core::{
handler::{ResourceError, ToolError},
protocol::ServerCapabilities,
resource::Resource,
tool::Tool,
Content,
};
use mcp_server::router::CapabilitiesBuilder;
use mcp_server::Router;
/// An extension designed for non-developers to help them with common tasks like
/// web scraping, data processing, and automation.
#[derive(Clone)]
pub struct NonDeveloperRouter {
tools: Vec<Tool>,
cache_dir: PathBuf,
active_resources: Arc<Mutex<HashMap<String, Resource>>>,
http_client: Client,
instructions: String,
}
impl Default for NonDeveloperRouter {
fn default() -> Self {
Self::new()
}
}
impl NonDeveloperRouter {
pub fn new() -> Self {
// Create tools for the system
let web_search_tool = Tool::new(
"web_search",
indoc! {r#"
Search the web for a single word (proper noun ideally) using DuckDuckGo's API. Returns results in JSON format.
The results are cached locally for future reference.
Be sparing as there is a limited number of api calls allowed.
"#},
json!({
"type": "object",
"required": ["query"],
"properties": {
"query": {
"type": "string",
"description": "A single word to search for, a topic, propernoun, brand name that you may not know about"
}
}
}),
);
let web_scrape_tool = Tool::new(
"web_scrape",
indoc! {r#"
Fetch and save content from a web page. The content can be saved as:
- text (for HTML pages)
- json (for API responses)
- binary (for images and other files)
The content is cached locally and can be accessed later using the cache_path
returned in the response.
"#},
json!({
"type": "object",
"required": ["url"],
"properties": {
"url": {
"type": "string",
"description": "The URL to fetch content from"
},
"save_as": {
"type": "string",
"enum": ["text", "json", "binary"],
"default": "text",
"description": "How to interpret and save the content"
}
}
}),
);
let computer_control_tool = Tool::new(
"computer_control",
indoc! {r#"
Control the computer using AppleScript (macOS only).
Allows automation of applications and system features through AppleScript.
Common uses:
- Control applications (Mail, Safari, iTunes, etc.)
- System settings and notifications
- UI automation
- File and folder operations
- Calendar and reminders
- Media management
- Combining with the screenshot tool to help user achieve tasks.
It allows users to control applications and system features programmatically.
Here's an overview of what AppleScript can automate:
Application Control
Launch, quit, or manage applications.
Interact with app-specific features (e.g., sending an email in Mail, creating a document in Pages, or editing photos in Preview).
Perform tasks in third-party apps that support AppleScript, such as Adobe Photoshop, Microsoft Office, or Safari.
User Interface Automation
Simulate user interactions like clicking buttons, selecting menu items, or typing text.
Fill out forms or automate repetitive tasks in apps.
System Settings and Utilities
Change system preferences (e.g., volume, screen brightness, Wi-Fi settings).
Automate tasks like shutting down, restarting, or putting the system to sleep.
Monitor system events or logs.
Web Automation
Open specific URLs in AppleScript-enabled browsers.
Automate web interactions (e.g., filling forms, navigating pages).
Scrape information from websites.
Email and Messaging
Automate sending and organizing emails in the Mail app.
Extract email contents or attachments.
Send messages via Messages, Slack etc
Media Management
Organize and edit iTunes/Music libraries (e.g., create playlists, change metadata).
Manage photos in Photos (e.g., creating albums, importing/exporting images).
Automate tasks in video or music production tools like Final Cut Pro or GarageBand.
Data Processing
Interact with spreadsheets (e.g., Numbers or Excel).
"#},
json!({
"type": "object",
"required": ["script"],
"properties": {
"script": {
"type": "string",
"description": "The AppleScript content to execute"
},
"save_output": {
"type": "boolean",
"default": false,
"description": "Whether to save the script output to a file"
}
}
}),
);
let quick_script_tool = Tool::new(
"automation_script",
indoc! {r#"
Create and run small scripts for automation tasks.
Supports Shell and Ruby (on macOS).
The script is saved to a temporary file and executed.
Consider using shell script (bash) for most simple tasks first.
Ruby is useful for text processing or when you need more sophisticated scripting capabilities.
Some examples of shell:
- create a sorted list of unique lines: sort file.txt | uniq
- extract 2nd column in csv: awk -F "," '{ print $2}'
- pattern matching: grep pattern file.txt
"#},
json!({
"type": "object",
"required": ["language", "script"],
"properties": {
"language": {
"type": "string",
"enum": ["shell", "ruby"],
"description": "The scripting language to use"
},
"script": {
"type": "string",
"description": "The script content"
},
"save_output": {
"type": "boolean",
"default": false,
"description": "Whether to save the script output to a file"
}
}
}),
);
let cache_tool = Tool::new(
"cache",
indoc! {r#"
Manage cached files and data:
- list: List all cached files
- view: View content of a cached file
- delete: Delete a cached file
- clear: Clear all cached files
"#},
json!({
"type": "object",
"required": ["command"],
"properties": {
"command": {
"type": "string",
"enum": ["list", "view", "delete", "clear"],
"description": "The command to perform"
},
"path": {
"type": "string",
"description": "Path to the cached file for view/delete commands"
}
}
}),
);
// Create cache directory in user's home directory
let cache_dir = dirs::cache_dir()
.unwrap_or_else(|| PathBuf::from("/tmp"))
.join("goose")
.join("non_developer");
fs::create_dir_all(&cache_dir).unwrap_or_else(|_| {
println!(
"Warning: Failed to create cache directory at {:?}",
cache_dir
)
});
let instructions = formatdoc! {r#"
You are a helpful assistant to a power user who is not a professional developer, but you may use devleopment tools to help assist them.
The user may not know how to break down tasks, so you will need to ensure that you do, and run things in batches as needed.
The NonDeveloperExtension helps you with common tasks like web scraping,
data processing, and automation and computer control without requiring programming expertise,
supplementing the Developer Extension.
You can use scripting as needed to work with text files of data, such as csvs, json, or text files etc.
Using the developer extension is allowed for more sophisticated tasks or instructed to (js or py can be helpful for more complex tasks if tools are available).
Accessing web sites, even apis, may be common (you can use bash scripting to do this) without troubling them too much (they won't know what limits are).
Try to do your best to find ways to complete a task without too many quesitons or offering options unless it is really unclear, find a way if you can.
You can also guide them steps if they can help out as you go along.
There is already a screenshot tool available you can use if needed to see what is on screen.
Here are some extra tools:
automation_script
- Create and run simple automation scripts
- Supports Shell (such as bash), AppleScript (on macos), Ruby (on macos)
- Scripts can save their output to files
- on macos, can use applescript to interact with the desktop, eg calendars, notes and more, anything apple script can do for apps that support it:
AppleScript is a powerful scripting language designed for automating tasks on macOS such as: Integration with Other Scripts
Execute shell scripts, Ruby scripts, or other automation scripts.
Combine workflows across scripting languages.
Complex Workflows
Automate multi-step tasks involving multiple apps or system features.
Create scheduled tasks using Calendar or other scheduling apps.
- use the screenshot tool if needed to help with tasks
computer_control
- Control the computer using AppleScript (macOS only)
- Consider the screenshot tool to work out what is on screen and what to do to help with the control task.
web_search
- Search the web using DuckDuckGo's API for general topics or keywords
web_scrape
- Fetch content from html websites and APIs
- Save as text, JSON, or binary files
- Content is cached locally for later use
- This is not optimised for complex websites, so don't use this as the first tool.
cache
- Manage your cached files
- List, view, delete files
- Clear all cached data
The extension automatically manages:
- Cache directory: {cache_dir}
- File organization and cleanup
"#,
cache_dir = cache_dir.display()
};
Self {
tools: vec![
web_search_tool,
web_scrape_tool,
quick_script_tool,
computer_control_tool,
cache_tool,
],
cache_dir,
active_resources: Arc::new(Mutex::new(HashMap::new())),
http_client: Client::builder().user_agent("Goose/1.0").build().unwrap(),
instructions: instructions.clone(),
}
}
// Helper function to generate a cache file path
fn get_cache_path(&self, prefix: &str, extension: &str) -> PathBuf {
let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S");
self.cache_dir
.join(format!("{}_{}.{}", prefix, timestamp, extension))
}
// Helper function to save content to cache
async fn save_to_cache(
&self,
content: &[u8],
prefix: &str,
extension: &str,
) -> Result<PathBuf, ToolError> {
let cache_path = self.get_cache_path(prefix, extension);
fs::write(&cache_path, content)
.map_err(|e| ToolError::ExecutionError(format!("Failed to write to cache: {}", e)))?;
Ok(cache_path)
}
// Helper function to register a file as a resource
fn register_as_resource(&self, cache_path: &PathBuf, mime_type: &str) -> Result<(), ToolError> {
let uri = Url::from_file_path(cache_path)
.map_err(|_| ToolError::ExecutionError("Invalid cache path".into()))?
.to_string();
let resource = Resource::new(
uri.clone(),
Some(mime_type.to_string()),
Some(cache_path.to_string_lossy().into_owned()),
)
.map_err(|e| ToolError::ExecutionError(e.to_string()))?;
self.active_resources.lock().unwrap().insert(uri, resource);
Ok(())
}
// Implement web_scrape tool functionality
async fn web_search(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let query = params
.get("query")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'query' parameter".into()))?;
// Create the DuckDuckGo API URL
let url = format!(
"https://api.duckduckgo.com/?q={}&format=json&pretty=1",
urlencoding::encode(query)
);
// Fetch the results
let response = self.http_client.get(&url).send().await.map_err(|e| {
ToolError::ExecutionError(format!("Failed to fetch search results: {}", e))
})?;
let status = response.status();
if !status.is_success() {
return Err(ToolError::ExecutionError(format!(
"HTTP request failed with status: {}",
status
)));
}
// Get the JSON response
let json_text = response.text().await.map_err(|e| {
ToolError::ExecutionError(format!("Failed to get response text: {}", e))
})?;
// Save to cache
let cache_path = self
.save_to_cache(json_text.as_bytes(), "search", "json")
.await?;
// Register as a resource
self.register_as_resource(&cache_path, "json")?;
Ok(vec![Content::text(format!(
"Search results saved to: {}",
cache_path.display()
))])
}
async fn web_scrape(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let url = params
.get("url")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'url' parameter".into()))?;
let save_as = params
.get("save_as")
.and_then(|v| v.as_str())
.unwrap_or("text");
// Fetch the content
let response = self
.http_client
.get(url)
.send()
.await
.map_err(|e| ToolError::ExecutionError(format!("Failed to fetch URL: {}", e)))?;
let status = response.status();
if !status.is_success() {
return Err(ToolError::ExecutionError(format!(
"HTTP request failed with status: {}",
status
)));
}
// Process based on save_as parameter
let (content, extension) =
match save_as {
"text" => {
let text = response.text().await.map_err(|e| {
ToolError::ExecutionError(format!("Failed to get text: {}", e))
})?;
(text.into_bytes(), "txt")
}
"json" => {
let text = response.text().await.map_err(|e| {
ToolError::ExecutionError(format!("Failed to get text: {}", e))
})?;
// Verify it's valid JSON
serde_json::from_str::<Value>(&text).map_err(|e| {
ToolError::ExecutionError(format!("Invalid JSON response: {}", e))
})?;
(text.into_bytes(), "json")
}
"binary" => {
let bytes = response.bytes().await.map_err(|e| {
ToolError::ExecutionError(format!("Failed to get bytes: {}", e))
})?;
(bytes.to_vec(), "bin")
}
_ => unreachable!(), // Prevented by enum in tool definition
};
// Save to cache
let cache_path = self.save_to_cache(&content, "web", extension).await?;
// Register as a resource
self.register_as_resource(&cache_path, save_as)?;
Ok(vec![Content::text(format!(
"Content saved to: {}",
cache_path.display()
))])
}
// Implement quick_script tool functionality
async fn quick_script(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let language = params
.get("language")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'language' parameter".into()))?;
let script = params
.get("script")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'script' parameter".into()))?;
let save_output = params
.get("save_output")
.and_then(|v| v.as_bool())
.unwrap_or(false);
// Create a temporary directory for the script
let script_dir = tempfile::tempdir().map_err(|e| {
ToolError::ExecutionError(format!("Failed to create temporary directory: {}", e))
})?;
let command = match language {
"shell" => {
let script_path = script_dir.path().join("script.sh");
fs::write(&script_path, script).map_err(|e| {
ToolError::ExecutionError(format!("Failed to write script: {}", e))
})?;
fs::set_permissions(&script_path, fs::Permissions::from_mode(0o755)).map_err(
|e| {
ToolError::ExecutionError(format!(
"Failed to set script permissions: {}",
e
))
},
)?;
script_path.display().to_string()
}
"ruby" => {
let script_path = script_dir.path().join("script.rb");
fs::write(&script_path, script).map_err(|e| {
ToolError::ExecutionError(format!("Failed to write script: {}", e))
})?;
format!("ruby {}", script_path.display())
}
_ => unreachable!(), // Prevented by enum in tool definition
};
// Run the script
let output = Command::new("bash")
.arg("-c")
.arg(&command)
.output()
.await
.map_err(|e| ToolError::ExecutionError(format!("Failed to run script: {}", e)))?;
let output_str = String::from_utf8_lossy(&output.stdout).into_owned();
let error_str = String::from_utf8_lossy(&output.stderr).into_owned();
let mut result = if output.status.success() {
format!("Script completed successfully.\n\nOutput:\n{}", output_str)
} else {
format!(
"Script failed with error code {}.\n\nError:\n{}\nOutput:\n{}",
output.status, error_str, output_str
)
};
// Save output if requested
if save_output && !output_str.is_empty() {
let cache_path = self
.save_to_cache(output_str.as_bytes(), "script_output", "txt")
.await?;
result.push_str(&format!("\n\nOutput saved to: {}", cache_path.display()));
// Register as a resource
self.register_as_resource(&cache_path, "text")?;
}
Ok(vec![Content::text(result)])
}
// Implement computer control (AppleScript) functionality
async fn computer_control(&self, params: Value) -> Result<Vec<Content>, ToolError> {
if std::env::consts::OS != "macos" {
return Err(ToolError::ExecutionError(
"Computer control (AppleScript) is only supported on macOS".into(),
));
}
let script = params
.get("script")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'script' parameter".into()))?;
let save_output = params
.get("save_output")
.and_then(|v| v.as_bool())
.unwrap_or(false);
// Create a temporary directory for the script
let script_dir = tempfile::tempdir().map_err(|e| {
ToolError::ExecutionError(format!("Failed to create temporary directory: {}", e))
})?;
let script_path = script_dir.path().join("script.scpt");
fs::write(&script_path, script)
.map_err(|e| ToolError::ExecutionError(format!("Failed to write script: {}", e)))?;
let command = format!("osascript {}", script_path.display());
// Run the script
let output = Command::new("bash")
.arg("-c")
.arg(&command)
.output()
.await
.map_err(|e| ToolError::ExecutionError(format!("Failed to run AppleScript: {}", e)))?;
let output_str = String::from_utf8_lossy(&output.stdout).into_owned();
let error_str = String::from_utf8_lossy(&output.stderr).into_owned();
let mut result = if output.status.success() {
format!(
"AppleScript completed successfully.\n\nOutput:\n{}",
output_str
)
} else {
format!(
"AppleScript failed with error code {}.\n\nError:\n{}\nOutput:\n{}",
output.status, error_str, output_str
)
};
// Save output if requested
if save_output && !output_str.is_empty() {
let cache_path = self
.save_to_cache(output_str.as_bytes(), "applescript_output", "txt")
.await?;
result.push_str(&format!("\n\nOutput saved to: {}", cache_path.display()));
// Register as a resource
self.register_as_resource(&cache_path, "text")?;
}
Ok(vec![Content::text(result)])
}
// Implement cache tool functionality
async fn cache(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let command = params
.get("command")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'command' parameter".into()))?;
match command {
"list" => {
let mut files = Vec::new();
for entry in fs::read_dir(&self.cache_dir).map_err(|e| {
ToolError::ExecutionError(format!("Failed to read cache directory: {}", e))
})? {
let entry = entry.map_err(|e| {
ToolError::ExecutionError(format!("Failed to read directory entry: {}", e))
})?;
files.push(format!("{}", entry.path().display()));
}
files.sort();
Ok(vec![Content::text(format!(
"Cached files:\n{}",
files.join("\n")
))])
}
"view" => {
let path = params.get("path").and_then(|v| v.as_str()).ok_or_else(|| {
ToolError::InvalidParameters("Missing 'path' parameter for view".into())
})?;
let content = fs::read_to_string(path).map_err(|e| {
ToolError::ExecutionError(format!("Failed to read file: {}", e))
})?;
Ok(vec![Content::text(format!(
"Content of {}:\n\n{}",
path, content
))])
}
"delete" => {
let path = params.get("path").and_then(|v| v.as_str()).ok_or_else(|| {
ToolError::InvalidParameters("Missing 'path' parameter for delete".into())
})?;
fs::remove_file(path).map_err(|e| {
ToolError::ExecutionError(format!("Failed to delete file: {}", e))
})?;
// Remove from active resources if present
if let Ok(url) = Url::from_file_path(path) {
self.active_resources
.lock()
.unwrap()
.remove(&url.to_string());
}
Ok(vec![Content::text(format!("Deleted file: {}", path))])
}
"clear" => {
fs::remove_dir_all(&self.cache_dir).map_err(|e| {
ToolError::ExecutionError(format!("Failed to clear cache directory: {}", e))
})?;
fs::create_dir_all(&self.cache_dir).map_err(|e| {
ToolError::ExecutionError(format!("Failed to recreate cache directory: {}", e))
})?;
// Clear active resources
self.active_resources.lock().unwrap().clear();
Ok(vec![Content::text("Cache cleared successfully.")])
}
_ => unreachable!(), // Prevented by enum in tool definition
}
}
}
impl Router for NonDeveloperRouter {
fn name(&self) -> String {
"NonDeveloperExtension".to_string()
}
fn instructions(&self) -> String {
self.instructions.clone()
}
fn capabilities(&self) -> ServerCapabilities {
CapabilitiesBuilder::new()
.with_tools(false)
.with_resources(false, false)
.build()
}
fn list_tools(&self) -> Vec<Tool> {
self.tools.clone()
}
fn call_tool(
&self,
tool_name: &str,
arguments: Value,
) -> Pin<Box<dyn Future<Output = Result<Vec<Content>, ToolError>> + Send + 'static>> {
let this = self.clone();
let tool_name = tool_name.to_string();
Box::pin(async move {
match tool_name.as_str() {
"web_search" => this.web_search(arguments).await,
"web_scrape" => this.web_scrape(arguments).await,
"automation_script" => this.quick_script(arguments).await,
"computer_control" => this.computer_control(arguments).await,
"cache" => this.cache(arguments).await,
_ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))),
}
})
}
fn list_resources(&self) -> Vec<Resource> {
let active_resources = self.active_resources.lock().unwrap();
let resources = active_resources.values().cloned().collect();
tracing::info!("Listing resources: {:?}", resources);
resources
}
fn read_resource(
&self,
uri: &str,
) -> Pin<Box<dyn Future<Output = Result<String, ResourceError>> + Send + 'static>> {
let uri = uri.to_string();
let this = self.clone();
Box::pin(async move {
let active_resources = this.active_resources.lock().unwrap();
let resource = active_resources
.get(&uri)
.ok_or_else(|| ResourceError::NotFound(format!("Resource not found: {}", uri)))?
.clone();
let url = Url::parse(&uri)
.map_err(|e| ResourceError::NotFound(format!("Invalid URI: {}", e)))?;
if url.scheme() != "file" {
return Err(ResourceError::NotFound(
"Only file:// URIs are supported".into(),
));
}
let path = url
.to_file_path()
.map_err(|_| ResourceError::NotFound("Invalid file path in URI".into()))?;
match resource.mime_type.as_str() {
"text" | "json" => fs::read_to_string(&path).map_err(|e| {
ResourceError::ExecutionError(format!("Failed to read file: {}", e))
}),
"binary" => {
let bytes = fs::read(&path).map_err(|e| {
ResourceError::ExecutionError(format!("Failed to read file: {}", e))
})?;
Ok(base64::prelude::BASE64_STANDARD.encode(bytes))
}
mime_type => Err(ResourceError::NotFound(format!(
"Unsupported mime type: {}",
mime_type
))),
}
})
}
}

View File

@@ -0,0 +1,40 @@
[package]
name = "goose-server"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
description.workspace = true
[dependencies]
goose = { path = "../goose" }
mcp-core = { path = "../mcp-core" }
goose-mcp = { path = "../goose-mcp" }
mcp-server = { path = "../mcp-server" }
axum = { version = "0.7", features = ["ws"] }
tokio = { version = "1.0", features = ["full"] }
chrono = "0.4"
tower-http = { version = "0.5", features = ["cors"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
futures = "0.3"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json", "time"] }
tracing-appender = "0.2"
tokio-stream = "0.1"
anyhow = "1.0"
bytes = "1.5"
http = "1.0"
config = { version = "0.14.1", features = ["toml"] }
thiserror = "1.0"
clap = { version = "4.4", features = ["derive"] }
once_cell = "1.18"
[[bin]]
name = "goosed"
path = "src/main.rs"
[dev-dependencies]
tower = "0.5"
async-trait = "0.1"

View File

@@ -0,0 +1,34 @@
use crate::configuration;
use crate::state;
use anyhow::Result;
use tower_http::cors::{Any, CorsLayer};
use tracing::info;
pub async fn run() -> Result<()> {
// Initialize logging
crate::logging::setup_logging(Some("goosed"))?;
// Load configuration
let settings = configuration::Settings::new()?;
// load secret key from GOOSE_SERVER__SECRET_KEY environment variable
let secret_key =
std::env::var("GOOSE_SERVER__SECRET_KEY").unwrap_or_else(|_| "test".to_string());
// Create app state - agent will start as None
let state = state::AppState::new(secret_key.clone()).await?;
// Create router with CORS support
let cors = CorsLayer::new()
.allow_origin(Any)
.allow_methods(Any)
.allow_headers(Any);
let app = crate::routes::configure(state).layer(cors);
// Run server
let listener = tokio::net::TcpListener::bind(settings.socket_addr()).await?;
info!("listening on {}", listener.local_addr()?);
axum::serve(listener, app).await?;
Ok(())
}

View File

@@ -0,0 +1,32 @@
use anyhow::Result;
use goose_mcp::{
DeveloperRouter, GoogleDriveRouter, JetBrainsRouter, MemoryRouter, NonDeveloperRouter,
};
use mcp_server::router::RouterService;
use mcp_server::{BoundedService, ByteTransport, Server};
use tokio::io::{stdin, stdout};
pub async fn run(name: &str) -> Result<()> {
// Initialize logging
crate::logging::setup_logging(Some(&format!("mcp-{name}")))?;
tracing::info!("Starting MCP server");
let router: Option<Box<dyn BoundedService>> = match name {
"developer" => Some(Box::new(RouterService(DeveloperRouter::new()))),
"nondeveloper" => Some(Box::new(RouterService(NonDeveloperRouter::new()))),
"jetbrains" => Some(Box::new(RouterService(JetBrainsRouter::new()))),
"google_drive" => {
let router = GoogleDriveRouter::new().await;
Some(Box::new(RouterService(router)))
}
"memory" => Some(Box::new(RouterService(MemoryRouter::new()))),
_ => None,
};
// Create and run the server
let server = Server::new(router.unwrap_or_else(|| panic!("Unknown server requested {}", name)));
let transport = ByteTransport::new(stdin(), stdout());
tracing::info!("Server initialized and ready to handle requests");
Ok(server.run(transport).await?)
}

View File

@@ -0,0 +1,2 @@
pub mod agent;
pub mod mcp;

View File

@@ -0,0 +1,90 @@
use crate::error::{to_env_var, ConfigError};
use config::{Config, Environment};
use serde::Deserialize;
use std::net::SocketAddr;
#[derive(Debug, Default, Deserialize)]
pub struct Settings {
#[serde(default = "default_host")]
pub host: String,
#[serde(default = "default_port")]
pub port: u16,
}
impl Settings {
pub fn socket_addr(&self) -> SocketAddr {
format!("{}:{}", self.host, self.port)
.parse()
.expect("Failed to parse socket address")
}
pub fn new() -> Result<Self, ConfigError> {
Self::load_and_validate()
}
fn load_and_validate() -> Result<Self, ConfigError> {
// Start with default configuration
let config = Config::builder()
// Server defaults
.set_default("host", default_host())?
.set_default("port", default_port())?
// Layer on the environment variables
.add_source(
Environment::with_prefix("GOOSE")
.prefix_separator("_")
.separator("__")
.try_parsing(true),
)
.build()?;
// Try to deserialize the configuration
let result: Result<Self, config::ConfigError> = config.try_deserialize();
// Handle missing field errors specially
match result {
Ok(settings) => Ok(settings),
Err(err) => {
tracing::debug!("Configuration error: {:?}", &err);
// Handle both NotFound and missing field message variants
let error_str = err.to_string();
if error_str.starts_with("missing field") {
// Extract field name from error message "missing field `type`"
let field = error_str
.trim_start_matches("missing field `")
.trim_end_matches("`");
let env_var = to_env_var(field);
Err(ConfigError::MissingEnvVar { env_var })
} else if let config::ConfigError::NotFound(field) = &err {
let env_var = to_env_var(field);
Err(ConfigError::MissingEnvVar { env_var })
} else {
Err(ConfigError::Other(err))
}
}
}
}
}
fn default_host() -> String {
"127.0.0.1".to_string()
}
fn default_port() -> u16 {
3000
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_socket_addr_conversion() {
let server_settings = Settings {
host: "127.0.0.1".to_string(),
port: 3000,
};
let addr = server_settings.socket_addr();
assert_eq!(addr.to_string(), "127.0.0.1:3000");
}
}

View File

@@ -0,0 +1,40 @@
use thiserror::Error;
#[derive(Error, Debug)]
pub enum ConfigError {
#[error("Missing required environment variable: {env_var}")]
MissingEnvVar { env_var: String },
#[error("Configuration error: {0}")]
Other(#[from] config::ConfigError),
}
// Helper function to format environment variable names
pub(crate) fn to_env_var(field_path: &str) -> String {
// Handle nested fields by converting dots to double underscores
// If the field is in the provider object, we need to prefix it appropriately
let normalized_path = if field_path == "type" {
"provider.type".to_string()
} else if field_path.starts_with("provider.") {
field_path.to_string()
} else {
format!("provider.{}", field_path)
};
format!(
"GOOSE_{}",
normalized_path.replace('.', "__").to_uppercase()
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_env_var_conversion() {
assert_eq!(to_env_var("type"), "GOOSE_PROVIDER__TYPE");
assert_eq!(to_env_var("api_key"), "GOOSE_PROVIDER__API_KEY");
assert_eq!(to_env_var("provider.host"), "GOOSE_PROVIDER__HOST");
assert_eq!(to_env_var("provider.api_key"), "GOOSE_PROVIDER__API_KEY");
}
}

View File

@@ -0,0 +1,106 @@
use anyhow::{Context, Result};
use std::fs;
use std::path::PathBuf;
use tracing_appender::rolling::Rotation;
use tracing_subscriber::{
filter::LevelFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer,
Registry,
};
use goose::tracing::langfuse_layer;
/// Returns the directory where log files should be stored.
/// Creates the directory structure if it doesn't exist.
fn get_log_directory() -> Result<PathBuf> {
let home = std::env::var("HOME").context("HOME environment variable not set")?;
let base_log_dir = PathBuf::from(home)
.join(".config")
.join("goose")
.join("logs")
.join("server"); // Add server-specific subdirectory
// Create date-based subdirectory
let now = chrono::Local::now();
let date_dir = base_log_dir.join(now.format("%Y-%m-%d").to_string());
// Ensure log directory exists
fs::create_dir_all(&date_dir).context("Failed to create log directory")?;
Ok(date_dir)
}
/// Sets up the logging infrastructure for the application.
/// This includes:
/// - File-based logging with JSON formatting (DEBUG level)
/// - Console output for development (INFO level)
/// - Optional Langfuse integration (DEBUG level)
pub fn setup_logging(name: Option<&str>) -> Result<()> {
// Set up file appender for goose module logs
let log_dir = get_log_directory()?;
let timestamp = chrono::Local::now().format("%Y%m%d_%H%M%S").to_string();
// Create log file name by prefixing with timestamp
let log_filename = if name.is_some() {
format!("{}-{}.log", timestamp, name.unwrap())
} else {
format!("{}.log", timestamp)
};
// Create non-rolling file appender for detailed logs
let file_appender =
tracing_appender::rolling::RollingFileAppender::new(Rotation::NEVER, log_dir, log_filename);
// Create JSON file logging layer
let file_layer = fmt::layer()
.with_target(true)
.with_level(true)
.with_writer(file_appender)
.with_ansi(false)
.with_file(true);
// Create console logging layer for development - INFO and above only
let console_layer = fmt::layer()
.with_target(true)
.with_level(true)
.with_ansi(true)
.with_file(true)
.with_line_number(true)
.pretty();
// Base filter for all logging
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| {
// Set default levels for different modules
EnvFilter::new("")
// Set mcp-server module to DEBUG
.add_directive("mcp_server=debug".parse().unwrap())
// Set mcp-client to DEBUG
.add_directive("mcp_client=debug".parse().unwrap())
// Set goose module to DEBUG
.add_directive("goose=debug".parse().unwrap())
// Set goose-server to INFO
.add_directive("goose_server=info".parse().unwrap())
// Set tower-http to INFO for request logging
.add_directive("tower_http=info".parse().unwrap())
// Set everything else to WARN
.add_directive(LevelFilter::WARN.into())
});
// Build the subscriber with required layers
let subscriber = Registry::default()
.with(file_layer.with_filter(env_filter))
.with(console_layer.with_filter(LevelFilter::INFO));
// Initialize with Langfuse if available
if let Some(langfuse) = langfuse_layer::create_langfuse_observer() {
subscriber
.with(langfuse.with_filter(LevelFilter::DEBUG))
.try_init()
.context("Failed to set global subscriber")?;
} else {
subscriber
.try_init()
.context("Failed to set global subscriber")?;
}
Ok(())
}

View File

@@ -0,0 +1,43 @@
mod commands;
mod configuration;
mod error;
mod logging;
mod routes;
mod state;
use clap::{Parser, Subcommand};
#[derive(Parser)]
#[command(author, version, about, long_about = None)]
#[command(propagate_version = true)]
struct Cli {
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand)]
enum Commands {
/// Run the agent server
Agent,
/// Run the MCP server
Mcp {
/// Name of the MCP server type
name: String,
},
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let cli = Cli::parse();
match &cli.command {
Commands::Agent => {
commands::agent::run().await?;
}
Commands::Mcp { name } => {
commands::mcp::run(name).await?;
}
}
Ok(())
}

View File

@@ -0,0 +1,137 @@
use crate::state::AppState;
use axum::{
extract::State,
http::{HeaderMap, StatusCode},
routing::{get, post},
Json, Router,
};
use goose::config::Config;
use goose::{agents::AgentFactory, model::ModelConfig, providers};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::env;
#[derive(Serialize)]
struct VersionsResponse {
available_versions: Vec<String>,
default_version: String,
}
#[derive(Deserialize)]
struct CreateAgentRequest {
version: Option<String>,
provider: String,
model: Option<String>,
}
#[derive(Serialize)]
struct CreateAgentResponse {
version: String,
}
#[derive(Deserialize)]
struct ProviderFile {
name: String,
description: String,
models: Vec<String>,
required_keys: Vec<String>,
}
#[derive(Serialize)]
struct ProviderDetails {
name: String,
description: String,
models: Vec<String>,
required_keys: Vec<String>,
}
#[derive(Serialize)]
struct ProviderList {
id: String,
details: ProviderDetails,
}
async fn get_versions() -> Json<VersionsResponse> {
let versions = AgentFactory::available_versions();
let default_version = AgentFactory::default_version().to_string();
Json(VersionsResponse {
available_versions: versions.iter().map(|v| v.to_string()).collect(),
default_version,
})
}
async fn create_agent(
State(state): State<AppState>,
headers: HeaderMap,
Json(payload): Json<CreateAgentRequest>,
) -> Result<Json<CreateAgentResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
// Set the environment variable for the model if provided
if let Some(model) = &payload.model {
let env_var_key = format!("{}_MODEL", payload.provider.to_uppercase());
env::set_var(env_var_key.clone(), model);
println!("Set environment variable: {}={}", env_var_key, model);
}
let config = Config::global();
let model = payload.model.unwrap_or_else(|| {
config
.get("GOOSE_MODEL")
.expect("Did not find a model on payload or in env")
});
let model_config = ModelConfig::new(model);
let provider =
providers::create(&payload.provider, model_config).expect("Failed to create provider");
let version = payload
.version
.unwrap_or_else(|| AgentFactory::default_version().to_string());
let new_agent = AgentFactory::create(&version, provider).expect("Failed to create agent");
let mut agent = state.agent.lock().await;
*agent = Some(new_agent);
Ok(Json(CreateAgentResponse { version }))
}
async fn list_providers() -> Json<Vec<ProviderList>> {
let contents = include_str!("providers_and_keys.json");
let providers: HashMap<String, ProviderFile> =
serde_json::from_str(contents).expect("Failed to parse providers_and_keys.json");
let response: Vec<ProviderList> = providers
.into_iter()
.map(|(id, provider)| ProviderList {
id,
details: ProviderDetails {
name: provider.name,
description: provider.description,
models: provider.models,
required_keys: provider.required_keys,
},
})
.collect();
// Return the response as JSON.
Json(response)
}
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/agent/versions", get(get_versions))
.route("/agent/providers", get(list_providers))
.route("/agent", post(create_agent))
.with_state(state)
}

View File

@@ -0,0 +1,208 @@
use std::collections::HashMap;
use crate::state::AppState;
use axum::{extract::State, routing::post, Json, Router};
use goose::{
agents::{extension::Envs, ExtensionConfig},
config::Config,
};
use http::{HeaderMap, StatusCode};
use serde::{Deserialize, Serialize};
/// Enum representing the different types of extension configuration requests.
#[derive(Deserialize)]
#[serde(tag = "type")]
enum ExtensionConfigRequest {
/// Server-Sent Events (SSE) extension.
#[serde(rename = "sse")]
Sse {
/// The name to identify this extension
name: String,
/// The URI endpoint for the SSE extension.
uri: String,
/// List of environment variable keys. The server will fetch their values from the keyring.
env_keys: Vec<String>,
},
/// Standard I/O (stdio) extension.
#[serde(rename = "stdio")]
Stdio {
/// The name to identify this extension
name: String,
/// The command to execute.
cmd: String,
/// Arguments for the command.
args: Vec<String>,
/// List of environment variable keys. The server will fetch their values from the keyring.
env_keys: Vec<String>,
},
/// Built-in extension that is part of the goose binary.
#[serde(rename = "builtin")]
Builtin {
/// The name of the built-in extension.
name: String,
},
}
/// Response structure for adding an extension.
///
/// - `error`: Indicates whether an error occurred (`true`) or not (`false`).
/// - `message`: Provides detailed error information when `error` is `true`.
#[derive(Serialize)]
struct ExtensionResponse {
error: bool,
message: Option<String>,
}
/// Handler for adding a new extension configuration.
async fn add_extension(
State(state): State<AppState>,
headers: HeaderMap,
Json(request): Json<ExtensionConfigRequest>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
// Verify the presence and validity of the secret key.
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
// Load the configuration
let config = Config::global();
// Initialize a vector to collect any missing keys.
let mut missing_keys = Vec::new();
// Construct ExtensionConfig with Envs populated from keyring based on provided env_keys.
let extension_config: ExtensionConfig = match request {
ExtensionConfigRequest::Sse {
name,
uri,
env_keys,
} => {
let mut env_map = HashMap::new();
for key in env_keys {
match config.get_secret(&key) {
Ok(value) => {
env_map.insert(key, value);
}
Err(_) => {
missing_keys.push(key);
}
}
}
if !missing_keys.is_empty() {
return Ok(Json(ExtensionResponse {
error: true,
message: Some(format!(
"Missing secrets for keys: {}",
missing_keys.join(", ")
)),
}));
}
ExtensionConfig::Sse {
name,
uri,
envs: Envs::new(env_map),
}
}
ExtensionConfigRequest::Stdio {
name,
cmd,
args,
env_keys,
} => {
let mut env_map = HashMap::new();
for key in env_keys {
match config.get_secret(&key) {
Ok(value) => {
env_map.insert(key, value);
}
Err(_) => {
missing_keys.push(key);
}
}
}
if !missing_keys.is_empty() {
return Ok(Json(ExtensionResponse {
error: true,
message: Some(format!(
"Missing secrets for keys: {}",
missing_keys.join(", ")
)),
}));
}
ExtensionConfig::Stdio {
name,
cmd,
args,
envs: Envs::new(env_map),
}
}
ExtensionConfigRequest::Builtin { name } => ExtensionConfig::Builtin { name },
};
// Acquire a lock on the agent and attempt to add the extension.
let mut agent = state.agent.lock().await;
let agent = agent.as_mut().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
let response = agent.add_extension(extension_config).await;
// Respond with the result.
match response {
Ok(_) => Ok(Json(ExtensionResponse {
error: false,
message: None,
})),
Err(e) => {
eprintln!("Failed to add extension configuration: {:?}", e);
Ok(Json(ExtensionResponse {
error: true,
message: Some(format!(
"Failed to add extension configuration, error: {:?}",
e
)),
}))
}
}
}
/// Handler for removing an extension by name
async fn remove_extension(
State(state): State<AppState>,
headers: HeaderMap,
Json(name): Json<String>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
// Verify the presence and validity of the secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
// Acquire a lock on the agent and attempt to remove the extension
let mut agent = state.agent.lock().await;
let agent = agent.as_mut().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
agent.remove_extension(&name).await;
Ok(Json(ExtensionResponse {
error: false,
message: None,
}))
}
/// Registers the extension management routes with the Axum router.
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/extensions/add", post(add_extension))
.route("/extensions/remove", post(remove_extension))
.with_state(state)
}

View File

@@ -0,0 +1,17 @@
use axum::{routing::get, Json, Router};
use serde::Serialize;
#[derive(Serialize)]
struct StatusResponse {
status: &'static str,
}
/// Simple status endpoint that returns 200 OK when the server is running
async fn status() -> Json<StatusResponse> {
Json(StatusResponse { status: "ok" })
}
/// Configure health check routes
pub fn routes() -> Router {
Router::new().route("/status", get(status))
}

View File

@@ -0,0 +1,18 @@
// Export route modules
pub mod agent;
pub mod extension;
pub mod health;
pub mod reply;
pub mod secrets;
use axum::Router;
// Function to configure all routes
pub fn configure(state: crate::state::AppState) -> Router {
Router::new()
.merge(health::routes())
.merge(reply::routes(state.clone()))
.merge(agent::routes(state.clone()))
.merge(extension::routes(state.clone()))
.merge(secrets::routes(state))
}

View File

@@ -0,0 +1,44 @@
{
"openai": {
"name": "OpenAI",
"description": "Use GPT-4 and other OpenAI models",
"models": ["gpt-4o", "gpt-4-turbo","o1"],
"required_keys": ["OPENAI_API_KEY"]
},
"anthropic": {
"name": "Anthropic",
"description": "Use Claude and other Anthropic models",
"models": ["claude-3.5-sonnet-2"],
"required_keys": ["ANTHROPIC_API_KEY"]
},
"databricks": {
"name": "Databricks",
"description": "Connect to LLMs via Databricks",
"models": ["claude-3-5-sonnet-2"],
"required_keys": ["DATABRICKS_HOST"]
},
"google": {
"name": "Google",
"description": "Lorem ipsum",
"models": ["gemini-1.5-flash"],
"required_keys": ["GOOGLE_API_KEY"]
},
"grok": {
"name": "Groq",
"description": "Lorem ipsum",
"models": ["llama-3.3-70b-versatile"],
"required_keys": ["GROQ_API_KEY"]
},
"ollama": {
"name": "Ollama",
"description": "Lorem ipsum",
"models": ["qwen2.5"],
"required_keys": []
},
"openrouter": {
"name": "OpenRouter",
"description": "Lorem ipsum",
"models": [],
"required_keys": ["OPENROUTER_API_KEY"]
}
}

View File

@@ -0,0 +1,604 @@
use crate::state::AppState;
use axum::{
extract::State,
http::{self, HeaderMap, StatusCode},
response::IntoResponse,
routing::post,
Json, Router,
};
use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::message::{Message, MessageContent};
use mcp_core::{content::Content, role::Role};
use serde::Deserialize;
use serde_json::{json, Value};
use std::{
convert::Infallible,
pin::Pin,
task::{Context, Poll},
time::Duration,
};
use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_stream::wrappers::ReceiverStream;
// Types matching the incoming JSON structure
#[derive(Debug, Deserialize)]
struct ChatRequest {
messages: Vec<IncomingMessage>,
}
#[derive(Debug, Deserialize)]
struct IncomingMessage {
role: String,
content: String,
#[serde(default)]
#[serde(rename = "toolInvocations")]
tool_invocations: Vec<ToolInvocation>,
}
#[derive(Debug, Deserialize)]
struct ToolInvocation {
state: String,
#[serde(rename = "toolCallId")]
tool_call_id: String,
#[serde(rename = "toolName")]
tool_name: String,
args: Value,
result: Option<Vec<Content>>,
}
// Custom SSE response type that implements the Vercel AI SDK protocol
pub struct SseResponse {
rx: ReceiverStream<String>,
}
impl SseResponse {
fn new(rx: ReceiverStream<String>) -> Self {
Self { rx }
}
}
impl Stream for SseResponse {
type Item = Result<Bytes, Infallible>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.rx)
.poll_next(cx)
.map(|opt| opt.map(|s| Ok(Bytes::from(s))))
}
}
impl IntoResponse for SseResponse {
fn into_response(self) -> axum::response::Response {
let stream = self;
let body = axum::body::Body::from_stream(stream);
http::Response::builder()
.header("Content-Type", "text/event-stream")
.header("Cache-Control", "no-cache")
.header("Connection", "keep-alive")
.header("x-vercel-ai-data-stream", "v1")
.body(body)
.unwrap()
}
}
// Convert incoming messages to our internal Message type
fn convert_messages(incoming: Vec<IncomingMessage>) -> Vec<Message> {
let mut messages = Vec::new();
for msg in incoming {
match msg.role.as_str() {
"user" => {
messages.push(Message::user().with_text(msg.content));
}
"assistant" => {
// First handle any tool invocations - each represents a complete request/response cycle
for tool in msg.tool_invocations {
if tool.state == "result" {
// Add the original tool request from assistant
let tool_call = mcp_core::tool::ToolCall {
name: tool.tool_name,
arguments: tool.args,
};
messages.push(
Message::assistant()
.with_tool_request(tool.tool_call_id.clone(), Ok(tool_call)),
);
// Add the tool response from user
if let Some(result) = &tool.result {
messages.push(
Message::user()
.with_tool_response(tool.tool_call_id, Ok(result.clone())),
);
}
}
}
// Then add the assistant's text response after tool interactions
if !msg.content.is_empty() {
messages.push(Message::assistant().with_text(msg.content));
}
}
_ => {
tracing::warn!("Unknown role: {}", msg.role);
}
}
}
messages
}
// Protocol-specific message formatting
struct ProtocolFormatter;
impl ProtocolFormatter {
fn format_text(text: &str) -> String {
let encoded_text = serde_json::to_string(text).unwrap_or_else(|_| String::new());
format!("0:{}\n", encoded_text)
}
fn format_tool_call(id: &str, name: &str, args: &Value) -> String {
// Tool calls start with "9:"
let tool_call = json!({
"toolCallId": id,
"toolName": name,
"args": args
});
format!("9:{}\n", tool_call)
}
fn format_tool_response(id: &str, result: &Vec<Content>) -> String {
// Tool responses start with "a:"
let response = json!({
"toolCallId": id,
"result": result,
});
format!("a:{}\n", response)
}
fn format_error(error: &str) -> String {
// Error messages start with "3:" in the new protocol.
let encoded_error = serde_json::to_string(error).unwrap_or_else(|_| String::new());
format!("3:{}\n", encoded_error)
}
fn format_finish(reason: &str) -> String {
// Finish messages start with "d:"
let finish = json!({
"finishReason": reason,
"usage": {
"promptTokens": 0,
"completionTokens": 0
}
});
format!("d:{}\n", finish)
}
}
async fn stream_message(
message: Message,
tx: &mpsc::Sender<String>,
) -> Result<(), mpsc::error::SendError<String>> {
match message.role {
Role::User => {
// Handle tool responses
for content in message.content {
// I believe with the protocol we aren't intended to pass back user messages, so we only deal with
// the tool responses here
if let MessageContent::ToolResponse(response) = content {
// We should return a result for either an error or a success
match response.tool_result {
Ok(result) => {
tx.send(ProtocolFormatter::format_tool_response(
&response.id,
&result,
))
.await?;
}
Err(err) => {
// Send an error message first
tx.send(ProtocolFormatter::format_error(&err.to_string()))
.await?;
// Then send an empty tool response to maintain the protocol
let result =
vec![Content::text(format!("Error: {}", err)).with_priority(0.0)];
tx.send(ProtocolFormatter::format_tool_response(
&response.id,
&result,
))
.await?;
}
}
}
}
}
Role::Assistant => {
for content in message.content {
match content {
MessageContent::ToolRequest(request) => {
match request.tool_call {
Ok(tool_call) => {
tx.send(ProtocolFormatter::format_tool_call(
&request.id,
&tool_call.name,
&tool_call.arguments,
))
.await?;
}
Err(err) => {
// Send a placeholder tool call to maintain protocol
tx.send(ProtocolFormatter::format_tool_call(
&request.id,
"invalid_tool",
&json!({"error": err.to_string()}),
))
.await?;
}
}
}
MessageContent::Text(text) => {
for line in text.text.lines() {
let modified_line = format!("{}\n", line);
tx.send(ProtocolFormatter::format_text(&modified_line))
.await?;
}
}
MessageContent::Image(_) => {
// TODO
continue;
}
MessageContent::ToolResponse(_) => {
// Tool responses should only come from the user
continue;
}
}
}
}
}
Ok(())
}
async fn handler(
State(state): State<AppState>,
headers: HeaderMap,
Json(request): Json<ChatRequest>,
) -> Result<SseResponse, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
// Check protocol header (optional in our case)
if let Some(protocol) = headers.get("x-protocol") {
if protocol.to_str().map(|p| p != "data").unwrap_or(true) {
return Err(StatusCode::BAD_REQUEST);
}
}
// Create channel for streaming
let (tx, rx) = mpsc::channel(100);
let stream = ReceiverStream::new(rx);
// Convert incoming messages
let messages = convert_messages(request.messages);
// Get a lock on the shared agent
let agent = state.agent.clone();
// Spawn task to handle streaming
tokio::spawn(async move {
let agent = agent.lock().await;
let agent = match agent.as_ref() {
Some(agent) => agent,
None => {
let _ = tx
.send(ProtocolFormatter::format_error("No agent configured"))
.await;
let _ = tx.send(ProtocolFormatter::format_finish("error")).await;
return;
}
};
let mut stream = match agent.reply(&messages).await {
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {}", e);
let _ = tx
.send(ProtocolFormatter::format_error(&e.to_string()))
.await;
let _ = tx.send(ProtocolFormatter::format_finish("error")).await;
return;
}
};
loop {
tokio::select! {
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(message))) => {
if let Err(e) = stream_message(message, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await;
break;
}
}
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await;
break;
}
Ok(None) => {
break;
}
Err(_) => { // Heartbeat, used to detect disconnected clients and then end running tools.
if tx.is_closed() {
// Kill any running processes when the client disconnects
// TODO is this used? I suspect post MCP this is on the server instead
// goose::process_store::kill_processes();
break;
}
continue;
}
}
}
}
}
// Send finish message
let _ = tx.send(ProtocolFormatter::format_finish("stop")).await;
});
Ok(SseResponse::new(stream))
}
#[derive(Debug, Deserialize, serde::Serialize)]
struct AskRequest {
prompt: String,
}
#[derive(Debug, serde::Serialize)]
struct AskResponse {
response: String,
}
// simple ask an AI for a response, non streaming
async fn ask_handler(
State(state): State<AppState>,
headers: HeaderMap,
Json(request): Json<AskRequest>,
) -> Result<Json<AskResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
let agent = state.agent.clone();
let agent = agent.lock().await;
let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
// Create a single message for the prompt
let messages = vec![Message::user().with_text(request.prompt)];
// Get response from agent
let mut response_text = String::new();
let mut stream = match agent.reply(&messages).await {
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
while let Some(response) = stream.next().await {
match response {
Ok(message) => {
if message.role == Role::Assistant {
for content in message.content {
if let MessageContent::Text(text) = content {
response_text.push_str(&text.text);
response_text.push('\n');
}
}
}
}
Err(e) => {
tracing::error!("Error processing as_ai message: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
}
}
Ok(Json(AskResponse {
response: response_text.trim().to_string(),
}))
}
// Configure routes for this module
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/reply", post(handler))
.route("/ask", post(ask_handler))
.with_state(state)
}
#[cfg(test)]
mod tests {
use super::*;
use goose::{
agents::AgentFactory,
model::ModelConfig,
providers::{
base::{Provider, ProviderUsage, Usage},
errors::ProviderError,
},
};
use mcp_core::tool::Tool;
// Mock Provider implementation for testing
#[derive(Clone)]
struct MockProvider {
model_config: ModelConfig,
}
#[async_trait::async_trait]
impl Provider for MockProvider {
fn metadata() -> goose::providers::base::ProviderMetadata {
goose::providers::base::ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete(
&self,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
Ok((
Message::assistant().with_text("Mock response"),
ProviderUsage::new("mock".to_string(), Usage::default()),
))
}
}
#[test]
fn test_convert_messages_user_only() {
let incoming = vec![IncomingMessage {
role: "user".to_string(),
content: "Hello".to_string(),
tool_invocations: vec![],
}];
let messages = convert_messages(incoming);
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, Role::User);
assert!(
matches!(&messages[0].content[0], MessageContent::Text(text) if text.text == "Hello")
);
}
#[test]
fn test_convert_messages_with_tool_invocation() {
let tool_result = vec![Content::text("tool response").with_priority(0.0)];
let incoming = vec![IncomingMessage {
role: "assistant".to_string(),
content: "".to_string(),
tool_invocations: vec![ToolInvocation {
state: "result".to_string(),
tool_call_id: "123".to_string(),
tool_name: "test_tool".to_string(),
args: json!({"key": "value"}),
result: Some(tool_result.clone()),
}],
}];
let messages = convert_messages(incoming);
assert_eq!(messages.len(), 2); // Tool request and response
// Check tool request
assert_eq!(messages[0].role, Role::Assistant);
assert!(
matches!(&messages[0].content[0], MessageContent::ToolRequest(req) if req.id == "123")
);
// Check tool response
assert_eq!(messages[1].role, Role::User);
assert!(
matches!(&messages[1].content[0], MessageContent::ToolResponse(resp) if resp.id == "123")
);
}
#[test]
fn test_protocol_formatter() {
// Test text formatting
let text = "Hello world";
let formatted = ProtocolFormatter::format_text(text);
assert_eq!(formatted, "0:\"Hello world\"\n");
// Test tool call formatting
let formatted =
ProtocolFormatter::format_tool_call("123", "test_tool", &json!({"key": "value"}));
assert!(formatted.starts_with("9:"));
assert!(formatted.contains("\"toolCallId\":\"123\""));
assert!(formatted.contains("\"toolName\":\"test_tool\""));
// Test tool response formatting
let result = vec![Content::text("response").with_priority(0.0)];
let formatted = ProtocolFormatter::format_tool_response("123", &result);
assert!(formatted.starts_with("a:"));
assert!(formatted.contains("\"toolCallId\":\"123\""));
// Test error formatting
let formatted = ProtocolFormatter::format_error("Test error");
println!("Formatted error: {}", formatted);
assert!(formatted.starts_with("3:"));
assert!(formatted.contains("Test error"));
// Test finish formatting
let formatted = ProtocolFormatter::format_finish("stop");
assert!(formatted.starts_with("d:"));
assert!(formatted.contains("\"finishReason\":\"stop\""));
}
mod integration_tests {
use super::*;
use axum::{body::Body, http::Request};
use std::sync::Arc;
use tokio::sync::Mutex;
use tower::ServiceExt;
// This test requires tokio runtime
#[tokio::test]
async fn test_ask_endpoint() {
// Create a mock app state with mock provider
let mock_model_config = ModelConfig::new("test-model".to_string());
let mock_provider = Box::new(MockProvider {
model_config: mock_model_config,
});
let agent = AgentFactory::create("reference", mock_provider).unwrap();
let state = AppState {
agent: Arc::new(Mutex::new(Some(agent))),
secret_key: "test-secret".to_string(),
};
// Build router
let app = routes(state);
// Create request
let request = Request::builder()
.uri("/ask")
.method("POST")
.header("content-type", "application/json")
.header("x-secret-key", "test-secret")
.body(Body::from(
serde_json::to_string(&AskRequest {
prompt: "test prompt".to_string(),
})
.unwrap(),
))
.unwrap();
// Send request
let response = app.oneshot(request).await.unwrap();
// Assert response status
assert_eq!(response.status(), StatusCode::OK);
}
}
}

View File

@@ -0,0 +1,191 @@
use crate::state::AppState;
use axum::{extract::State, routing::delete, routing::post, Json, Router};
use goose::config::Config;
use http::{HeaderMap, StatusCode};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
#[derive(Serialize)]
struct SecretResponse {
error: bool,
}
#[derive(Deserialize)]
struct SecretRequest {
key: String,
value: String,
}
async fn store_secret(
State(state): State<AppState>,
headers: HeaderMap,
Json(request): Json<SecretRequest>,
) -> Result<Json<SecretResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
match Config::global().set_secret(&request.key, Value::String(request.value)) {
Ok(_) => Ok(Json(SecretResponse { error: false })),
Err(_) => Ok(Json(SecretResponse { error: true })),
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ProviderSecretRequest {
pub providers: Vec<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SecretStatus {
pub is_set: bool,
pub location: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ProviderResponse {
pub supported: bool,
pub name: Option<String>,
pub description: Option<String>,
pub models: Option<Vec<String>>,
pub secret_status: HashMap<String, SecretStatus>,
}
#[derive(Debug, Serialize, Deserialize)]
struct ProviderConfig {
name: String,
description: String,
models: Vec<String>,
required_keys: Vec<String>,
}
static PROVIDER_ENV_REQUIREMENTS: Lazy<HashMap<String, ProviderConfig>> = Lazy::new(|| {
let contents = include_str!("providers_and_keys.json");
serde_json::from_str(contents).expect("Failed to parse providers_and_keys.json")
});
fn check_key_status(key: &str) -> (bool, Option<String>) {
if let Ok(_value) = std::env::var(key) {
(true, Some("env".to_string()))
} else if Config::global().get_secret::<String>(key).is_ok() {
(true, Some("keyring".to_string()))
} else {
(false, None)
}
}
async fn check_provider_secrets(
Json(request): Json<ProviderSecretRequest>,
) -> Result<Json<HashMap<String, ProviderResponse>>, StatusCode> {
let mut response = HashMap::new();
for provider_name in request.providers {
if let Some(provider_config) = PROVIDER_ENV_REQUIREMENTS.get(&provider_name) {
let mut secret_status = HashMap::new();
for key in &provider_config.required_keys {
let (key_set, key_location) = check_key_status(key);
secret_status.insert(
key.to_string(),
SecretStatus {
is_set: key_set,
location: key_location,
},
);
}
response.insert(
provider_name,
ProviderResponse {
supported: true,
name: Some(provider_config.name.clone()),
description: Some(provider_config.description.clone()),
models: Some(provider_config.models.clone()),
secret_status,
},
);
} else {
response.insert(
provider_name,
ProviderResponse {
supported: false,
name: None,
description: None,
models: None,
secret_status: HashMap::new(),
},
);
}
}
Ok(Json(response))
}
#[derive(Deserialize)]
struct DeleteSecretRequest {
key: String,
}
async fn delete_secret(
State(state): State<AppState>,
headers: HeaderMap,
Json(request): Json<DeleteSecretRequest>,
) -> Result<StatusCode, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
// Attempt to delete the key
match Config::global().delete_secret(&request.key) {
Ok(_) => Ok(StatusCode::NO_CONTENT),
Err(_) => Err(StatusCode::NOT_FOUND),
}
}
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/secrets/providers", post(check_provider_secrets))
.route("/secrets/store", post(store_secret))
.route("/secrets/delete", delete(delete_secret))
.with_state(state)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_unsupported_provider() {
// Setup
let request = ProviderSecretRequest {
providers: vec!["unsupported_provider".to_string()],
};
// Execute
let result = check_provider_secrets(Json(request)).await;
// Assert
assert!(result.is_ok());
let Json(response) = result.unwrap();
let provider_status = response
.get("unsupported_provider")
.expect("Provider should exist");
assert!(!provider_status.supported);
assert!(provider_status.secret_status.is_empty());
}
}

View File

@@ -0,0 +1,21 @@
use anyhow::Result;
use goose::agents::Agent;
use std::sync::Arc;
use tokio::sync::Mutex;
/// Shared application state
#[allow(dead_code)]
#[derive(Clone)]
pub struct AppState {
pub agent: Arc<Mutex<Option<Box<dyn Agent>>>>,
pub secret_key: String,
}
impl AppState {
pub async fn new(secret_key: String) -> Result<Self> {
Ok(Self {
agent: Arc::new(Mutex::new(None)),
secret_key,
})
}
}

1
crates/goose/.gitignore vendored Normal file
View File

@@ -0,0 +1 @@
.env

78
crates/goose/Cargo.toml Normal file
View File

@@ -0,0 +1,78 @@
[package]
name = "goose"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
description.workspace = true
[build-dependencies]
tokio = { version = "1.36", features = ["full"] }
reqwest = { version = "0.12.9", features = ["json", "rustls-tls"], default-features = false }
[dependencies]
mcp-client = { path = "../mcp-client" }
mcp-core = { path = "../mcp-core" }
anyhow = "1.0"
thiserror = "1.0"
futures = "0.3"
reqwest = { version = "0.12.9", features = [
"rustls-tls",
"json",
"cookies",
"gzip",
"brotli",
"deflate",
"zstd",
"charset",
"http2",
"stream"
], default-features = false }
tokio = { version = "1.0", features = ["full"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde_urlencoded = "0.7"
uuid = { version = "1.0", features = ["v4"] }
regex = "1.11.1"
async-trait = "0.1"
async-stream = "0.3"
tera = "1.20.0"
tokenizers = "0.20.3"
include_dir = "0.7.4"
chrono = { version = "0.4.38", features = ["serde"] }
indoc = "2.0.5"
nanoid = "0.4"
sha2 = "0.10"
base64 = "0.21"
url = "2.5"
axum = "0.7"
webbrowser = "0.8"
dotenv = "0.15"
lazy_static = "1.5"
tracing = "0.1"
tracing-subscriber = "0.3"
wiremock = "0.6.0"
keyring = { version = "3.6.1", features = ["apple-native", "windows-native", "sync-secret-service"] }
ctor = "0.2.7"
paste = "1.0"
serde_yaml = "0.9.34"
once_cell = "1.20.2"
dirs = "6.0.0"
[dev-dependencies]
criterion = "0.5"
tempfile = "3.15.0"
serial_test = "3.2.0"
[[example]]
name = "agent"
path = "examples/agent.rs"
[[example]]
name = "databricks_oauth"
path = "examples/databricks_oauth.rs"
[[bench]]
name = "tokenization_benchmark"
harness = false

View File

@@ -0,0 +1,20 @@
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use goose::token_counter::TokenCounter;
fn benchmark_tokenization(c: &mut Criterion) {
let lengths = [1_000, 5_000, 10_000, 50_000, 100_000, 124_000, 200_000];
let tokenizer_names = ["Xenova--gpt-4o", "Xenova--claude-tokenizer"];
for tokenizer_name in tokenizer_names {
let counter = TokenCounter::new(tokenizer_name);
for &length in &lengths {
let text = "hello ".repeat(length);
c.bench_function(&format!("{}_{}_tokens", tokenizer_name, length), |b| {
b.iter(|| counter.count_tokens(black_box(&text)))
});
}
}
}
criterion_group!(benches, benchmark_tokenization);
criterion_main!(benches);

53
crates/goose/build.rs Normal file
View File

@@ -0,0 +1,53 @@
use std::error::Error;
use std::fs;
use std::path::Path;
const BASE_DIR: &str = "../../tokenizer_files";
const TOKENIZERS: &[&str] = &["Xenova/gpt-4o", "Xenova/claude-tokenizer"];
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
// Create base directory
fs::create_dir_all(BASE_DIR)?;
println!("cargo:rerun-if-changed=build.rs");
println!("cargo:rerun-if-changed={}", BASE_DIR);
for tokenizer_name in TOKENIZERS {
download_tokenizer(tokenizer_name).await?;
}
Ok(())
}
async fn download_tokenizer(repo_id: &str) -> Result<(), Box<dyn Error>> {
let dir_name = repo_id.replace('/', "--");
let download_dir = format!("{}/{}", BASE_DIR, dir_name);
let file_url = format!(
"https://huggingface.co/{}/resolve/main/tokenizer.json",
repo_id
);
let file_path = format!("{}/tokenizer.json", download_dir);
// Create directory if it doesn't exist
fs::create_dir_all(&download_dir)?;
// Check if file already exists
if Path::new(&file_path).exists() {
println!("Tokenizer for {} already exists, skipping...", repo_id);
return Ok(());
}
println!("Downloading tokenizer for {}...", repo_id);
// Download the file
let response = reqwest::get(&file_url).await?;
if !response.status().is_success() {
return Err(format!("Failed to download tokenizer for {}", repo_id).into());
}
let content = response.bytes().await?;
fs::write(&file_path, content)?;
println!("Downloaded {} to {}", repo_id, file_path);
Ok(())
}

View File

@@ -0,0 +1,36 @@
use dotenv::dotenv;
use futures::StreamExt;
use goose::agents::{AgentFactory, ExtensionConfig};
use goose::message::Message;
use goose::providers::databricks::DatabricksProvider;
#[tokio::main]
async fn main() {
// Setup a model provider from env vars
let _ = dotenv();
let provider = Box::new(DatabricksProvider::default());
// Setup an agent with the developer extension
let mut agent = AgentFactory::create("reference", provider).expect("default should exist");
let config = ExtensionConfig::stdio("developer", "./target/debug/developer");
agent.add_extension(config).await.unwrap();
println!("Extensions:");
for extension in agent.list_extensions().await {
println!(" {}", extension);
}
let messages = vec![Message::user()
.with_text("can you summarize the readme.md in this dir using just a haiku?")];
let mut stream = agent.reply(&messages).await.unwrap();
while let Some(message) = stream.next().await {
println!(
"{}",
serde_json::to_string_pretty(&message.unwrap()).unwrap()
);
println!("\n");
}
}

View File

@@ -0,0 +1,40 @@
use anyhow::Result;
use dotenv::dotenv;
use goose::{
message::Message,
providers::{base::Provider, databricks::DatabricksProvider},
};
#[tokio::main]
async fn main() -> Result<()> {
// Load environment variables from .env file
dotenv().ok();
// Clear any token to force OAuth
std::env::remove_var("DATABRICKS_TOKEN");
// Create the provider
let provider = DatabricksProvider::default();
// Create a simple message
let message = Message::user().with_text("Tell me a short joke about programming.");
// Get a response
let (response, usage) = provider
.complete("You are a helpful assistant.", &[message], &[])
.await?;
// Print the response and usage statistics
println!("\nResponse from AI:");
println!("---------------");
for content in response.content {
dbg!(content);
}
println!("\nToken Usage:");
println!("------------");
println!("Input tokens: {:?}", usage.usage.input_tokens);
println!("Output tokens: {:?}", usage.usage.output_tokens);
println!("Total tokens: {:?}", usage.usage.total_tokens);
Ok(())
}

View File

@@ -0,0 +1,79 @@
use anyhow::Result;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use dotenv::dotenv;
use goose::{
message::Message,
providers::{databricks::DatabricksProvider, openai::OpenAiProvider},
};
use mcp_core::{
content::Content,
tool::{Tool, ToolCall},
};
use serde_json::json;
use std::fs;
#[tokio::main]
async fn main() -> Result<()> {
// Load environment variables from .env file
dotenv().ok();
// Create providers
let providers: Vec<Box<dyn goose::providers::base::Provider + Send + Sync>> = vec![
Box::new(DatabricksProvider::default()),
Box::new(OpenAiProvider::default()),
];
for provider in providers {
// Read and encode test image
let image_data = fs::read("crates/goose/examples/test_assets/test_image.png")?;
let base64_image = BASE64.encode(image_data);
// Create a message sequence that includes a tool response with both text and image
let messages = vec![
Message::user().with_text("Read the image at ./test_image.png please"),
Message::assistant().with_tool_request(
"000",
Ok(ToolCall::new(
"view_image",
json!({"path": "./test_image.png"}),
)),
),
Message::user()
.with_tool_response("000", Ok(vec![Content::image(base64_image, "image/png")])),
];
// Get a response from the model about the image
let input_schema = json!({
"type": "object",
"required": ["path"],
"properties": {
"path": {
"type": "string",
"default": null,
"description": "The path to the image"
},
}
});
let (response, usage) = provider
.complete(
"You are a helpful assistant. Please describe any text you see in the image.",
&messages,
&[Tool::new("view_image", "View an image", input_schema)],
)
.await?;
// Print the response and usage statistics
println!("\nResponse from AI:");
println!("---------------");
for content in response.content {
println!("{:?}", content);
}
println!("\nToken Usage:");
println!("------------");
println!("Input tokens: {:?}", usage.usage.input_tokens);
println!("Output tokens: {:?}", usage.usage.output_tokens);
println!("Total tokens: {:?}", usage.usage.total_tokens);
}
Ok(())
}

Binary file not shown.

After

Width:  |  Height:  |  Size: 4.2 KiB

View File

@@ -0,0 +1,31 @@
use anyhow::Result;
use async_trait::async_trait;
use futures::stream::BoxStream;
use serde_json::Value;
use super::extension::{ExtensionConfig, ExtensionResult};
use crate::message::Message;
use crate::providers::base::ProviderUsage;
/// Core trait defining the behavior of an Agent
#[async_trait]
pub trait Agent: Send + Sync {
/// Create a stream that yields each message as it's generated by the agent
async fn reply(&self, messages: &[Message]) -> Result<BoxStream<'_, Result<Message>>>;
/// Add a new MCP client to the agent
async fn add_extension(&mut self, config: ExtensionConfig) -> ExtensionResult<()>;
/// Remove an extension by name
async fn remove_extension(&mut self, name: &str);
/// List all extensions
// TODO this needs to also include status so we can tell if extensions are dropped
async fn list_extensions(&self) -> Vec<String>;
/// Pass through a JSON-RPC request to a specific extension
async fn passthrough(&self, extension: &str, request: Value) -> ExtensionResult<Value>;
/// Get the total usage of the agent
async fn usage(&self) -> Vec<ProviderUsage>;
}

View File

@@ -0,0 +1,734 @@
use chrono::{DateTime, TimeZone, Utc};
use futures::stream::{FuturesUnordered, StreamExt};
use mcp_client::McpService;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::LazyLock;
use std::time::Duration;
use tokio::sync::Mutex;
use tracing::{debug, instrument};
use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult};
use crate::prompt_template::load_prompt_file;
use crate::providers::base::{Provider, ProviderUsage};
use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
use mcp_client::transport::{SseTransport, StdioTransport, Transport};
use mcp_core::{Content, Tool, ToolCall, ToolError, ToolResult};
use serde_json::Value;
// By default, we set it to Jan 1, 2020 if the resource does not have a timestamp
// This is to ensure that the resource is considered less important than resources with a more recent timestamp
static DEFAULT_TIMESTAMP: LazyLock<DateTime<Utc>> =
LazyLock::new(|| Utc.with_ymd_and_hms(2020, 1, 1, 0, 0, 0).unwrap());
type McpClientBox = Arc<Mutex<Box<dyn McpClientTrait>>>;
/// Manages MCP clients and their interactions
pub struct Capabilities {
clients: HashMap<String, McpClientBox>,
instructions: HashMap<String, String>,
resource_capable_extensions: HashSet<String>,
provider: Box<dyn Provider>,
provider_usage: Mutex<Vec<ProviderUsage>>,
}
/// A flattened representation of a resource used by the agent to prepare inference
#[derive(Debug, Clone)]
pub struct ResourceItem {
pub client_name: String, // The name of the client that owns the resource
pub uri: String, // The URI of the resource
pub name: String, // The name of the resource
pub content: String, // The content of the resource
pub timestamp: DateTime<Utc>, // The timestamp of the resource
pub priority: f32, // The priority of the resource
pub token_count: Option<u32>, // The token count of the resource (filled in by the agent)
}
impl ResourceItem {
pub fn new(
client_name: String,
uri: String,
name: String,
content: String,
timestamp: DateTime<Utc>,
priority: f32,
) -> Self {
Self {
client_name,
uri,
name,
content,
timestamp,
priority,
token_count: None,
}
}
}
/// Sanitizes a string by replacing invalid characters with underscores.
/// Valid characters match [a-zA-Z0-9_-]
fn normalize(input: String) -> String {
let mut result = String::with_capacity(input.len());
for c in input.chars() {
result.push(match c {
c if c.is_ascii_alphanumeric() || c == '_' || c == '-' => c,
c if c.is_whitespace() => continue, // effectively "strip" whitespace
_ => '_', // Replace any other non-ASCII character with '_'
});
}
result.to_lowercase()
}
impl Capabilities {
/// Create a new Capabilities with the specified provider
pub fn new(provider: Box<dyn Provider>) -> Self {
Self {
clients: HashMap::new(),
instructions: HashMap::new(),
resource_capable_extensions: HashSet::new(),
provider,
provider_usage: Mutex::new(Vec::new()),
}
}
pub fn supports_resources(&self) -> bool {
!self.resource_capable_extensions.is_empty()
}
/// Add a new MCP extension based on the provided client type
// TODO IMPORTANT need to ensure this times out if the extension command is broken!
pub async fn add_extension(&mut self, config: ExtensionConfig) -> ExtensionResult<()> {
let mut client: Box<dyn McpClientTrait> = match &config {
ExtensionConfig::Sse { uri, envs, .. } => {
let transport = SseTransport::new(uri, envs.get_env());
let handle = transport.start().await?;
let service = McpService::with_timeout(handle, Duration::from_secs(300));
Box::new(McpClient::new(service))
}
ExtensionConfig::Stdio {
cmd, args, envs, ..
} => {
let transport = StdioTransport::new(cmd, args.to_vec(), envs.get_env());
let handle = transport.start().await?;
let service = McpService::with_timeout(handle, Duration::from_secs(300));
Box::new(McpClient::new(service))
}
ExtensionConfig::Builtin { name } => {
// For builtin extensions, we run the current executable with mcp and extension name
let cmd = std::env::current_exe()
.expect("should find the current executable")
.to_str()
.expect("should resolve executable to string path")
.to_string();
let transport = StdioTransport::new(
&cmd,
vec!["mcp".to_string(), name.clone()],
HashMap::new(),
);
let handle = transport.start().await?;
let service = McpService::with_timeout(handle, Duration::from_secs(300));
Box::new(McpClient::new(service))
}
};
// Initialize the client with default capabilities
let info = ClientInfo {
name: "goose".to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
};
let capabilities = ClientCapabilities::default();
let init_result = client
.initialize(info, capabilities)
.await
.map_err(|e| ExtensionError::Initialization(config.clone(), e))?;
let sanitized_name = normalize(config.name().to_string());
// Store instructions if provided
if let Some(instructions) = init_result.instructions {
self.instructions
.insert(sanitized_name.clone(), instructions);
}
// if the server is capable if resources we track it
if init_result.capabilities.resources.is_some() {
self.resource_capable_extensions
.insert(sanitized_name.clone());
}
// Store the client using the provided name
self.clients
.insert(sanitized_name.clone(), Arc::new(Mutex::new(client)));
Ok(())
}
/// Get a reference to the provider
pub fn provider(&self) -> &dyn Provider {
&*self.provider
}
/// Record provider usage
// TODO consider moving this off to the provider or as a form of logging
pub async fn record_usage(&self, usage: ProviderUsage) {
self.provider_usage.lock().await.push(usage);
}
/// Get aggregated usage statistics
pub async fn remove_extension(&mut self, name: &str) -> ExtensionResult<()> {
let sanitized_name = normalize(name.to_string());
self.clients.remove(&sanitized_name);
self.instructions.remove(&sanitized_name);
self.resource_capable_extensions.remove(&sanitized_name);
Ok(())
}
pub async fn list_extensions(&self) -> ExtensionResult<Vec<String>> {
Ok(self.clients.keys().cloned().collect())
}
pub async fn get_usage(&self) -> Vec<ProviderUsage> {
let provider_usage = self.provider_usage.lock().await.clone();
let mut usage_map: HashMap<String, ProviderUsage> = HashMap::new();
provider_usage.iter().for_each(|usage| {
usage_map
.entry(usage.model.clone())
.and_modify(|e| {
e.usage.input_tokens = Some(
e.usage.input_tokens.unwrap_or(0) + usage.usage.input_tokens.unwrap_or(0),
);
e.usage.output_tokens = Some(
e.usage.output_tokens.unwrap_or(0) + usage.usage.output_tokens.unwrap_or(0),
);
e.usage.total_tokens = Some(
e.usage.total_tokens.unwrap_or(0) + usage.usage.total_tokens.unwrap_or(0),
);
})
.or_insert_with(|| usage.clone());
});
usage_map.into_values().collect()
}
/// Get all tools from all clients with proper prefixing
pub async fn get_prefixed_tools(&mut self) -> ExtensionResult<Vec<Tool>> {
let mut tools = Vec::new();
for (name, client) in &self.clients {
let client_guard = client.lock().await;
let mut client_tools = client_guard.list_tools(None).await?;
loop {
for tool in client_tools.tools {
tools.push(Tool::new(
format!("{}__{}", name, tool.name),
&tool.description,
tool.input_schema,
));
}
// exit loop when there are no more pages
if client_tools.next_cursor.is_none() {
break;
}
client_tools = client_guard.list_tools(client_tools.next_cursor).await?;
}
}
Ok(tools)
}
/// Get client resources and their contents
pub async fn get_resources(&self) -> ExtensionResult<Vec<ResourceItem>> {
let mut result: Vec<ResourceItem> = Vec::new();
for (name, client) in &self.clients {
let client_guard = client.lock().await;
let resources = client_guard.list_resources(None).await?;
for resource in resources.resources {
// Skip reading the resource if it's not marked active
// This avoids blowing up the context with inactive resources
if !resource.is_active() {
continue;
}
if let Ok(contents) = client_guard.read_resource(&resource.uri).await {
for content in contents.contents {
let (uri, content_str) = match content {
mcp_core::resource::ResourceContents::TextResourceContents {
uri,
text,
..
} => (uri, text),
mcp_core::resource::ResourceContents::BlobResourceContents {
uri,
blob,
..
} => (uri, blob),
};
result.push(ResourceItem::new(
name.clone(),
uri,
resource.name.clone(),
content_str,
resource.timestamp().unwrap_or(*DEFAULT_TIMESTAMP),
resource.priority().unwrap_or(0.0),
));
}
}
}
}
Ok(result)
}
/// Get the extension prompt including client instructions
pub async fn get_system_prompt(&self) -> String {
let mut context: HashMap<&str, Vec<ExtensionInfo>> = HashMap::new();
let extensions_info: Vec<ExtensionInfo> = self
.clients
.keys()
.map(|name| {
let instructions = self.instructions.get(name).cloned().unwrap_or_default();
let has_resources = self.resource_capable_extensions.contains(name);
ExtensionInfo::new(name, &instructions, has_resources)
})
.collect();
context.insert("extensions", extensions_info);
load_prompt_file("system.md", &context).expect("Prompt should render")
}
/// Find and return a reference to the appropriate client for a tool call
fn get_client_for_tool(&self, prefixed_name: &str) -> Option<(&str, McpClientBox)> {
self.clients
.iter()
.find(|(key, _)| prefixed_name.starts_with(*key))
.map(|(name, client)| (name.as_str(), Arc::clone(client)))
}
// Function that gets executed for read_resource tool
async fn read_resource(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let uri = params
.get("uri")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'uri' parameter".to_string()))?;
let extension_name = params.get("extension_name").and_then(|v| v.as_str());
// If extension name is provided, we can just look it up
if extension_name.is_some() {
let result = self
.read_resource_from_extension(uri, extension_name.unwrap())
.await?;
return Ok(result);
}
// If extension name is not provided, we need to search for the resource across all extensions
// Loop through each extension and try to read the resource, don't raise an error if the resource is not found
// TODO: do we want to find if a provided uri is in multiple extensions?
// currently it will return the first match and skip any others
for extension_name in self.resource_capable_extensions.iter() {
let result = self.read_resource_from_extension(uri, extension_name).await;
match result {
Ok(result) => return Ok(result),
Err(_) => continue,
}
}
// None of the extensions had the resource so we raise an error
let available_extensions = self
.clients
.keys()
.map(|s| s.as_str())
.collect::<Vec<&str>>()
.join(", ");
let error_msg = format!(
"Resource with uri '{}' not found. Here are the available extensions: {}",
uri, available_extensions
);
Err(ToolError::InvalidParameters(error_msg))
}
async fn read_resource_from_extension(
&self,
uri: &str,
extension_name: &str,
) -> Result<Vec<Content>, ToolError> {
let available_extensions = self
.clients
.keys()
.map(|s| s.as_str())
.collect::<Vec<&str>>()
.join(", ");
let error_msg = format!(
"Extension '{}' not found. Here are the available extensions: {}",
extension_name, available_extensions
);
let client = self
.clients
.get(extension_name)
.ok_or(ToolError::InvalidParameters(error_msg))?;
let client_guard = client.lock().await;
let read_result = client_guard.read_resource(uri).await.map_err(|_| {
ToolError::ExecutionError(format!("Could not read resource with uri: {}", uri))
})?;
let mut result = Vec::new();
for content in read_result.contents {
// Only reading the text resource content; skipping the blob content cause it's too long
if let mcp_core::resource::ResourceContents::TextResourceContents { text, .. } = content
{
let content_str = format!("{}\n\n{}", uri, text);
result.push(Content::text(content_str));
}
}
Ok(result)
}
async fn list_resources_from_extension(
&self,
extension_name: &str,
) -> Result<Vec<Content>, ToolError> {
let client = self.clients.get(extension_name).ok_or_else(|| {
ToolError::InvalidParameters(format!("Extension {} is not valid", extension_name))
})?;
let client_guard = client.lock().await;
client_guard
.list_resources(None)
.await
.map_err(|e| {
ToolError::ExecutionError(format!(
"Unable to list resources for {}, {:?}",
extension_name, e
))
})
.map(|lr| {
let resource_list = lr
.resources
.into_iter()
.map(|r| format!("{} - {}, uri: ({})", extension_name, r.name, r.uri))
.collect::<Vec<String>>()
.join("\n");
vec![Content::text(resource_list)]
})
}
async fn list_resources(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let extension = params.get("extension").and_then(|v| v.as_str());
match extension {
Some(extension_name) => {
// Handle single extension case
self.list_resources_from_extension(extension_name).await
}
None => {
// Handle all extensions case using FuturesUnordered
let mut futures = FuturesUnordered::new();
// Create futures for each resource_capable_extension
for extension_name in &self.resource_capable_extensions {
futures.push(async move {
self.list_resources_from_extension(extension_name).await
});
}
let mut all_resources = Vec::new();
let mut errors = Vec::new();
// Process results as they complete
while let Some(result) = futures.next().await {
match result {
Ok(content) => {
all_resources.extend(content);
}
Err(tool_error) => {
errors.push(tool_error);
}
}
}
// Log any errors that occurred
if !errors.is_empty() {
tracing::error!(
errors = ?errors
.into_iter()
.map(|e| format!("{:?}", e))
.collect::<Vec<_>>(),
"errors from listing resources"
);
}
Ok(all_resources)
}
}
}
/// Dispatch a single tool call to the appropriate client
#[instrument(skip(self, tool_call), fields(input, output))]
pub async fn dispatch_tool_call(&self, tool_call: ToolCall) -> ToolResult<Vec<Content>> {
let result = if tool_call.name == "platform__read_resource" {
// Check if the tool is read_resource and handle it separately
self.read_resource(tool_call.arguments.clone()).await
} else if tool_call.name == "platform__list_resources" {
self.list_resources(tool_call.arguments.clone()).await
} else {
// Else, dispatch tool call based on the prefix naming convention
let (client_name, client) = self
.get_client_for_tool(&tool_call.name)
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
// rsplit returns the iterator in reverse, tool_name is then at 0
let tool_name = tool_call
.name
.strip_prefix(client_name)
.and_then(|s| s.strip_prefix("__"))
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
let client_guard = client.lock().await;
client_guard
.call_tool(tool_name, tool_call.clone().arguments)
.await
.map(|result| result.content)
.map_err(|e| ToolError::ExecutionError(e.to_string()))
};
debug!(
"input" = serde_json::to_string(&tool_call).unwrap(),
"output" = serde_json::to_string(&result).unwrap(),
);
result
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::Message;
use crate::model::ModelConfig;
use crate::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
use crate::providers::errors::ProviderError;
use mcp_client::client::Error;
use mcp_client::client::McpClientTrait;
use mcp_core::protocol::{
CallToolResult, InitializeResult, ListResourcesResult, ListToolsResult, ReadResourceResult,
};
use serde_json::json;
// Mock Provider implementation for testing
#[derive(Clone)]
struct MockProvider {
model_config: ModelConfig,
}
#[async_trait::async_trait]
impl Provider for MockProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete(
&self,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
Ok((
Message::assistant().with_text("Mock response"),
ProviderUsage::new("mock".to_string(), Usage::default()),
))
}
}
struct MockClient {}
#[async_trait::async_trait]
impl McpClientTrait for MockClient {
async fn initialize(
&mut self,
_info: ClientInfo,
_capabilities: ClientCapabilities,
) -> Result<InitializeResult, Error> {
Err(Error::NotInitialized)
}
async fn list_resources(
&self,
_next_cursor: Option<String>,
) -> Result<ListResourcesResult, Error> {
Err(Error::NotInitialized)
}
async fn read_resource(&self, _uri: &str) -> Result<ReadResourceResult, Error> {
Err(Error::NotInitialized)
}
async fn list_tools(&self, _next_cursor: Option<String>) -> Result<ListToolsResult, Error> {
Err(Error::NotInitialized)
}
async fn call_tool(&self, name: &str, _arguments: Value) -> Result<CallToolResult, Error> {
match name {
"tool" | "test__tool" => Ok(CallToolResult {
content: vec![],
is_error: None,
}),
_ => Err(Error::NotInitialized),
}
}
}
#[test]
fn test_get_client_for_tool() {
let mock_model_config =
ModelConfig::new("test-model".to_string()).with_context_limit(200_000.into());
let mut capabilities = Capabilities::new(Box::new(MockProvider {
model_config: mock_model_config,
}));
// Add some mock clients
capabilities.clients.insert(
normalize("test_client".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
capabilities.clients.insert(
normalize("__client".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
capabilities.clients.insert(
normalize("__cli__ent__".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
capabilities.clients.insert(
normalize("client 🚀".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
// Test basic case
assert!(capabilities
.get_client_for_tool("test_client__tool")
.is_some());
// Test leading underscores
assert!(capabilities.get_client_for_tool("__client__tool").is_some());
// Test multiple underscores in client name, and ending with __
assert!(capabilities
.get_client_for_tool("__cli__ent____tool")
.is_some());
// Test unicode in tool name, "client 🚀" should become "client_"
assert!(capabilities.get_client_for_tool("client___tool").is_some());
}
#[tokio::test]
async fn test_dispatch_tool_call() {
// test that dispatch_tool_call parses out the sanitized name correctly, and extracts
// tool_names
let mock_model_config =
ModelConfig::new("test-model".to_string()).with_context_limit(200_000.into());
let mut capabilities = Capabilities::new(Box::new(MockProvider {
model_config: mock_model_config,
}));
// Add some mock clients
capabilities.clients.insert(
normalize("test_client".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
capabilities.clients.insert(
normalize("__cli__ent__".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
capabilities.clients.insert(
normalize("client 🚀".to_string()),
Arc::new(Mutex::new(Box::new(MockClient {}))),
);
// verify a normal tool call
let tool_call = ToolCall {
name: "test_client__tool".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(tool_call).await;
assert!(result.is_ok());
let tool_call = ToolCall {
name: "test_client__test__tool".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(tool_call).await;
assert!(result.is_ok());
// verify a multiple underscores dispatch
let tool_call = ToolCall {
name: "__cli__ent____tool".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(tool_call).await;
assert!(result.is_ok());
// Test unicode in tool name, "client 🚀" should become "client_"
let tool_call = ToolCall {
name: "client___tool".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(tool_call).await;
assert!(result.is_ok());
let tool_call = ToolCall {
name: "client___test__tool".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(tool_call).await;
assert!(result.is_ok());
// this should error out, specifically for an ToolError::ExecutionError
let invalid_tool_call = ToolCall {
name: "client___tools".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(invalid_tool_call).await;
assert!(matches!(
result.err().unwrap(),
ToolError::ExecutionError(_)
));
// this should error out, specifically with an ToolError::NotFound
// this client doesn't exist
let invalid_tool_call = ToolCall {
name: "_client__tools".to_string(),
arguments: json!({}),
};
let result = capabilities.dispatch_tool_call(invalid_tool_call).await;
assert!(matches!(result.err().unwrap(), ToolError::NotFound(_)));
}
}

View File

@@ -0,0 +1,158 @@
use std::collections::HashMap;
use mcp_client::client::Error as ClientError;
use serde::{Deserialize, Serialize};
use thiserror::Error;
/// Errors from Extension operation
#[derive(Error, Debug)]
pub enum ExtensionError {
#[error("Failed to start the MCP server from configuration `{0}` `{1}`")]
Initialization(ExtensionConfig, ClientError),
#[error("Failed a client call to an MCP server: {0}")]
Client(#[from] ClientError),
#[error("User Message exceeded context-limit. History could not be truncated to accomodate.")]
ContextLimit,
#[error("Transport error: {0}")]
Transport(#[from] mcp_client::transport::Error),
}
pub type ExtensionResult<T> = Result<T, ExtensionError>;
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
pub struct Envs {
/// A map of environment variables to set, e.g. API_KEY -> some_secret, HOST -> host
#[serde(default)]
#[serde(flatten)]
map: HashMap<String, String>,
}
impl Envs {
pub fn new(map: HashMap<String, String>) -> Self {
Self { map }
}
pub fn get_env(&self) -> HashMap<String, String> {
self.map
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect()
}
}
/// Represents the different types of MCP extensions that can be added to the manager
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type")]
pub enum ExtensionConfig {
/// Server-sent events client with a URI endpoint
#[serde(rename = "sse")]
Sse {
/// The name used to identify this extension
name: String,
uri: String,
#[serde(default)]
envs: Envs,
},
/// Standard I/O client with command and arguments
#[serde(rename = "stdio")]
Stdio {
/// The name used to identify this extension
name: String,
cmd: String,
args: Vec<String>,
#[serde(default)]
envs: Envs,
},
/// Built-in extension that is part of the goose binary
#[serde(rename = "builtin")]
Builtin {
/// The name used to identify this extension
name: String,
},
}
impl Default for ExtensionConfig {
fn default() -> Self {
Self::Builtin {
name: String::from("default"),
}
}
}
impl ExtensionConfig {
pub fn sse<S: Into<String>>(name: S, uri: S) -> Self {
Self::Sse {
name: name.into(),
uri: uri.into(),
envs: Envs::default(),
}
}
pub fn stdio<S: Into<String>>(name: S, cmd: S) -> Self {
Self::Stdio {
name: name.into(),
cmd: cmd.into(),
args: vec![],
envs: Envs::default(),
}
}
pub fn with_args<I, S>(self, args: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
match self {
Self::Stdio {
name, cmd, envs, ..
} => Self::Stdio {
name,
cmd,
envs,
args: args.into_iter().map(Into::into).collect(),
},
other => other,
}
}
/// Get the extension name regardless of variant
pub fn name(&self) -> &str {
match self {
Self::Sse { name, .. } => name,
Self::Stdio { name, .. } => name,
Self::Builtin { name } => name,
}
}
}
impl std::fmt::Display for ExtensionConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ExtensionConfig::Sse { name, uri, .. } => write!(f, "SSE({}: {})", name, uri),
ExtensionConfig::Stdio {
name, cmd, args, ..
} => {
write!(f, "Stdio({}: {} {})", name, cmd, args.join(" "))
}
ExtensionConfig::Builtin { name } => write!(f, "Builtin({})", name),
}
}
}
/// Information about the extension used for building prompts
#[derive(Clone, Debug, Serialize)]
pub struct ExtensionInfo {
name: String,
instructions: String,
has_resources: bool,
}
impl ExtensionInfo {
pub fn new(name: &str, instructions: &str, has_resources: bool) -> Self {
Self {
name: name.to_string(),
instructions: instructions.to_string(),
has_resources,
}
}
}

View File

@@ -0,0 +1,69 @@
use std::collections::HashMap;
use std::sync::{OnceLock, RwLock};
pub use super::Agent;
use crate::providers::base::Provider;
type AgentConstructor = Box<dyn Fn(Box<dyn Provider>) -> Box<dyn Agent> + Send + Sync>;
// Use std::sync::RwLock for interior mutability
static AGENT_REGISTRY: OnceLock<RwLock<HashMap<&'static str, AgentConstructor>>> = OnceLock::new();
/// Initialize the registry if it hasn't been initialized
fn registry() -> &'static RwLock<HashMap<&'static str, AgentConstructor>> {
AGENT_REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
}
/// Register a new agent version
pub fn register_agent(
version: &'static str,
constructor: impl Fn(Box<dyn Provider>) -> Box<dyn Agent> + Send + Sync + 'static,
) {
let registry = registry();
if let Ok(mut map) = registry.write() {
map.insert(version, Box::new(constructor));
}
}
pub struct AgentFactory;
impl AgentFactory {
/// Create a new agent instance of the specified version
pub fn create(version: &str, provider: Box<dyn Provider>) -> Option<Box<dyn Agent>> {
let registry = registry();
let map = registry
.read()
.expect("should be able to read the registry");
let constructor = map.get(version)?;
Some(constructor(provider))
}
/// Get a list of all available agent versions
pub fn available_versions() -> Vec<&'static str> {
registry()
.read()
.map(|map| map.keys().copied().collect())
.unwrap_or_default()
}
/// Get the default version name
pub fn default_version() -> &'static str {
"truncate"
}
}
/// Macro to help with agent registration
#[macro_export]
macro_rules! register_agent {
($version:expr, $agent_type:ty) => {
paste::paste! {
#[ctor::ctor]
#[allow(non_snake_case)]
fn [<__register_agent_ $version>]() {
$crate::agents::factory::register_agent($version, |provider| {
Box::new(<$agent_type>::new(provider))
});
}
}
};
}

View File

@@ -0,0 +1,11 @@
mod agent;
mod capabilities;
pub mod extension;
mod factory;
mod reference;
mod truncate;
pub use agent::Agent;
pub use capabilities::Capabilities;
pub use extension::ExtensionConfig;
pub use factory::{register_agent, AgentFactory};

View File

@@ -0,0 +1,189 @@
/// A simplified agent implementation used as a reference
/// It makes no attempt to handle context limits, and cannot read resources
use async_trait::async_trait;
use futures::stream::BoxStream;
use tokio::sync::Mutex;
use tracing::{debug, instrument};
use super::Agent;
use crate::agents::capabilities::Capabilities;
use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::message::{Message, ToolRequest};
use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
use crate::register_agent;
use crate::token_counter::TokenCounter;
use indoc::indoc;
use mcp_core::tool::Tool;
use serde_json::{json, Value};
/// Reference implementation of an Agent
pub struct ReferenceAgent {
capabilities: Mutex<Capabilities>,
_token_counter: TokenCounter,
}
impl ReferenceAgent {
pub fn new(provider: Box<dyn Provider>) -> Self {
let token_counter = TokenCounter::new(provider.get_model_config().tokenizer_name());
Self {
capabilities: Mutex::new(Capabilities::new(provider)),
_token_counter: token_counter,
}
}
}
#[async_trait]
impl Agent for ReferenceAgent {
async fn add_extension(&mut self, extension: ExtensionConfig) -> ExtensionResult<()> {
let mut capabilities = self.capabilities.lock().await;
capabilities.add_extension(extension).await
}
async fn remove_extension(&mut self, name: &str) {
let mut capabilities = self.capabilities.lock().await;
capabilities
.remove_extension(name)
.await
.expect("Failed to remove extension");
}
async fn list_extensions(&self) -> Vec<String> {
let capabilities = self.capabilities.lock().await;
capabilities
.list_extensions()
.await
.expect("Failed to list extensions")
}
async fn passthrough(&self, _extension: &str, _request: Value) -> ExtensionResult<Value> {
// TODO implement
Ok(Value::Null)
}
#[instrument(skip(self, messages), fields(user_message))]
async fn reply(
&self,
messages: &[Message],
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
let mut capabilities = self.capabilities.lock().await;
let mut tools = capabilities.get_prefixed_tools().await?;
// we add in the read_resource tool by default
// TODO: make sure there is no collision with another extension's tool name
let read_resource_tool = Tool::new(
"platform__read_resource".to_string(),
indoc! {r#"
Read a resource from an extension.
Resources allow extensions to share data that provide context to LLMs, such as
files, database schemas, or application-specific information. This tool searches for the
resource URI in the provided extension, and reads in the resource content. If no extension
is provided, the tool will search all extensions for the resource.
"#}.to_string(),
json!({
"type": "object",
"required": ["uri"],
"properties": {
"uri": {"type": "string", "description": "Resource URI"},
"extension_name": {"type": "string", "description": "Optional extension name"}
}
}),
);
let list_resources_tool = Tool::new(
"platform__list_resources".to_string(),
indoc! {r#"
List resources from an extension(s).
Resources allow extensions to share data that provide context to LLMs, such as
files, database schemas, or application-specific information. This tool lists resources
in the provided extension, and returns a list for the user to browse. If no extension
is provided, the tool will search all extensions for the resource.
"#}.to_string(),
json!({
"type": "object",
"properties": {
"extension_name": {"type": "string", "description": "Optional extension name"}
}
}),
);
if capabilities.supports_resources() {
tools.push(read_resource_tool);
tools.push(list_resources_tool);
}
let system_prompt = capabilities.get_system_prompt().await;
// Set the user_message field in the span instead of creating a new event
if let Some(content) = messages
.last()
.and_then(|msg| msg.content.first())
.and_then(|c| c.as_text())
{
debug!("user_message" = &content);
}
Ok(Box::pin(async_stream::try_stream! {
let _reply_guard = reply_span.enter();
loop {
// Get completion from provider
let (response, usage) = capabilities.provider().complete(
&system_prompt,
&messages,
&tools,
).await?;
capabilities.record_usage(usage).await;
// Yield the assistant's response
yield response.clone();
tokio::task::yield_now().await;
// First collect any tool requests
let tool_requests: Vec<&ToolRequest> = response.content
.iter()
.filter_map(|content| content.as_tool_request())
.collect();
if tool_requests.is_empty() {
break;
}
// Then dispatch each in parallel
let futures: Vec<_> = tool_requests
.iter()
.filter_map(|request| request.tool_call.clone().ok())
.map(|tool_call| capabilities.dispatch_tool_call(tool_call))
.collect();
// Process all the futures in parallel but wait until all are finished
let outputs = futures::future::join_all(futures).await;
// Create a message with the responses
let mut message_tool_response = Message::user();
// Now combine these into MessageContent::ToolResponse using the original ID
for (request, output) in tool_requests.iter().zip(outputs.into_iter()) {
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
output,
);
}
yield message_tool_response.clone();
messages.push(response);
messages.push(message_tool_response);
}
}))
}
async fn usage(&self) -> Vec<ProviderUsage> {
let capabilities = self.capabilities.lock().await;
capabilities.get_usage().await
}
}
register_agent!("reference", ReferenceAgent);

View File

@@ -0,0 +1,274 @@
/// A truncate agent that truncates the conversation history when it exceeds the model's context limit
/// It makes no attempt to handle context limits, and cannot read resources
use async_trait::async_trait;
use futures::stream::BoxStream;
use tokio::sync::Mutex;
use tracing::{debug, error, instrument, warn};
use super::Agent;
use crate::agents::capabilities::Capabilities;
use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::message::{Message, ToolRequest};
use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
use crate::providers::errors::ProviderError;
use crate::register_agent;
use crate::token_counter::TokenCounter;
use crate::truncate::{truncate_messages, OldestFirstTruncation};
use indoc::indoc;
use mcp_core::tool::Tool;
use serde_json::{json, Value};
const MAX_TRUNCATION_ATTEMPTS: usize = 3;
const ESTIMATE_FACTOR_DECAY: f32 = 0.9;
/// Truncate implementation of an Agent
pub struct TruncateAgent {
capabilities: Mutex<Capabilities>,
token_counter: TokenCounter,
}
impl TruncateAgent {
pub fn new(provider: Box<dyn Provider>) -> Self {
let token_counter = TokenCounter::new(provider.get_model_config().tokenizer_name());
Self {
capabilities: Mutex::new(Capabilities::new(provider)),
token_counter,
}
}
/// Truncates the messages to fit within the model's context window
/// Ensures the last message is a user message and removes tool call-response pairs
async fn truncate_messages(
&self,
messages: &mut Vec<Message>,
estimate_factor: f32,
) -> anyhow::Result<()> {
// Model's actual context limit
let context_limit = self
.capabilities
.lock()
.await
.provider()
.get_model_config()
.context_limit();
// Our conservative estimate of the **target** context limit
// Our token count is an estimate since model providers often don't provide the tokenizer (eg. Claude)
let context_limit = (context_limit as f32 * estimate_factor) as usize;
// Calculate current token count
let mut token_counts: Vec<usize> = messages
.iter()
.map(|msg| self.token_counter.count_tokens(&msg.as_concat_text()))
.collect();
let _ = truncate_messages(
messages,
&mut token_counts,
context_limit,
&OldestFirstTruncation,
);
Ok(())
}
}
#[async_trait]
impl Agent for TruncateAgent {
async fn add_extension(&mut self, extension: ExtensionConfig) -> ExtensionResult<()> {
let mut capabilities = self.capabilities.lock().await;
capabilities.add_extension(extension).await
}
async fn remove_extension(&mut self, name: &str) {
let mut capabilities = self.capabilities.lock().await;
capabilities
.remove_extension(name)
.await
.expect("Failed to remove extension");
}
async fn list_extensions(&self) -> Vec<String> {
let capabilities = self.capabilities.lock().await;
capabilities
.list_extensions()
.await
.expect("Failed to list extensions")
}
async fn passthrough(&self, _extension: &str, _request: Value) -> ExtensionResult<Value> {
// TODO implement
Ok(Value::Null)
}
#[instrument(skip(self, messages), fields(user_message))]
async fn reply(
&self,
messages: &[Message],
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
let mut capabilities = self.capabilities.lock().await;
let mut tools = capabilities.get_prefixed_tools().await?;
let mut truncation_attempt: usize = 0;
// we add in the read_resource tool by default
// TODO: make sure there is no collision with another extension's tool name
let read_resource_tool = Tool::new(
"platform__read_resource".to_string(),
indoc! {r#"
Read a resource from an extension.
Resources allow extensions to share data that provide context to LLMs, such as
files, database schemas, or application-specific information. This tool searches for the
resource URI in the provided extension, and reads in the resource content. If no extension
is provided, the tool will search all extensions for the resource.
"#}.to_string(),
json!({
"type": "object",
"required": ["uri"],
"properties": {
"uri": {"type": "string", "description": "Resource URI"},
"extension_name": {"type": "string", "description": "Optional extension name"}
}
}),
);
let list_resources_tool = Tool::new(
"platform__list_resources".to_string(),
indoc! {r#"
List resources from an extension(s).
Resources allow extensions to share data that provide context to LLMs, such as
files, database schemas, or application-specific information. This tool lists resources
in the provided extension, and returns a list for the user to browse. If no extension
is provided, the tool will search all extensions for the resource.
"#}.to_string(),
json!({
"type": "object",
"properties": {
"extension_name": {"type": "string", "description": "Optional extension name"}
}
}),
);
if capabilities.supports_resources() {
tools.push(read_resource_tool);
tools.push(list_resources_tool);
}
let system_prompt = capabilities.get_system_prompt().await;
// Set the user_message field in the span instead of creating a new event
if let Some(content) = messages
.last()
.and_then(|msg| msg.content.first())
.and_then(|c| c.as_text())
{
debug!("user_message" = &content);
}
Ok(Box::pin(async_stream::try_stream! {
let _reply_guard = reply_span.enter();
loop {
// Attempt to get completion from provider
match capabilities.provider().complete(
&system_prompt,
&messages,
&tools,
).await {
Ok((response, usage)) => {
capabilities.record_usage(usage).await;
// Reset truncation attempt
truncation_attempt = 0;
// Yield the assistant's response
yield response.clone();
tokio::task::yield_now().await;
// First collect any tool requests
let tool_requests: Vec<&ToolRequest> = response.content
.iter()
.filter_map(|content| content.as_tool_request())
.collect();
if tool_requests.is_empty() {
break;
}
// Then dispatch each in parallel
let futures: Vec<_> = tool_requests
.iter()
.filter_map(|request| request.tool_call.clone().ok())
.map(|tool_call| capabilities.dispatch_tool_call(tool_call))
.collect();
// Process all the futures in parallel but wait until all are finished
let outputs = futures::future::join_all(futures).await;
// Create a message with the responses
let mut message_tool_response = Message::user();
// Now combine these into MessageContent::ToolResponse using the original ID
for (request, output) in tool_requests.iter().zip(outputs.into_iter()) {
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
output,
);
}
yield message_tool_response.clone();
messages.push(response);
messages.push(message_tool_response);
},
Err(ProviderError::ContextLengthExceeded(_)) => {
if truncation_attempt >= MAX_TRUNCATION_ATTEMPTS {
// Create an error message & terminate the stream
// the previous message would have been a user message (e.g. before any tool calls, this is just after the input message.
// at the start of a loop after a tool call, it would be after a tool_use assistant followed by a tool_result user)
yield Message::assistant().with_text("Error: Context length exceeds limits even after multiple attempts to truncate.");
break;
}
truncation_attempt += 1;
warn!("Context length exceeded. Truncation Attempt: {}/{}.", truncation_attempt, MAX_TRUNCATION_ATTEMPTS);
// Decay the estimate factor as we make more truncation attempts
// Estimate factor decays like this over time: 0.9, 0.81, 0.729, ...
let estimate_factor: f32 = ESTIMATE_FACTOR_DECAY.powi(truncation_attempt as i32);
// release the lock before truncation to prevent deadlock
drop(capabilities);
self.truncate_messages(&mut messages, estimate_factor).await?;
// Re-acquire the lock
capabilities = self.capabilities.lock().await;
// Retry the loop after truncation
continue;
},
Err(e) => {
// Create an error message & terminate the stream
error!("Error: {}", e);
yield Message::assistant().with_text(format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error."));
break;
}
}
// Yield control back to the scheduler to prevent blocking
tokio::task::yield_now().await;
}
}))
}
async fn usage(&self) -> Vec<ProviderUsage> {
let capabilities = self.capabilities.lock().await;
capabilities.get_usage().await
}
}
register_agent!("truncate", TruncateAgent);

View File

@@ -0,0 +1,467 @@
use keyring::Entry;
use once_cell::sync::OnceCell;
use serde::Deserialize;
use serde_json::Value;
use std::collections::HashMap;
use std::env;
use std::path::{Path, PathBuf};
use thiserror::Error;
const KEYRING_SERVICE: &str = "goose";
const KEYRING_USERNAME: &str = "secrets";
#[cfg(test)]
const TEST_KEYRING_SERVICE: &str = "goose-test";
#[derive(Error, Debug)]
pub enum ConfigError {
#[error("Configuration value not found: {0}")]
NotFound(String),
#[error("Failed to deserialize value: {0}")]
DeserializeError(String),
#[error("Failed to read config file: {0}")]
FileError(#[from] std::io::Error),
#[error("Failed to create config directory: {0}")]
DirectoryError(String),
#[error("Failed to access keyring: {0}")]
KeyringError(String),
}
impl From<serde_json::Error> for ConfigError {
fn from(err: serde_json::Error) -> Self {
ConfigError::DeserializeError(err.to_string())
}
}
impl From<serde_yaml::Error> for ConfigError {
fn from(err: serde_yaml::Error) -> Self {
ConfigError::DeserializeError(err.to_string())
}
}
impl From<keyring::Error> for ConfigError {
fn from(err: keyring::Error) -> Self {
ConfigError::KeyringError(err.to_string())
}
}
/// Configuration management for Goose.
///
/// This module provides a flexible configuration system that supports:
/// - Dynamic configuration keys
/// - Multiple value types through serde deserialization
/// - Environment variable overrides
/// - YAML-based configuration file storage
/// - Hot reloading of configuration changes
/// - Secure secret storage in system keyring
///
/// Configuration values are loaded with the following precedence:
/// 1. Environment variables (exact key match)
/// 2. Configuration file (~/.config/goose/config.yaml by default)
///
/// Secrets are loaded with the following precedence:
/// 1. Environment variables (exact key match)
/// 2. System keyring
///
/// # Examples
///
/// ```no_run
/// use goose::config::Config;
/// use serde::Deserialize;
///
/// // Get a string value
/// let config = Config::global();
/// let api_key: String = config.get("OPENAI_API_KEY").unwrap();
///
/// // Get a complex type
/// #[derive(Deserialize)]
/// struct ServerConfig {
/// host: String,
/// port: u16,
/// }
///
/// let server_config: ServerConfig = config.get("server").unwrap();
/// ```
///
/// # Naming Convention
/// we recommend snake_case for keys, and will convert to UPPERCASE when
/// checking for environment overrides. e.g. openai_api_key will check for an
/// environment variable OPENAI_API_KEY
///
/// For Goose-specific configuration, consider prefixing with "goose_" to avoid conflicts.
pub struct Config {
config_path: PathBuf,
keyring_service: String,
}
// Global instance
static GLOBAL_CONFIG: OnceCell<Config> = OnceCell::new();
impl Default for Config {
fn default() -> Self {
let config_dir = dirs::home_dir()
.expect("goose requires a home dir")
.join(".config")
.join("goose");
std::fs::create_dir_all(&config_dir).expect("Failed to create config directory");
let config_path = config_dir.join("config.yaml");
Config {
config_path,
keyring_service: KEYRING_SERVICE.to_string(),
}
}
}
impl Config {
/// Get the global configuration instance.
///
/// This will initialize the configuration with the default path (~/.config/goose/config.yaml)
/// if it hasn't been initialized yet.
pub fn global() -> &'static Config {
GLOBAL_CONFIG.get_or_init(Config::default)
}
/// Create a new configuration instance with custom paths
///
/// This is primarily useful for testing or for applications that need
/// to manage multiple configuration files.
pub fn new<P: AsRef<Path>>(config_path: P, service: &str) -> Result<Self, ConfigError> {
Ok(Config {
config_path: config_path.as_ref().to_path_buf(),
keyring_service: service.to_string(),
})
}
/// Check if this config already exists
pub fn exists(&self) -> bool {
self.config_path.exists()
}
/// Check if this config already exists
pub fn clear(&self) -> Result<(), ConfigError> {
Ok(std::fs::remove_file(&self.config_path)?)
}
/// Get the path to the configuration file
pub fn path(&self) -> String {
self.config_path.to_string_lossy().to_string()
}
// Load current values from the config file
fn load_values(&self) -> Result<HashMap<String, Value>, ConfigError> {
if self.config_path.exists() {
let file_content = std::fs::read_to_string(&self.config_path)?;
// Parse YAML into JSON Value for consistent internal representation
let yaml_value: serde_yaml::Value = serde_yaml::from_str(&file_content)?;
let json_value: Value = serde_json::to_value(yaml_value)?;
match json_value {
Value::Object(map) => Ok(map.into_iter().collect()),
_ => Ok(HashMap::new()),
}
} else {
Ok(HashMap::new())
}
}
// Load current secrets from the keyring
fn load_secrets(&self) -> Result<HashMap<String, Value>, ConfigError> {
let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?;
match entry.get_password() {
Ok(content) => {
let values: HashMap<String, Value> = serde_json::from_str(&content)?;
Ok(values)
}
Err(keyring::Error::NoEntry) => Ok(HashMap::new()),
Err(e) => Err(ConfigError::KeyringError(e.to_string())),
}
}
/// Get a configuration value.
///
/// This will attempt to get the value from:
/// 1. Environment variable with the exact key name
/// 2. Configuration file
///
/// The value will be deserialized into the requested type. This works with
/// both simple types (String, i32, etc.) and complex types that implement
/// serde::Deserialize.
///
/// # Errors
///
/// Returns a ConfigError if:
/// - The key doesn't exist in either environment or config file
/// - The value cannot be deserialized into the requested type
/// - There is an error reading the config file
pub fn get<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T, ConfigError> {
// First check environment variables (convert to uppercase)
let env_key = key.to_uppercase();
if let Ok(val) = env::var(&env_key) {
// Parse the environment variable value into a serde_json::Value
let value: Value = serde_json::from_str(&val).unwrap_or(Value::String(val));
return Ok(serde_json::from_value(value)?);
}
// Load current values from file
let values = self.load_values()?;
// Then check our stored values
values
.get(key)
.ok_or_else(|| ConfigError::NotFound(key.to_string()))
.and_then(|v| Ok(serde_json::from_value(v.clone())?))
}
/// Set a configuration value in the config file.
///
/// This will immediately write the value to the config file. The value
/// can be any type that can be serialized to JSON/YAML.
///
/// Note that this does not affect environment variables - those can only
/// be set through the system environment.
///
/// # Errors
///
/// Returns a ConfigError if:
/// - There is an error reading or writing the config file
/// - There is an error serializing the value
pub fn set(&self, key: &str, value: Value) -> Result<(), ConfigError> {
let mut values = self.load_values()?;
values.insert(key.to_string(), value);
// Convert to YAML for storage
let yaml_value = serde_yaml::to_string(&values)?;
// Ensure the directory exists
if let Some(parent) = self.config_path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| ConfigError::DirectoryError(e.to_string()))?;
}
std::fs::write(&self.config_path, yaml_value)?;
Ok(())
}
/// Get a secret value.
///
/// This will attempt to get the value from:
/// 1. Environment variable with the exact key name
/// 2. System keyring
///
/// The value will be deserialized into the requested type. This works with
/// both simple types (String, i32, etc.) and complex types that implement
/// serde::Deserialize.
///
/// # Errors
///
/// Returns a ConfigError if:
/// - The key doesn't exist in either environment or keyring
/// - The value cannot be deserialized into the requested type
/// - There is an error accessing the keyring
pub fn get_secret<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T, ConfigError> {
// First check environment variables (convert to uppercase)
let env_key = key.to_uppercase();
if let Ok(val) = env::var(&env_key) {
let value: Value = serde_json::from_str(&val).unwrap_or(Value::String(val));
return Ok(serde_json::from_value(value)?);
}
// Then check keyring
let values = self.load_secrets()?;
values
.get(key)
.ok_or_else(|| ConfigError::NotFound(key.to_string()))
.and_then(|v| Ok(serde_json::from_value(v.clone())?))
}
/// Set a secret value in the system keyring.
///
/// This will store the value in a single JSON object in the system keyring,
/// alongside any other secrets. The value can be any type that can be
/// serialized to JSON.
///
/// Note that this does not affect environment variables - those can only
/// be set through the system environment.
///
/// # Errors
///
/// Returns a ConfigError if:
/// - There is an error accessing the keyring
/// - There is an error serializing the value
pub fn set_secret(&self, key: &str, value: Value) -> Result<(), ConfigError> {
let mut values = self.load_secrets()?;
values.insert(key.to_string(), value);
let json_value = serde_json::to_string(&values)?;
let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?;
entry.set_password(&json_value)?;
Ok(())
}
/// Delete a secret from the system keyring.
///
/// This will remove the specified key from the JSON object in the system keyring.
/// Other secrets will remain unchanged.
///
/// # Errors
///
/// Returns a ConfigError if:
/// - There is an error accessing the keyring
/// - There is an error serializing the remaining values
pub fn delete_secret(&self, key: &str) -> Result<(), ConfigError> {
let mut values = self.load_secrets()?;
values.remove(key);
let json_value = serde_json::to_string(&values)?;
let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?;
entry.set_password(&json_value)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use serial_test::serial;
use tempfile::NamedTempFile;
fn cleanup_keyring() -> Result<(), ConfigError> {
let entry = Entry::new(TEST_KEYRING_SERVICE, KEYRING_USERNAME)?;
match entry.delete_credential() {
Ok(_) => Ok(()),
Err(keyring::Error::NoEntry) => Ok(()),
Err(e) => Err(ConfigError::KeyringError(e.to_string())),
}
}
#[test]
fn test_basic_config() -> Result<(), ConfigError> {
let temp_file = NamedTempFile::new().unwrap();
let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?;
// Set a simple string value
config.set("test_key", Value::String("test_value".to_string()))?;
// Test simple string retrieval
let value: String = config.get("test_key")?;
assert_eq!(value, "test_value");
// Test with environment variable override
std::env::set_var("TEST_KEY", "env_value");
let value: String = config.get("test_key")?;
assert_eq!(value, "env_value");
Ok(())
}
#[test]
fn test_complex_type() -> Result<(), ConfigError> {
#[derive(Deserialize, Debug, PartialEq)]
struct TestStruct {
field1: String,
field2: i32,
}
let temp_file = NamedTempFile::new().unwrap();
let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?;
// Set a complex value
config.set(
"complex_key",
serde_json::json!({
"field1": "hello",
"field2": 42
}),
)?;
let value: TestStruct = config.get("complex_key")?;
assert_eq!(value.field1, "hello");
assert_eq!(value.field2, 42);
Ok(())
}
#[test]
fn test_missing_value() {
let temp_file = NamedTempFile::new().unwrap();
let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE).unwrap();
let result: Result<String, ConfigError> = config.get("nonexistent_key");
assert!(matches!(result, Err(ConfigError::NotFound(_))));
}
#[test]
fn test_yaml_formatting() -> Result<(), ConfigError> {
let temp_file = NamedTempFile::new().unwrap();
let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?;
config.set("key1", Value::String("value1".to_string()))?;
config.set("key2", Value::Number(42.into()))?;
// Read the file directly to check YAML formatting
let content = std::fs::read_to_string(temp_file.path())?;
assert!(content.contains("key1: value1"));
assert!(content.contains("key2: 42"));
Ok(())
}
#[test]
#[serial]
fn test_secret_management() -> Result<(), ConfigError> {
cleanup_keyring()?;
let temp_file = NamedTempFile::new().unwrap();
let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?;
// Test setting and getting a simple secret
config.set_secret("api_key", Value::String("secret123".to_string()))?;
let value: String = config.get_secret("api_key")?;
assert_eq!(value, "secret123");
// Test environment variable override
std::env::set_var("API_KEY", "env_secret");
let value: String = config.get_secret("api_key")?;
assert_eq!(value, "env_secret");
std::env::remove_var("API_KEY");
// Test deleting a secret
config.delete_secret("api_key")?;
let result: Result<String, ConfigError> = config.get_secret("api_key");
assert!(matches!(result, Err(ConfigError::NotFound(_))));
cleanup_keyring()?;
Ok(())
}
#[test]
#[serial]
fn test_multiple_secrets() -> Result<(), ConfigError> {
cleanup_keyring()?;
let temp_file = NamedTempFile::new().unwrap();
let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?;
// Set multiple secrets
config.set_secret("key1", Value::String("secret1".to_string()))?;
config.set_secret("key2", Value::String("secret2".to_string()))?;
// Verify both exist
let value1: String = config.get_secret("key1")?;
let value2: String = config.get_secret("key2")?;
assert_eq!(value1, "secret1");
assert_eq!(value2, "secret2");
// Delete one secret
config.delete_secret("key1")?;
// Verify key1 is gone but key2 remains
let result1: Result<String, ConfigError> = config.get_secret("key1");
let value2: String = config.get_secret("key2")?;
assert!(matches!(result1, Err(ConfigError::NotFound(_))));
assert_eq!(value2, "secret2");
cleanup_keyring()?;
Ok(())
}
}

View File

@@ -0,0 +1,118 @@
use super::base::Config;
use crate::agents::ExtensionConfig;
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
const DEFAULT_EXTENSION: &str = "developer";
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct ExtensionEntry {
pub enabled: bool,
#[serde(flatten)]
pub config: ExtensionConfig,
}
/// Extension configuration management
pub struct ExtensionManager;
impl ExtensionManager {
/// Get the extension configuration if enabled
pub fn get_config(name: &str) -> Result<Option<ExtensionConfig>> {
let config = Config::global();
// Try to get the extension entry
let extensions: HashMap<String, ExtensionEntry> = match config.get("extensions") {
Ok(exts) => exts,
Err(super::ConfigError::NotFound(_)) => {
// Initialize with default developer extension
let defaults = HashMap::from([(
DEFAULT_EXTENSION.to_string(),
ExtensionEntry {
enabled: true,
config: ExtensionConfig::Builtin {
name: DEFAULT_EXTENSION.to_string(),
},
},
)]);
config.set("extensions", serde_json::to_value(&defaults)?)?;
defaults
}
Err(e) => return Err(e.into()),
};
Ok(extensions.get(name).and_then(|entry| {
if entry.enabled {
Some(entry.config.clone())
} else {
None
}
}))
}
/// Set or update an extension configuration
pub fn set(entry: ExtensionEntry) -> Result<()> {
let config = Config::global();
let mut extensions: HashMap<String, ExtensionEntry> =
config.get("extensions").unwrap_or_else(|_| HashMap::new());
extensions.insert(entry.config.name().parse()?, entry);
config.set("extensions", serde_json::to_value(extensions)?)?;
Ok(())
}
/// Remove an extension configuration
pub fn remove(name: &str) -> Result<()> {
let config = Config::global();
let mut extensions: HashMap<String, ExtensionEntry> =
config.get("extensions").unwrap_or_else(|_| HashMap::new());
extensions.remove(name);
config.set("extensions", serde_json::to_value(extensions)?)?;
Ok(())
}
/// Enable or disable an extension
pub fn set_enabled(name: &str, enabled: bool) -> Result<()> {
let config = Config::global();
let mut extensions: HashMap<String, ExtensionEntry> =
config.get("extensions").unwrap_or_else(|_| HashMap::new());
if let Some(entry) = extensions.get_mut(name) {
entry.enabled = enabled;
config.set("extensions", serde_json::to_value(extensions)?)?;
}
Ok(())
}
/// Get all extensions and their configurations
pub fn get_all() -> Result<Vec<ExtensionEntry>> {
let config = Config::global();
let extensions: HashMap<String, ExtensionEntry> =
config.get("extensions").unwrap_or_default();
Ok(Vec::from_iter(extensions.values().cloned()))
}
/// Get all extension names
pub fn get_all_names() -> Result<Vec<String>> {
let config = Config::global();
Ok(config
.get("extensions")
.unwrap_or_else(|_| get_keys(Default::default())))
}
/// Check if an extension is enabled
pub fn is_enabled(name: &str) -> Result<bool> {
let config = Config::global();
let extensions: HashMap<String, ExtensionEntry> =
config.get("extensions").unwrap_or_else(|_| HashMap::new());
Ok(extensions.get(name).map(|e| e.enabled).unwrap_or(false))
}
}
fn get_keys(entries: HashMap<String, ExtensionEntry>) -> Vec<String> {
entries.into_keys().collect()
}

View File

@@ -0,0 +1,6 @@
mod base;
mod extensions;
pub use crate::agents::ExtensionConfig;
pub use base::{Config, ConfigError};
pub use extensions::{ExtensionEntry, ExtensionManager};

9
crates/goose/src/lib.rs Normal file
View File

@@ -0,0 +1,9 @@
pub mod agents;
pub mod config;
pub mod message;
pub mod model;
pub mod prompt_template;
pub mod providers;
pub mod token_counter;
pub mod tracing;
pub mod truncate;

250
crates/goose/src/message.rs Normal file
View File

@@ -0,0 +1,250 @@
use std::collections::HashSet;
/// Messages which represent the content sent back and forth to LLM provider
///
/// We use these messages in the agent code, and interfaces which interact with
/// the agent. That let's us reuse message histories across different interfaces.
///
/// The content of the messages uses MCP types to avoid additional conversions
/// when interacting with MCP servers.
use chrono::Utc;
use mcp_core::content::{Content, ImageContent, TextContent};
use mcp_core::handler::ToolResult;
use mcp_core::role::Role;
use mcp_core::tool::ToolCall;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ToolRequest {
pub id: String,
pub tool_call: ToolResult<ToolCall>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ToolResponse {
pub id: String,
pub tool_result: ToolResult<Vec<Content>>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
/// Content passed inside a message, which can be both simple content and tool content
pub enum MessageContent {
Text(TextContent),
Image(ImageContent),
ToolRequest(ToolRequest),
ToolResponse(ToolResponse),
}
impl MessageContent {
pub fn text<S: Into<String>>(text: S) -> Self {
MessageContent::Text(TextContent {
text: text.into(),
annotations: None,
})
}
pub fn image<S: Into<String>, T: Into<String>>(data: S, mime_type: T) -> Self {
MessageContent::Image(ImageContent {
data: data.into(),
mime_type: mime_type.into(),
annotations: None,
})
}
pub fn tool_request<S: Into<String>>(id: S, tool_call: ToolResult<ToolCall>) -> Self {
MessageContent::ToolRequest(ToolRequest {
id: id.into(),
tool_call,
})
}
pub fn tool_response<S: Into<String>>(id: S, tool_result: ToolResult<Vec<Content>>) -> Self {
MessageContent::ToolResponse(ToolResponse {
id: id.into(),
tool_result,
})
}
pub fn as_tool_request(&self) -> Option<&ToolRequest> {
if let MessageContent::ToolRequest(ref tool_request) = self {
Some(tool_request)
} else {
None
}
}
pub fn as_tool_response(&self) -> Option<&ToolResponse> {
if let MessageContent::ToolResponse(ref tool_response) = self {
Some(tool_response)
} else {
None
}
}
pub fn as_tool_response_text(&self) -> Option<String> {
if let Some(tool_response) = self.as_tool_response() {
if let Ok(contents) = &tool_response.tool_result {
let texts: Vec<String> = contents
.iter()
.filter_map(|content| content.as_text().map(String::from))
.collect();
if !texts.is_empty() {
return Some(texts.join("\n"));
}
}
}
None
}
/// Get the text content if this is a TextContent variant
pub fn as_text(&self) -> Option<&str> {
match self {
MessageContent::Text(text) => Some(&text.text),
_ => None,
}
}
}
impl From<Content> for MessageContent {
fn from(content: Content) -> Self {
match content {
Content::Text(text) => MessageContent::Text(text),
Content::Image(image) => MessageContent::Image(image),
Content::Resource(resource) => MessageContent::Text(TextContent {
text: resource.get_text(),
annotations: None,
}),
}
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
/// A message to or from an LLM
pub struct Message {
pub role: Role,
pub created: i64,
pub content: Vec<MessageContent>,
}
impl Message {
/// Create a new user message with the current timestamp
pub fn user() -> Self {
Message {
role: Role::User,
created: Utc::now().timestamp(),
content: Vec::new(),
}
}
/// Create a new assistant message with the current timestamp
pub fn assistant() -> Self {
Message {
role: Role::Assistant,
created: Utc::now().timestamp(),
content: Vec::new(),
}
}
/// Add any MessageContent to the message
pub fn with_content(mut self, content: MessageContent) -> Self {
self.content.push(content);
self
}
/// Add text content to the message
pub fn with_text<S: Into<String>>(self, text: S) -> Self {
self.with_content(MessageContent::text(text))
}
/// Add image content to the message
pub fn with_image<S: Into<String>, T: Into<String>>(self, data: S, mime_type: T) -> Self {
self.with_content(MessageContent::image(data, mime_type))
}
/// Add a tool request to the message
pub fn with_tool_request<S: Into<String>>(
self,
id: S,
tool_call: ToolResult<ToolCall>,
) -> Self {
self.with_content(MessageContent::tool_request(id, tool_call))
}
/// Add a tool response to the message
pub fn with_tool_response<S: Into<String>>(
self,
id: S,
result: ToolResult<Vec<Content>>,
) -> Self {
self.with_content(MessageContent::tool_response(id, result))
}
/// Get the concatenated text content of the message, separated by newlines
pub fn as_concat_text(&self) -> String {
self.content
.iter()
.filter_map(|c| c.as_text())
.collect::<Vec<_>>()
.join("\n")
}
/// Check if the message is a tool call
pub fn is_tool_call(&self) -> bool {
self.content
.iter()
.any(|c| matches!(c, MessageContent::ToolRequest(_)))
}
/// Check if the message is a tool response
pub fn is_tool_response(&self) -> bool {
self.content
.iter()
.any(|c| matches!(c, MessageContent::ToolResponse(_)))
}
/// Retrieves all tool `id` from the message
pub fn get_tool_ids(&self) -> HashSet<&str> {
self.content
.iter()
.filter_map(|content| match content {
MessageContent::ToolRequest(req) => Some(req.id.as_str()),
MessageContent::ToolResponse(res) => Some(res.id.as_str()),
_ => None,
})
.collect()
}
/// Retrieves all tool `id` from ToolRequest messages
pub fn get_tool_request_ids(&self) -> HashSet<&str> {
self.content
.iter()
.filter_map(|content| {
if let MessageContent::ToolRequest(req) = content {
Some(req.id.as_str())
} else {
None
}
})
.collect()
}
/// Retrieves all tool `id` from ToolResponse messages
pub fn get_tool_response_ids(&self) -> HashSet<&str> {
self.content
.iter()
.filter_map(|content| {
if let MessageContent::ToolResponse(res) = content {
Some(res.id.as_str())
} else {
None
}
})
.collect()
}
/// Check if the message has only TextContent
pub fn has_only_text_content(&self) -> bool {
self.content
.iter()
.all(|c| matches!(c, MessageContent::Text(_)))
}
}

142
crates/goose/src/model.rs Normal file
View File

@@ -0,0 +1,142 @@
use serde::{Deserialize, Serialize};
const DEFAULT_CONTEXT_LIMIT: usize = 128_000;
// Tokenizer names, used to infer from model name
pub const GPT_4O_TOKENIZER: &str = "Xenova--gpt-4o";
pub const CLAUDE_TOKENIZER: &str = "Xenova--claude-tokenizer";
/// Configuration for model-specific settings and limits
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
/// The name of the model to use
pub model_name: String,
// Optional tokenizer name (corresponds to the sanitized HuggingFace tokenizer name)
// "Xenova/gpt-4o" -> "Xenova/gpt-4o"
// If not provided, best attempt will be made to infer from model name or default
pub tokenizer_name: String,
/// Optional explicit context limit that overrides any defaults
pub context_limit: Option<usize>,
/// Optional temperature setting (0.0 - 1.0)
pub temperature: Option<f32>,
/// Optional maximum tokens to generate
pub max_tokens: Option<i32>,
}
impl ModelConfig {
/// Create a new ModelConfig with the specified model name
///
/// The context limit is set with the following precedence:
/// 1. Explicit context_limit if provided in config
/// 2. Model-specific default based on model name
/// 3. Global default (128_000) (in get_context_limit)
pub fn new(model_name: String) -> Self {
let context_limit = Self::get_model_specific_limit(&model_name);
let tokenizer_name = Self::infer_tokenizer_name(&model_name);
Self {
model_name,
tokenizer_name: tokenizer_name.to_string(),
context_limit,
temperature: None,
max_tokens: None,
}
}
fn infer_tokenizer_name(model_name: &str) -> &'static str {
if model_name.contains("claude") {
CLAUDE_TOKENIZER
} else {
// Default tokenizer
GPT_4O_TOKENIZER
}
}
/// Get model-specific context limit based on model name
fn get_model_specific_limit(model_name: &str) -> Option<usize> {
// Implement some sensible defaults
match model_name {
// OpenAI models, https://platform.openai.com/docs/models#models-overview
name if name.contains("gpt-4o") => Some(128_000),
name if name.contains("gpt-4-turbo") => Some(128_000),
// Anthropic models, https://docs.anthropic.com/en/docs/about-claude/models
name if name.contains("claude-3") => Some(200_000),
// Meta Llama models, https://github.com/meta-llama/llama-models/tree/main?tab=readme-ov-file#llama-models-1
name if name.contains("llama3.2") => Some(128_000),
name if name.contains("llama3.3") => Some(128_000),
_ => None,
}
}
/// Set an explicit context limit
pub fn with_context_limit(mut self, limit: Option<usize>) -> Self {
// Default is None and therefore DEFAULT_CONTEXT_LIMIT, only set
// if input is Some to allow passing through with_context_limit in
// configuration cases
if limit.is_some() {
self.context_limit = limit;
}
self
}
/// Set the temperature
pub fn with_temperature(mut self, temp: Option<f32>) -> Self {
self.temperature = temp;
self
}
/// Set the max tokens
pub fn with_max_tokens(mut self, tokens: Option<i32>) -> Self {
self.max_tokens = tokens;
self
}
// Get the tokenizer name
pub fn tokenizer_name(&self) -> &str {
&self.tokenizer_name
}
/// Get the context_limit for the current model
/// If none are defined, use the DEFAULT_CONTEXT_LIMIT
pub fn context_limit(&self) -> usize {
self.context_limit.unwrap_or(DEFAULT_CONTEXT_LIMIT)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_model_config_context_limits() {
// Test explicit limit
let config =
ModelConfig::new("claude-3-opus".to_string()).with_context_limit(Some(150_000));
assert_eq!(config.context_limit(), 150_000);
// Test model-specific defaults
let config = ModelConfig::new("claude-3-opus".to_string());
assert_eq!(config.context_limit(), 200_000);
let config = ModelConfig::new("gpt-4-turbo".to_string());
assert_eq!(config.context_limit(), 128_000);
// Test fallback to default
let config = ModelConfig::new("unknown-model".to_string());
assert_eq!(config.context_limit(), DEFAULT_CONTEXT_LIMIT);
}
#[test]
fn test_model_config_settings() {
let config = ModelConfig::new("test-model".to_string())
.with_temperature(Some(0.7))
.with_max_tokens(Some(1000))
.with_context_limit(Some(50_000));
assert_eq!(config.temperature, Some(0.7));
assert_eq!(config.max_tokens, Some(1000));
assert_eq!(config.context_limit, Some(50_000));
}
}

View File

@@ -0,0 +1,139 @@
use include_dir::{include_dir, Dir};
use serde::Serialize;
use std::path::PathBuf;
use tera::{Context, Error as TeraError, Tera};
// The prompts directory needs to be embedded in the binary (so it works when distributed)
static PROMPTS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/prompts");
pub fn load_prompt<T: Serialize>(template: &str, context_data: &T) -> Result<String, TeraError> {
let mut tera = Tera::default();
tera.add_raw_template("inline_template", template)?;
let context = Context::from_serialize(context_data)?;
let rendered = tera.render("inline_template", &context)?;
Ok(rendered.trim().to_string())
}
pub fn load_prompt_file<T: Serialize>(
template_file: impl Into<PathBuf>,
context_data: &T,
) -> Result<String, TeraError> {
let template_path = template_file.into();
// Get the file content from the embedded directory
let template_content = if let Some(file) = PROMPTS_DIR.get_file(template_path.to_str().unwrap())
{
String::from_utf8_lossy(file.contents()).into_owned()
} else {
return Err(TeraError::chain(
"Failed to find template file",
std::io::Error::new(
std::io::ErrorKind::NotFound,
"Template file not found in embedded directory",
),
));
};
load_prompt(&template_content, context_data)
}
#[cfg(test)]
mod tests {
use super::*;
use mcp_core::tool::Tool;
use serde_json::json;
use std::collections::HashMap;
#[test]
fn test_load_prompt() {
let template = "Hello, {{ name }}! You are {{ age }} years old.";
let mut context = HashMap::new();
context.insert("name".to_string(), "Alice".to_string());
context.insert("age".to_string(), 30.to_string());
let result = load_prompt(template, &context).unwrap();
assert_eq!(result, "Hello, Alice! You are 30 years old.");
}
#[test]
fn test_load_prompt_missing_variable() {
let template = "Hello, {{ name }}! You are {{ age }} years old.";
let mut context = HashMap::new();
context.insert("name".to_string(), "Alice".to_string());
// 'age' is missing from context
let result = load_prompt(template, &context);
assert!(result.is_err());
}
#[test]
fn test_load_prompt_file() {
// since we are embedding the prompts directory, the file path needs to be relative to the prompts directory
let file_path = PathBuf::from("mock.md");
let mut context = HashMap::new();
context.insert("name".to_string(), "Alice".to_string());
context.insert("age".to_string(), 30.to_string());
let result = load_prompt_file(file_path, &context).unwrap();
assert_eq!(
result,
"This prompt is only used for testing.\n\nHello, Alice! You are 30 years old."
);
}
#[test]
fn test_load_prompt_file_missing_file() {
let file_path = PathBuf::from("non_existent_template.txt");
let context: HashMap<String, String> = HashMap::new(); // Add type annotation here
let result = load_prompt_file(file_path, &context);
assert!(result.is_err());
}
#[test]
fn test_load_prompt_with_tools() {
let template = "### Tool Descriptions\n{% for tool in tools %}\n{{tool.name}}: {{tool.description}}{% endfor %}";
let tools = vec![
Tool::new(
"calculator",
"Performs basic math operations",
json!({
"type": "object",
"properties": {
"operation": {"type": "string"},
"numbers": {"type": "array"}
}
}),
),
Tool::new(
"weather",
"Gets weather information",
json!({
"type": "object",
"properties": {
"location": {"type": "string"}
}
}),
),
];
let mut context = HashMap::new();
context.insert("tools".to_string(), tools);
let result = load_prompt(template, &context).unwrap();
let expected = "### Tool Descriptions\n\ncalculator: Performs basic math operations\nweather: Gets weather information";
assert_eq!(result, expected);
}
#[test]
fn test_load_prompt_with_empty_tools() {
let template = "### Tool Descriptions\n{% for tool in tools %}\n{{tool.name}}: {{tool.description}}{% endfor %}";
let tools: Vec<Tool> = vec![];
let mut context = HashMap::new();
context.insert("tools".to_string(), tools);
let result = load_prompt(template, &context).unwrap();
let expected = "### Tool Descriptions";
assert_eq!(result, expected);
}
}

View File

@@ -0,0 +1,3 @@
This prompt is only used for testing.
Hello, {{ name }}! You are {{ age }} years old.

View File

@@ -1,11 +1,10 @@
{{synopsis.current_summary}}
You prepare plans for an agent system. You will recieve the current system
status as well as in an incoming request from the human. Your plan will be used by an AI agent,
who is taking actions on behalf of the human.
# Instructions
The agent currently has access to the following tools
Prepare a plan that an agent will use to followup to the request described above. Your plan
will be carried out by an agent with access to the following tools:
{% for tool in exchange.tools %}
{% for tool in tools %}
{{tool.name}}: {{tool.description}}{% endfor %}
If the request is simple, such as a greeting or a request for information or advice, the plan can simply be:
@@ -23,26 +22,9 @@ Your plan needs to use the following format, but can have any number of tasks.
]
```
# Hints
{{synopsis.hints}}
# Context
The current state of the agent is:
{{system.info()}}
The agent already has access to content of the following files, your plan does not need to include finding or reading these.
However if a file or code object is not here but needs to be viewed to complete the goal, include plan steps to
use `rg` (ripgrep) to find the relevant references.
{% for file in system.active_files %}
{{file.path}}{% endfor %}
# Examples
These examples show the format you should follow
These examples show the format you should follow. *Do not reply with any other text, just the json plan*
```json
[

Some files were not shown because too many files have changed in this diff Show More