# AGENTS.md

This file provides guidance to coding agents when working with code in this repository.

## Project Overview

Cog is a tool that packages machine learning models in production-ready containers. 

It consists of:
- **Cog CLI** (`cmd/cog/`) - Command-line interface for building, running, and deploying models, written in Go
- **Python SDK** (`python/cog/`) - Python library for defining model predictors and training in Python
- **Coglet** (`crates/`) - Rust-based prediction server that runs inside containers, with Python bindings via PyO3

Documentation for the CLI and SDK is available by reading ./docs/llms.txt.

## Development Commands

Development tasks are managed with [mise](https://mise.jdx.dev/). Run `mise tasks` to see all available tasks.

### Quick Reference

| Task | Description |
|------|-------------|
| `mise run fmt` | Check formatting (all languages) |
| `mise run fmt:fix` | Fix formatting (all languages) |
| `mise run lint` | Run linters (all languages) |
| `mise run lint:fix` | Fix lint issues (all languages) |
| `mise run test:go` | Run Go tests |
| `mise run test:rust` | Run Rust tests |
| `mise run test:python` | Run Python tests |
| `mise run test:integration` | Run integration tests |
| `mise run build:cog` | Build cog CLI binary |
| `mise run build:coglet` | Build coglet wheel (dev) |
| `mise run build:sdk` | Build SDK wheel |
| `mise run install` | Build and symlink cog to /usr/local/bin |
| `mise run docs:llm` | **IMPORTANT:** Regenerate `docs/llms.txt` after editing docs |
| `mise run docs:cli` | Generate CLI reference docs from Go source code |
| `mise run version` | Show current version from VERSION.txt |
| `mise run version:bump <ver>` | Bump version everywhere and commit |

### Task Naming Convention

Tasks follow a consistent naming pattern:

- **Language-based tasks** for fmt/lint/test/typecheck: `task:go`, `task:rust`, `task:python`
- **Component-based tasks** for build: `build:cog`, `build:coglet`, `build:sdk`
- **Check vs Fix**: `fmt` and `lint` default to check mode (non-destructive); use `:fix` suffix to auto-fix

### All Tasks by Category

**Format:**
- `mise run fmt` / `mise run fmt:check` - Check all (alias)
- `mise run fmt:fix` - Fix all
- `mise run fmt:go` / `mise run fmt:rust` / `mise run fmt:python` - Per-language

**Lint:**
- `mise run lint` / `mise run lint:check` - Check all (alias)
- `mise run lint:fix` - Fix all
- `mise run lint:go` / `mise run lint:rust` / `mise run lint:python` - Per-language
- `mise run lint:rust:deny` - Check Rust licenses/advisories

**Test:**
- `mise run test:go` - Go unit tests
- `mise run test:rust` - Rust unit tests
- `mise run test:python` - Python unit tests (via tox)
- `mise run test:coglet:python` - Coglet Python binding tests
- `mise run test:integration` - Integration tests

**Build:**
- `mise run build:cog` - Build cog CLI (development)
- `mise run build:cog:release` - Build cog CLI (release)
- `mise run build:coglet` - Build coglet wheel (dev install)
- `mise run build:coglet:wheel` - Build coglet wheel (native platform)
- `mise run build:coglet:wheel:linux-x64` - Build for Linux x86_64
- `mise run build:coglet:wheel:linux-arm64` - Build for Linux ARM64
- `mise run build:sdk` - Build SDK wheel

**Install:**
- `mise run install` - Symlink cog CLI to `/usr/local/bin` (requires `build:cog` first)
- `PREFIX=/custom/path mise run install` - Symlink to custom location

**Version:**
- `mise run version` - Show current version from VERSION.txt
- `mise run version:bump <ver>` - Bump version everywhere and commit
- `mise run version:check` - Verify VERSION.txt matches Cargo.toml

**Other:**
- `mise run typecheck` - Type check all languages
- `mise run generate` - Run code generation
- `mise run clean` - Clean all build artifacts
- `mise run docs` - Build documentation
- `mise run docs:serve` - Serve docs locally

## Code Style Guidelines

### Go
- **Imports**: Organize in three groups separated by blank lines: (1) Standard library, (2) Third-party packages, (3) Internal packages (`github.com/replicate/cog/pkg/...`)
- **Formatting**: Use `mise run fmt:go:fix`
- **Linting**: Must pass golangci-lint with: errcheck, gocritic, gosec, govet, ineffassign, misspell, revive, staticcheck, unused
- **Error Handling**: Return errors as values; use `pkg/errors.CodedError` for user-facing errors with error codes
- **Naming**: CamelCase for exported, camelCase for unexported
- **Testing**: Use `testify/require` for assertions; prefer table-driven tests

Example import block:
```go
import (
    "fmt"
    
    "github.com/spf13/cobra"
    
    "github.com/replicate/cog/pkg/config"
)
```

### Python
- **Imports**: Automatically organized by ruff/isort (stdlib → third-party → local)
- **Formatting**: Use `mise run fmt:python:fix`
- **Linting**: Must pass ruff checks: E (pycodestyle), F (Pyflakes), I (isort), W (warnings), S (bandit), B (bugbear), ANN (annotations)
- **Type Annotations**: Required on all function signatures; use `typing_extensions` for compatibility; avoid `Any` where possible
- **Error Handling**: Raise exceptions with descriptive messages; avoid generic exception catching
- **Naming**: snake_case for functions/variables/modules, PascalCase for classes
- **Testing**: Use pytest with fixtures; async tests with pytest-asyncio
- **Compatibility**: Must support Python 3.10-3.13

### Rust
- **Formatting**: Use `mise run fmt:rust:fix`
- **Linting**: Must pass `mise run lint:rust` (clippy)
- **Dependencies**: Audited with `cargo-deny` (see `crates/deny.toml`); run `mise run lint:rust:deny`
- **Error Handling**: Use `thiserror` for typed errors, `anyhow` for application errors
- **Naming**: snake_case for functions/variables, PascalCase for types
- **Testing**: Use `cargo test`; snapshot tests use `insta`
- **Async**: tokio runtime; async/await patterns

## Working on the CLI and support tooling
The CLI code is in the `cmd/cog/` and `pkg/` directories. Support tooling is in the `tools/` directory. 

The main commands for working on the CLI are:
- `go run ./cmd/cog` - Runs the Cog CLI directly from source (requires wheel to be built first)
- `mise run build:cog` - Builds the Cog CLI binary
- `mise run install` - Symlinks the built binary to `/usr/local/bin` (run `build:cog` first), or to a custom path with `PREFIX=/custom/path mise run install`
- `mise run test:go` - Runs all Go unit tests
- `go test ./pkg/...` - Runs tests directly with `go test`

## Working on the Python SDK
The Python SDK is developed in the `python/cog/` directory. It uses `uv` for virtual environments and `tox` for testing across multiple Python versions.

The main commands for working on the SDK are:
- `mise run build:sdk` - Builds the Python wheel
- `mise run test:python` - Runs Python tests across all supported versions

## Working on Coglet (Rust)
Coglet is the Rust-based prediction server that runs inside Cog containers, handling HTTP requests, worker process management, and prediction execution.

The code is in the `crates/` directory:
- `crates/coglet/` - Core Rust library (HTTP server, worker orchestration, IPC)
- `crates/coglet-python/` - PyO3 bindings for Python predictor integration (requires Python 3.10+)

For detailed architecture documentation, see `crates/README.md` and `crates/coglet/README.md`.

The main commands for working on Coglet are:
- `mise run build:coglet` - Build and install coglet wheel for development (macOS, for local Rust/Python tests)
- `mise run build:coglet:wheel:linux-x64` - Build Linux x86_64 wheel (required to test Rust changes in Docker containers via `cog predict`/`cog train`)
- `mise run test:rust` - Run Rust unit tests
- `mise run lint:rust` - Run clippy linter
- `mise run fmt:rust:fix` - Format code

### Testing
Go code is tested using the built-in `go test` framework:
- `go test ./pkg/... -run <name>` - Runs specific Go tests by name
- `mise run test:go` - Runs all Go unit tests

Python code is tested using `tox`, which allows testing across multiple Python versions and configurations:
- `mise run test:python` - Runs all Python unit tests
- `uv run tox -e py312-tests -- python/tests/server/test_http.py::test_openapi_specification_with_yield` - Runs a specific Python test

The integration test suite in `integration-tests/` tests the end-to-end functionality of the Cog CLI and Python SDK using Go's testscript framework:
- `mise run test:integration` - Runs the integration tests
- `mise run test:integration string_predictor` - Runs a specific integration test

The integration tests require a built Cog binary, which defaults to the first `cog` in `PATH`. Run tests against a specific binary with the `COG_BINARY` environment variable:
```bash
mise run build:cog
COG_BINARY=dist/go/*/cog mise run test:integration
```

### Development Workflow
1. Run `mise install` to set up the development environment
2. Run `mise run build:sdk` after making changes to the `./python` directory
3. Run `mise run build:coglet:wheel:linux-x64` after making changes to the `./crates` directory (needed for Docker testing)
4. Run `mise run build:cog` to build the CLI (wheels are picked up from `dist/` at Docker build time, not embedded in the binary)
5. Run `mise run fmt:fix` to format code
6. Run `mise run lint` to check code quality
7. Run `mise run docs:llm` to regenerate `docs/llms.txt` after changing `README.md` or any `docs/*.md` file
8. Read the `./docs` directory and make sure the documentation is up to date

**IMPORTANT:** Always run `mise run lint` (or the language-specific variant, e.g. `mise run lint:go`) before committing to catch linter errors early. CI will reject PRs that fail lint checks.

## Architecture

### CLI Architecture (Go)
The CLI follows a command pattern with subcommands. The main components are:
- `pkg/cli/` - Command definitions (build, run, predict, serve, etc.)
- `pkg/docker/` - Docker client and container management
- `pkg/dockerfile/` - Dockerfile generation and templating
- `pkg/config/` - cog.yaml parsing and validation
- `pkg/image/` - Image building and pushing logic

### Python SDK Architecture
- `python/cog/` - Core SDK
  - `base_predictor.py` - Base class for model predictors
  - `types.py` - Input/output type definitions
  - `server/` - HTTP/queue server implementation
  - `command/` - Runner implementations for predict/train

### Coglet Architecture (Rust)
The prediction server that runs inside Cog containers. Uses a two-process architecture: a parent process (HTTP server + orchestrator) and a worker subprocess (Python predictor execution).

See `crates/README.md` for detailed architecture documentation.

- `crates/coglet/` - Core Rust library (HTTP server, worker orchestration, IPC bridge)
- `crates/coglet-python/` - PyO3 bindings for Python predictor integration

### Key Design Patterns
1. **Local Wheel Resolution**: The CLI discovers SDK and coglet wheels from `dist/` at Docker build time (not embedded in the binary)
2. **Docker SDK Integration**: Uses Docker Go SDK for container operations
3. **Type Safety**: Dataclasses for Python type validation, strongly typed Go interfaces
4. **Compatibility Matrix**: Automated CUDA/PyTorch/TensorFlow compatibility management

For comprehensive architecture documentation, see [`architecture/`](./architecture/00-overview.md).

## Common Tasks

### Adding a new CLI command
1. Create command file in `pkg/cli/`
2. Add command to `pkg/cli/root.go`
3. Implement business logic in appropriate `pkg/` subdirectory
4. Add tests

### Modifying Python SDK behavior
1. Edit files in `python/cog/`
2. Run `mise run build:sdk` to rebuild wheel
3. Test with `mise run test:python`
4. Integration test with `mise run test:integration`

### Updating ML framework compatibility
1. See `tools/compatgen/` for compatibility matrix generation
2. Update framework versions in relevant Dockerfile templates
3. Test with various framework combinations

### Updating the docs
- Documentation is in the `docs/` directory, written in Markdown and generated into HTML using `mkdocs`.
- **IMPORTANT:** After editing any file in `docs/` or `README.md`, you MUST run `mise run docs:llm` to regenerate `docs/llms.txt`. This file is used by coding agents and should be kept in sync with the documentation.
- **IMPORTANT:** CLI reference docs (`docs/cli.md`) are auto-generated from Go source code. After modifying CLI commands in `cmd/` or `pkg/cli/`, run `mise run docs:cli` to regenerate, and ensure `mise run docs:cli:check` passes before committing.

## CI Tool Dependencies

Development tools are managed in **two places** that must be kept in sync:

1. **`mise.toml`** — Tool versions for local development (uses aqua backend for prebuilt binaries)
2. **`.github/workflows/ci.yaml`** — Tool installation for CI (uses dedicated GitHub Actions)

CI deliberately avoids aqua downloads from GitHub Releases to prevent transient 502 failures. Instead, it uses:

| Tool | CI installation method | Why |
|------|----------------------|-----|
| gotestsum | `go install` | Uses Go module proxy, not GitHub Releases |
| cargo-deny | `taiki-e/install-action` | Prebuilt with checksum verification |
| cargo-nextest | `taiki-e/install-action` | Prebuilt with checksum verification |
| coglet wheel (maturin+zig) | `PyO3/maturin-action` | Bundles maturin and zig |
| golangci-lint | `golangci/golangci-lint-action` | Built-in caching |
| Rust toolchain | `dtolnay/rust-toolchain` | Guaranteed ordering |

Tools disabled in CI are listed in `MISE_DISABLE_TOOLS` in `ci.yaml`.

**When updating a tool version**, update both:
- The version in `mise.toml` (for local dev)
- The corresponding version pin in `.github/workflows/ci.yaml` (for CI)

## Important Files
- `VERSION.txt` - Canonical version (single source of truth)
- `cog.yaml` - User-facing model configuration
- `pkg/config/config.go` - Go code for parsing and validating `cog.yaml`
- `pkg/config/data/config_schema_v1.0.json` - JSON schema for `cog.yaml`
- `python/cog/base_predictor.py` - Predictor interface
- `crates/Cargo.toml` - Rust workspace configuration (version must match VERSION.txt)
- `crates/README.md` - Coglet architecture overview
- `mise.toml` - Task definitions for development workflow

## Testing Philosophy
- Unit tests for individual components (Go and Python)
- Integration tests for end-to-end workflows
- Tests use real Docker operations (no mocking Docker API)
- Always run `mise run build:sdk` after making Python changes before testing Go code
- Python 3.10-3.13 compatibility is required

### Go Test Conventions
All Go tests must use [testify](https://github.com/stretchr/testify) for assertions. Do **not** use raw `if` checks with `t.Fatal`/`t.Errorf` — use `require` and `assert` instead.

- **`require`** — for fatal assertions that should stop the test (setup failures, preconditions):
  ```go
  require.NoError(t, err, "failed to create client")
  require.Equal(t, expected, actual)
  require.True(t, condition, "server should be ready")
  ```
- **`assert`** — for non-fatal checks where the test should continue (e.g. validating multiple fields in a loop):
  ```go
  assert.Equal(t, http.StatusOK, resp.StatusCode)
  assert.Contains(t, output, "expected substring")
  assert.NoError(t, err, "prediction %d failed", i)
  ```
- Use `require` for errors in setup/teardown and `assert` for the actual test expectations
- Prefer specific assertions (`Equal`, `Contains`, `NoError`, `Len`, `Less`) over generic `True`/`False` — they produce better failure messages
- Prefer table-driven tests for testing multiple similar cases
