Agent Tools & Budget

• Tool
• internal/agent/tools.go
• internal/agent/budget.go

When astrolabe plans and drafts a wiki page, the model doesn't work from a one-shot context dump — it can look the codebase up while it thinks. That lookup layer lives in internal/agent/tools.go: a small catalogue of ReAct-style tools that are executed by astrolabe itself against the in-memory codegraph index, with results fed back into the conversation tools.go, file header. Its companion, internal/agent/budget.go, guards the other side of the prompt: the digest of the whole graph that each stage prompt carries, kept small enough for medium-sized models by a deterministic trimming ladder.

How it works

The tool interface

Everything the model can call is a Tool: a name, a description, a JSON-Schema string for its arguments, and a Run function that turns raw JSON arguments into a string result Tool. Output is capped before it reaches the model — capResult truncates anything past toolResultCap (8000 characters) with an explicit "(truncated)" marker toolResultCap. Tool implementations are methods on toolEnv, a small struct that binds one agent run's graph data (*graphData), the repository source reader, and the repo root toolEnv.

The catalogue itself is toolEnv::allTools, which returns the five tools below toolEnv::allTools:

Each implementation is deliberately plain. toolEnv::searchSymbols lowercases the query and does substring matching over name, qualified name and file path, skipping import and file nodes, sorts by path, and caps the list at 20 toolEnv::searchSymbols. toolEnv::exploreArea first tries an exact qualified-name lookup in the graph's byQN index; failing that it treats the query as a directory prefix and collects every non-import node under it, then appends non-contains edges between the collected nodes (up to 30) so the model sees not just what is in a package but how it connects toolEnv::exploreArea. toolEnv::readSource delegates to the provider's source reader, defaulting end to start+40 when omitted and hard-clamping any request to 200 lines toolEnv::readSource. toolEnv::symbolNeighbors resolves a qualified name (and tells the model to use search_symbols first if it doesn't resolve) and renders each incident edge with an arrow and edge kind, capped at 40 toolEnv::symbolNeighbors.

How calls reach the model

Tools travel as native function calls, not as prompt text. toolDefs converts the catalogue into llm.ToolDef values — name, description, JSON-Schema parameters — with an empty parameter schema substituted for no-argument tools like list_files toolDefs. The system prompt's TOOLS section therefore says nothing about call formats; it only coaches when to look up and reminds the model to batch independent lookups into one turn, because the wire format already carries the catalogue toolsSection. Backends without native tool support (ACP agents) adapt plain chat to this interface themselves — see ACP Agent Transport and LLM Client Layer.

The round-trip is driven by Agent::runToolLoop in Agent Core: it requires a client implementing llm.ToolClient (and fails fast otherwise), then alternates model turns with tool execution up to maxRounds Agent::runToolLoop. The final round sends no tool definitions at all — with tool_choice: "none" alone some chat templates still list the tools, and models like Qwen then emit their native tool-call markup as text instead of answering Agent::runToolLoop. Each structured ToolCall is executed by Agent::runToolCall, which wraps execution in an execute_tool {name} span and returns tool errors as text so the model can self-correct Agent::runToolCall. Rejected or empty replies are nudged with a corrective user message, and the final round gets one extra tools-free attempt before giving up Agent::runToolLoop.

The loop also budgets rounds, not just characters: after each batch of tool results, the last result message is appended with a countdown like [round 3 of 20: 16 lookup rounds left. Batch independent lookups into one turn.], or a "that was your last lookup" notice, so the model can plan ahead instead of running into the cap Agent::runToolLoop.

The graph behind the tools

All five tools read from a single graphData snapshot loaded once per run by Agent::loadGraph: files, nodes of the meaningful kinds (functions, methods, structs, interfaces, routes, variables), and edges, indexed by qualified name and node ID. Anything under the output directory is stripped so a bundle written inside the repo doesn't document its own source snapshot Agent::loadGraph. This snapshot is also the raw material for the stage prompt digest — graphData::Digest flattens it into Digest/NodeDigest/EdgeDigest values Agent::loadGraph.

The prompt budget

internal/agent/budget.go keeps the digest small enough for medium-sized models. The defaults assume ~4 characters per token: DefaultPlanBudget is 48 000 characters (≈12k tokens for the plan stage) and DefaultDraftBudget is 24 000 (≈6k for drafts) DefaultPlanBudget.

Size is measured by serializing the payload: digestSize marshals a Digest to JSON and returns its length digestSize. When it doesn't fit, fitDigest applies a deterministic trimming ladder — always in the same order, never probabilistically fitDigest:

1. clear docstrings
2. clear signatures
3. drop variables
4. drop unexported functions, methods and structs
5. drop methods
6. drop functions

Each step is re-measured, and the first one that fits wins. If even the last step doesn't fit, the trimmed result is returned anyway: the model context is the hard limit and the budget is only a guard. Every dropping step prunes edges whose endpoints no longer exist via dropDanglingEdges, so the digest never contains references to nodes the model can't see dropDanglingEdges.

Draft prompts are treated differently. Agent::warnOversizeDraft never trims a draft request — a page must always see every symbol it covers, or its citations would be starved. Instead it logs a warning suggesting fewer symbols per page when the serialized DraftRequest exceeds the draft budget warnOversizeDraft. The verification side of the pipeline that consumes those citations is described in Verification (Checker & Review).

Relationships

flowchart LR
    A["Agent::runToolLoop"] -->|"catalogue as defs"| TD["toolDefs"]
    A -->|"executes calls"| RC["Agent::runToolCall"]
    RC -->|"runs"| T["Tool"]
    TE["toolEnv"] -->|"implements"| T
    TE --> AT["toolEnv::allTools"]
    SS["toolEnv::searchSymbols"] --> TE
    EA["toolEnv::exploreArea"] --> TE
    RS["toolEnv::readSource"] --> TE
    SN["toolEnv::symbolNeighbors"] --> TE
    LF["toolEnv::listFiles"] --> TE
    TE --> G["graphData"]
    G -->|"loaded by"| LG["Agent::loadGraph"]
    G -->|"Digest()"| D["Digest"]
    FD["fitDigest"] -->|"trims"| D
    DS["digestSize"] -->|"measures"| D
    WOD["Agent::warnOversizeDraft"] -->|"reports over budget"| D
    TC["llm.ToolClient"] -->|"ChatTools"| A

The layer connects outward to the graph model (see Graph Model & Providers), to the client interface it requires (see LLM Client Layer), and to the stage loop that consumes it (see Agent Core).

Reference

Generated by astrolabe from commit ecc180b3 (tag 0.4.0).