diff --git a/.env.example b/.env.example index fb4fab8..e7462b3 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,9 @@ MINICODE_MODEL=deepseek-v4-pro OPENAI_BASE_URL=https://api.deepseek.com MINICODE_MODEL_TIMEOUT=120 MINICODE_MAX_STEPS=12 +# Optional cost controls; leave blank to disable either limit. +MINICODE_SESSION_TOKEN_BUDGET= +MINICODE_MAX_OUTPUT_TOKENS= # MINICODE_SYSTEM_PROMPT=You are a careful local coding assistant. MINICODE_CONTEXT_TOKENS=16000 MINICODE_CONTEXT_TRIGGER=0.8 diff --git a/README.md b/README.md index 262bfec..797a4b2 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ MiniCode Rebuild 是一个从零、分阶段实现的本地终端 AI Coding Agen ## 当前状态 -阶段 0“仓库初始化与工程基线”至阶段 10“可观测性、质量与发布准备”已经完成；阶段 11 已选择并实现独立高级能力“长期记忆与检索”。 +阶段 0“仓库初始化与工程基线”至阶段 10“可观测性、质量与发布准备”已经完成；阶段 11 已实现“长期记忆与检索”，阶段 12 已实现独立高级能力“成本控制”。 目前已经具备： @@ -38,11 +38,13 @@ MiniCode Rebuild 是一个从零、分阶段实现的本地终端 AI Coding Agen - 将脱敏生命周期元数据写入工作区 JSONL 日志，并通过时间线查看运行过程； - 在工作区本地显式保存长期记忆，并通过有界词法检索按需召回； - 将记忆结果标记为不可信历史数据，模型写入和删除仍经过权限边界； +- 为会话设置可选 token 预算，并在请求前估算输入和工具协议成本； +- 按剩余额度限制单次模型输出，预算不足时不调用 Provider； - 离线检查 Python、运行配置、Provider 配置、会话存储、Skills 与记忆存储 readiness； - 使用 Ruff、Mypy、分支覆盖率、构建、安装和跨平台 CI 作为发布质量门禁； - 执行自动化测试。 -阶段 0 至阶段 10 的基础路线和阶段 11 的“长期记忆与检索”已经完成。其余高级能力仍必须单独选择、设计、测试和提交。 +阶段 0 至阶段 10 的基础路线、阶段 11“长期记忆与检索”和阶段 12“成本控制”已经完成。其余高级能力仍必须单独选择、设计、测试和提交。 ## 长期记忆与检索 @@ -61,6 +63,19 @@ MiniCode Rebuild 是一个从零、分阶段实现的本地终端 AI Coding Agen 存储最多 500 条记忆；单条内容、标签、查询、返回数量和结果预览均有上限。检索使用无网络、无第三方依赖的确定性词法评分，适合项目约定、用户明确偏好和长期任务事实，不等同于 embedding 语义搜索。检索结果始终带有“不可信历史数据”边界，不得覆盖当前系统或用户指令，也可能已经过时。 +## Token 成本控制 + +成本控制是显式启用的 Provider 无关 token 门禁。可以限制整个持久化会话的累计 token 用量，也可以单独限制每次模型响应的最大输出： + +```powershell +minicode-rebuild --token-budget 50000 --max-output-tokens 2000 "分析当前项目" +minicode-rebuild --interactive --resume latest --token-budget 50000 +``` + +请求发出前，运行时会估算消息与工具声明占用的输入 token，并从会话剩余额度中扣除输入预留，再把允许的输出上限映射到 OpenAI-compatible `max_tokens`。如果请求至少需要的输入和一个输出 token 都无法容纳，Agent 以 `budget_exhausted` 停止，不调用 Provider。恢复会话时，门禁会继续使用已经持久化的 Provider token 统计。交互模式可用 `/budget` 查看限制、已用量和剩余额度。 + +该能力用于阻止失控的多步调用，不等同于精确账单上限。输入预算使用跨 Provider 启发式估算；服务端实际计费、缓存 token、推理 token 和价格规则由 Provider 决定。当前请求的真实用量只能在响应返回后得知，因此可能小幅越过估算值，但后续请求会使用更新后的服务端统计重新检查。若 Provider 不返回 usage，累计会话用量也无法精确增长；需要硬货币限额时仍应在 Provider 账户侧设置配额。 + ## 可观测性与 Readiness 每次 CLI 会话默认把生命周期元数据追加到工作区 `.minicode-rebuild/events.jsonl`。日志只包含时间、事件名、session ID、工具名、成功状态、错误代码和停止原因；不保存用户提示、工具参数、工具输出或 API Key。该目录已从 Git 和模型通用文件工具中隔离。 @@ -144,6 +159,8 @@ python -m minicode_rebuild --help | `OPENAI_BASE_URL` | `https://api.deepseek.com` | API 基址或完整 `/chat/completions` 地址 | | `MINICODE_MODEL_TIMEOUT` | `120` | 请求超时秒数，必须是正整数 | | `MINICODE_MAX_STEPS` | `12` | 每轮最大模型调用步数，必须是正整数 | +| `MINICODE_SESSION_TOKEN_BUDGET` | 无 | 可选的持久化会话累计 token 预算，必须是正整数 | +| `MINICODE_MAX_OUTPUT_TOKENS` | 无 | 可选的单次模型响应 token 上限，必须是正整数 | | `MINICODE_SYSTEM_PROMPT` | 内置安全提示 | 覆盖本进程使用的系统提示 | | `MINICODE_CONTEXT_TOKENS` | `16000` | 单轮输入的启发式上下文预算 | | `MINICODE_CONTEXT_TRIGGER` | `0.8` | 达到预算比例后自动压缩，范围 `(0, 1]` | @@ -190,11 +207,11 @@ minicode-rebuild --interactive --resume latest minicode-rebuild --resume "继续上次任务" ``` -交互模式提供 `/help`、`/session`、`/sessions`、`/transcript`、`/checkpoints`、`/rewind-preview [checkpoint-id]`、`/rewind [checkpoint-id]`、`/skills`、`/memory`、`/timeline`、`/stats`、`/compact` 和 `/exit`。`/rewind` 总会先显示预览，只有随后完整输入 `yes` 才修改文件；发现 Agent 写入后又有外部修改时会拒绝覆盖。 +交互模式提供 `/help`、`/session`、`/sessions`、`/transcript`、`/checkpoints`、`/rewind-preview [checkpoint-id]`、`/rewind [checkpoint-id]`、`/skills`、`/memory`、`/timeline`、`/budget`、`/stats`、`/compact` 和 `/exit`。`/rewind` 总会先显示预览，只有随后完整输入 `yes` 才修改文件；发现 Agent 写入后又有外部修改时会拒绝覆盖。 会话 JSON、长期记忆和事件日志位于工作区 `.minicode-rebuild/`，已从 Git 与内置通用文件工具中隔离。Checkpoint 只覆盖 `write_file`、`edit_file` 和 `patch_file` 的 UTF-8 文件变更；`run_command` 的任意副作用和专用记忆存储不在 Rewind 范围内。写文件、运行命令和模型发起的记忆变更仍会显示风险与操作详情，并要求选择一次允许、会话允许或拒绝。Headless 模式默认拒绝所有变更；只有明确传入 `--allow-mutations` 才会在本次运行内逐项自动批准，并在标准错误输出警告。 -每轮会输出模型步数、工具次数、模型返回的 token 用量和压缩次数。上下文估算是跨 Provider 的保守启发式，不等同于服务端精确 tokenizer；工具结果会优先裁剪，旧轮次按用户输入边界摘要，并始终保留最近完整轮次和主系统提示。会话恢复加载的是受预算约束的工作历史，`/transcript` 则保留完整、未压缩的用户消息、assistant 工具调用和工具结果。 +每轮会输出模型步数、工具次数、模型返回的 token 用量和压缩次数。上下文估算是跨 Provider 的保守启发式，不等同于服务端精确 tokenizer；工具结果会优先裁剪，旧轮次按用户输入边界摘要，并始终保留最近完整轮次和主系统提示。会话恢复加载的是受上下文预算约束的工作历史，并继续累计可选的会话 token 成本预算；`/transcript` 则保留完整、未压缩的用户消息、assistant 工具调用和工具结果。 最小的库调用边界如下： diff --git a/docs/REBUILD_LOG.md b/docs/REBUILD_LOG.md index 12a8f25..bc6ea8e 100644 --- a/docs/REBUILD_LOG.md +++ b/docs/REBUILD_LOG.md @@ -8,12 +8,12 @@ | 项目 | 内容 | |---|---| -| 当前阶段 | 阶段 11：长期记忆与检索（已完成） | -| 最近完成 | 阶段 11：长期记忆与检索 | -| 当前分支 | `codex/phase-11-memory` | -| 最新阶段实现提交 | `d06ad40 feat(phase-11): add workspace long-term memory` | -| 测试状态 | 阶段相关回归 `61 passed, 1 skipped`；全量回归 `276 passed, 3 skipped`；分支覆盖率 `85.47%` | -| 下一步 | 审核并由用户合并 Draft PR #5；其他高级能力继续保持独立阶段 | +| 当前阶段 | 阶段 12：成本控制（已完成） | +| 最近完成 | 阶段 12：Token 成本控制 | +| 当前分支 | `codex/phase-12-cost-control` | +| 最新阶段实现提交 | `fa243dc feat(phase-12): add token cost controls` | +| 测试状态 | 阶段相关回归 `107 passed`；全量回归 `298 passed, 3 skipped`；分支覆盖率 `85.58%` | +| 下一步 | 审核并由用户合并 Draft PR #6；其他高级能力继续保持独立阶段 | ## 总体架构 @@ -1803,3 +1803,71 @@ Draft PR #5 的首轮 Windows/Ubuntu、Python 3.11/3.13 四组任务都在同一 修复只在 `MemoryStore` 构造阶段区分该稳定错误码，继续拒绝操作，并增加不依赖主机符号链接权限的错误映射单元测试。修复后本地完整发布门禁为 `276 passed, 3 skipped`、覆盖率 `85.47%`，Ruff、Mypy、编译、构建与 Mock 演示全部通过；最终跨平台结果以重新触发的 PR #5 CI 为准。 `ee6f237` 推送后，GitHub Actions 的 Ubuntu 3.11、Ubuntu 3.13、Windows 3.11、Windows 3.13 四组任务全部通过。PR #5 保持 Draft、以 `master` 为基线且可合并；阶段 11 不自动修改或合并主分支。 + +## 阶段 12：成本控制 + +### 1. 开发前计划 + +- 只实现阶段 11 高级能力清单中的“成本控制”，不同时引入模型路由、MCP、多 Agent、Worktree 编排或 TUI。 +- 使用 Provider 无关的 token 数量作为稳定控制单位，不内置会随时间变化的模型价格，也不宣称计算精确货币账单。 +- 支持可选的持久化会话累计预算和单次响应输出上限；两者都未配置时保持此前行为。 +- 在每次 Provider 调用前估算消息与工具声明的输入 token，预算不足时失败关闭，不发送网络请求。 +- 将当前剩余额度映射为规范化 `ModelRequest.max_output_tokens`，OpenAI-compatible 适配器再写入 `max_tokens`。 +- 复用阶段 6 的会话 token 统计，使恢复会话继续消费同一预算；交互 CLI 通过 `/budget` 展示当前状态。 +- 为配置校验、请求估算、输出收紧、调用前拒绝、跨模型步骤累计、会话恢复、CLI 和适配器序列化补齐测试。 + +### 2. 威胁模型与非目标 + +- 防止失控循环持续调用模型：每个模型步骤都重新执行预算门禁，而不是只在一轮开始时检查一次。 +- 防止“大上下文 + 大输出上限”突破预留：输入估算先占用剩余额度，输出上限只能使用其余空间。 +- 防止无效配置静默失效：环境变量和 CLI 参数都只接受正整数，零、负数和非整数直接返回配置错误。 +- 防止恢复会话绕过累计限制：门禁使用持久化的 input/output token 统计作为已用量。 +- 本阶段不维护 Provider 价格表，不计算人民币或美元，不解析缓存、推理等厂商专有 token，也不替代 Provider 账户配额。 + +### 3. 预算模型 + +`TokenBudgetPolicy` 包含两个独立可选限制：`session_tokens` 控制一个持久化会话累计的 Provider 报告 token，`max_output_tokens` 控制每次响应的最大输出。`estimate_request_tokens()` 复用阶段 7 的中英文启发式，并额外计入工具名称、描述和 JSON Schema。 + +每次模型调用前计算： + +```text +remaining = session_budget - persisted_and_current_usage +available_output = remaining - estimated_input +request.max_output_tokens = min(configured_output_limit, available_output) +``` + +若 `available_output < 1`，Agent 以 `budget_exhausted` 停止，模型适配器不会收到请求。若只配置输出上限，则每个请求都使用固定上限；若只配置会话预算，则输出上限根据剩余额度动态收紧。 + +### 4. Agent、会话与 CLI 集成 + +`run_agent_turn()` 在消息压缩完成、构造 Provider 请求之后执行门禁，因此估算针对实际即将发送的消息。工具返回后进入下一模型步骤时会再次检查，并计入本轮前序响应的 usage。门禁拒绝不执行新的 Provider 请求，也不伪造模型回答。 + +`AgentSession` 把已经持久化的 input/output token 作为本轮起始用量，因而 `--resume` 无法重置会话预算。`/budget` 直接显示预算、Provider 报告用量、剩余额度和单次输出上限，不调用模型。Headless 和交互模式都可使用 `--token-budget`、`--max-output-tokens`，也可通过 `MINICODE_SESSION_TOKEN_BUDGET`、`MINICODE_MAX_OUTPUT_TOKENS` 配置。 + +### 5. 验收与安全回归 + +- 阶段相关回归：`107 passed`。 +- 全量回归：`298 passed, 3 skipped`。 +- 分支覆盖率：`85.58%`，达到 `85%` 门槛。 +- Ruff：`All checks passed!`。 +- Mypy：`Success: no issues found in 29 source files`。 +- 测试证明预算不足时 Provider 零调用、单次输出上限正确下传、多步工具循环重新检查、恢复会话沿用已报告用量、无效配置被拒绝。 + +当前 Windows 环境的 `3 skipped` 仍是缺少目录符号链接权限的安全测试；Ubuntu CI 会执行真实符号链接路径。该环境限制与成本控制无关。 + +### 6. 限制与后续边界 + +输入 token 是确定但近似的跨 Provider 估算，不是服务端 tokenizer。实际 usage 只能在响应后获得，所以单次请求可能因估算偏差略微越过会话预算；之后的模型步骤会使用更新后的真实统计拒绝继续调用。若 Provider 不返回 usage，累计统计无法精确增长，但单次输出上限和请求前估算仍然生效。 + +不同 Provider 对 `max_tokens`、隐藏推理 token、缓存命中和计费规则的解释可能不同。需要不可突破的货币限额时，必须同时使用 Provider 账户侧预算、限流或预付额度。本阶段不通过硬编码价格或猜测隐藏用量制造虚假的精确性。 + +### 7. Git 记录 + +- 基线：阶段 11 的 PR #5 已合并至 `master`，合并提交为 `4bc6e26`。 +- 分支：`codex/phase-12-cost-control`。 +- 实现提交：`fa243dc feat(phase-12): add token cost controls`。 +- 文档收口使用独立提交；分支推送后创建以 `master` 为基线的 Draft PR，不自动合并。 + +### 8. 跨平台验证记录 + +阶段 12 分支推送后创建 Draft PR #6。GitHub Actions 运行 `32097364115` 的 Ubuntu 3.11、Ubuntu 3.13、Windows 3.11、Windows 3.13 四组任务全部通过；PR 以 `master` 为基线并保持 Draft，不自动合并主分支。 diff --git a/src/minicode_rebuild/agent.py b/src/minicode_rebuild/agent.py index 2e8c597..522622e 100644 --- a/src/minicode_rebuild/agent.py +++ b/src/minicode_rebuild/agent.py @@ -16,6 +16,7 @@ TokenUsage, ToolCall, ) +from minicode_rebuild.cost import TokenBudgetPolicy, evaluate_budget from minicode_rebuild.tooling import ToolContext, ToolRegistry, ToolResult DEFAULT_MAX_STEPS = 12 @@ -30,6 +31,7 @@ class AgentStopReason(str, Enum): EMPTY_RESPONSE = "empty_response" MAX_STEPS = "max_steps" MODEL_ERROR = "model_error" + BUDGET_EXHAUSTED = "budget_exhausted" @dataclass(frozen=True, slots=True) @@ -136,6 +138,8 @@ def run_agent_turn( max_steps: int = DEFAULT_MAX_STEPS, tool_observer: ToolObserver | None = None, message_preparer: MessagePreparer | None = None, + token_budget: TokenBudgetPolicy | None = None, + used_tokens: int = 0, ) -> AgentResult: """Run one bounded turn until final text or an explicit stop condition.""" @@ -144,6 +148,13 @@ def run_agent_turn( raise TypeError("tools must be a ToolRegistry") if not isinstance(context, ToolContext): raise TypeError("context must be a ToolContext") + if isinstance(used_tokens, bool) or not isinstance(used_tokens, int): + raise TypeError("used_tokens must be an integer") + if used_tokens < 0: + raise ValueError("used_tokens must not be negative") + budget = token_budget or TokenBudgetPolicy() + if not isinstance(budget, TokenBudgetPolicy): + raise TypeError("token_budget must be a TokenBudgetPolicy") messages = _initial_messages( user_message=user_message, @@ -168,6 +179,25 @@ def run_agent_turn( messages=request_messages, tools=tools.model_tools(), ) + decision = evaluate_budget( + request, + budget, + used_tokens=used_tokens + usage.total_tokens, + ) + if not decision.allowed: + return _result( + content=decision.reason or "Token budget exhausted.", + stop_reason=AgentStopReason.BUDGET_EXHAUSTED, + messages=messages, + steps=step - 1, + tool_calls=tool_call_count, + usage=usage, + ) + request = ModelRequest( + messages=request_messages, + tools=request.tools, + max_output_tokens=decision.max_output_tokens, + ) try: response = model.complete(request) if not isinstance(response, ModelResponse): diff --git a/src/minicode_rebuild/cli.py b/src/minicode_rebuild/cli.py index 7ed4ea1..6796f8d 100644 --- a/src/minicode_rebuild/cli.py +++ b/src/minicode_rebuild/cli.py @@ -44,6 +44,16 @@ EXIT_INTERRUPTED = 130 +def _positive_int(value: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer") from exc + if parsed < 1: + raise argparse.ArgumentTypeError("must be greater than zero") + return parsed + + def build_parser() -> argparse.ArgumentParser: """Create the command-line parser without performing side effects.""" @@ -86,6 +96,20 @@ def build_parser() -> argparse.ArgumentParser: metavar="N", help="maximum model steps per turn (or MINICODE_MAX_STEPS)", ) + parser.add_argument( + "--token-budget", + type=_positive_int, + default=None, + metavar="N", + help="maximum reported tokens for this session (or MINICODE_SESSION_TOKEN_BUDGET)", + ) + parser.add_argument( + "--max-output-tokens", + type=_positive_int, + default=None, + metavar="N", + help="maximum output tokens per model request (or MINICODE_MAX_OUTPUT_TOKENS)", + ) parser.add_argument( "--system-prompt", default=None, @@ -137,6 +161,12 @@ def _runtime_settings( settings = replace(settings, max_steps=args.max_steps) if args.system_prompt is not None: settings = replace(settings, system_prompt=args.system_prompt) + budget = settings.token_budget_policy + if args.token_budget is not None: + budget = replace(budget, session_tokens=args.token_budget) + if args.max_output_tokens is not None: + budget = replace(budget, max_output_tokens=args.max_output_tokens) + settings = replace(settings, token_budget_policy=budget) return settings @@ -238,7 +268,7 @@ def _run_interactive( f"Session: {session.session_id}\n" "Commands: /help, /session, /sessions, /transcript, /checkpoints, " "/rewind-preview [id], /rewind [id], /skills, /memory, /timeline, " - "/stats, /compact, /exit\n" + "/budget, /stats, /compact, /exit\n" ) output.flush() while True: @@ -258,7 +288,7 @@ def _run_interactive( output.write( "Commands: /help, /session, /sessions, /transcript, /checkpoints, " "/rewind-preview [id], /rewind [id], /skills, /memory, /timeline, " - "/stats, /compact, /exit\n" + "/budget, /stats, /compact, /exit\n" ) continue if user_message == "/session": @@ -361,6 +391,9 @@ def _run_interactive( if user_message == "/timeline": output.write(session.timeline() + "\n") continue + if user_message == "/budget": + output.write(f"[budget] {session.budget_status()}\n") + continue if user_message == "/rewind-preview" or user_message.startswith("/rewind-preview "): checkpoint_id = user_message[len("/rewind-preview") :].strip() or None try: diff --git a/src/minicode_rebuild/cli_runtime.py b/src/minicode_rebuild/cli_runtime.py index bab39f0..b3450ee 100644 --- a/src/minicode_rebuild/cli_runtime.py +++ b/src/minicode_rebuild/cli_runtime.py @@ -11,6 +11,7 @@ from minicode_rebuild.config import RuntimeSettings from minicode_rebuild.context import CompactionResult, ContextManager from minicode_rebuild.core import Message, MessageRole, ModelAdapter, ToolCall +from minicode_rebuild.cost import format_budget from minicode_rebuild.hooks import HookEvent, HookManager, HookReport from minicode_rebuild.memory import ( MemoryRecord, @@ -259,6 +260,8 @@ def run(self, user_message: str) -> AgentResult: max_steps=self.settings.max_steps, tool_observer=self._observe_tool, message_preparer=self._prepare_messages, + token_budget=self.settings.token_budget_policy, + used_tokens=self.stats.input_tokens + self.stats.output_tokens, ) raw_history = tuple( message @@ -378,6 +381,14 @@ def timeline(self, *, limit: int = 100) -> str: return "Runtime event logging is disabled." return format_timeline(self.event_log.read(limit=limit)) + def budget_status(self) -> str: + """Render the configured token controls and reported session usage.""" + + return format_budget( + self.settings.token_budget_policy, + used_tokens=self.stats.input_tokens + self.stats.output_tokens, + ) + def preview_rewind(self, checkpoint_id: str | None = None) -> RewindPlan: if self.session_store is None or self.session_record is None: raise RuntimeError("Session persistence is disabled") diff --git a/src/minicode_rebuild/config.py b/src/minicode_rebuild/config.py index bc52455..5e9017a 100644 --- a/src/minicode_rebuild/config.py +++ b/src/minicode_rebuild/config.py @@ -8,6 +8,7 @@ from urllib.parse import urlsplit from minicode_rebuild.context import ContextPolicy +from minicode_rebuild.cost import TokenBudgetPolicy DEFAULT_MODEL = "deepseek-v4-pro" DEFAULT_OPENAI_BASE_URL = "https://api.deepseek.com" @@ -100,6 +101,7 @@ class RuntimeSettings: max_steps: int = DEFAULT_MAX_STEPS system_prompt: str = DEFAULT_SYSTEM_PROMPT context_policy: ContextPolicy = field(default_factory=ContextPolicy) + token_budget_policy: TokenBudgetPolicy = field(default_factory=TokenBudgetPolicy) def __post_init__(self) -> None: if isinstance(self.max_steps, bool) or not isinstance(self.max_steps, int): @@ -112,6 +114,10 @@ def __post_init__(self) -> None: raise ModelConfigurationError("MINICODE_SYSTEM_PROMPT must be text") if not isinstance(self.context_policy, ContextPolicy): raise ModelConfigurationError("context_policy must be a ContextPolicy") + if not isinstance(self.token_budget_policy, TokenBudgetPolicy): + raise ModelConfigurationError( + "token_budget_policy must be a TokenBudgetPolicy" + ) object.__setattr__(self, "system_prompt", self.system_prompt.strip()) @classmethod @@ -142,6 +148,20 @@ def integer(name: str, default: int) -> int: ) return value + def optional_integer(name: str) -> int | None: + raw = env.get(name, "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError as error: + raise ModelConfigurationError(f"{name} must be an integer") from error + if value < 1: + raise ModelConfigurationError( + f"{name} must be greater than zero" + ) + return value + trigger_name = "MINICODE_CONTEXT_TRIGGER" trigger_raw = env.get( trigger_name, str(ContextPolicy().trigger_ratio) @@ -177,4 +197,8 @@ def integer(name: str, default: int) -> int: "MINICODE_SUMMARY_TOKENS", ContextPolicy().summary_tokens ), ), + token_budget_policy=TokenBudgetPolicy( + session_tokens=optional_integer("MINICODE_SESSION_TOKEN_BUDGET"), + max_output_tokens=optional_integer("MINICODE_MAX_OUTPUT_TOKENS"), + ), ) diff --git a/src/minicode_rebuild/core.py b/src/minicode_rebuild/core.py index 1b93a4b..825e6b1 100644 --- a/src/minicode_rebuild/core.py +++ b/src/minicode_rebuild/core.py @@ -108,6 +108,7 @@ class ModelRequest: messages: tuple[Message, ...] tools: tuple[ModelTool, ...] = () + max_output_tokens: int | None = None def __post_init__(self) -> None: messages = tuple(self.messages) @@ -115,6 +116,13 @@ def __post_init__(self) -> None: raise ValueError("a model request requires at least one message") object.__setattr__(self, "messages", messages) object.__setattr__(self, "tools", tuple(self.tools)) + if self.max_output_tokens is not None: + if ( + isinstance(self.max_output_tokens, bool) + or not isinstance(self.max_output_tokens, int) + or self.max_output_tokens < 1 + ): + raise ValueError("max_output_tokens must be a positive integer or None") @dataclass(frozen=True, slots=True) diff --git a/src/minicode_rebuild/cost.py b/src/minicode_rebuild/cost.py new file mode 100644 index 0000000..fb6dd9c --- /dev/null +++ b/src/minicode_rebuild/cost.py @@ -0,0 +1,136 @@ +"""Provider-independent token budget controls for model requests.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass + +from minicode_rebuild.context import estimate_messages_tokens, estimate_text_tokens +from minicode_rebuild.core import ModelRequest + + +@dataclass(frozen=True, slots=True) +class TokenBudgetPolicy: + """Optional session and per-response token limits.""" + + session_tokens: int | None = None + max_output_tokens: int | None = None + + def __post_init__(self) -> None: + for name in ("session_tokens", "max_output_tokens"): + value = getattr(self, name) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer or None") + if value < 1: + raise ValueError(f"{name} must be greater than zero") + + @property + def enabled(self) -> bool: + return self.session_tokens is not None or self.max_output_tokens is not None + + +@dataclass(frozen=True, slots=True) +class BudgetDecision: + """One deterministic request admission decision.""" + + allowed: bool + estimated_input_tokens: int + remaining_tokens: int | None + max_output_tokens: int | None + reason: str | None = None + + +def estimate_request_tokens(request: ModelRequest) -> int: + """Estimate messages and tool declarations without a provider tokenizer.""" + + total = estimate_messages_tokens(request.messages) + for tool in request.tools: + schema = json.dumps( + dict(tool.parameters), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ) + total += 8 + total += estimate_text_tokens(tool.name) + total += estimate_text_tokens(tool.description) + total += estimate_text_tokens(schema) + return max(1, total) + + +def evaluate_budget( + request: ModelRequest, + policy: TokenBudgetPolicy, + *, + used_tokens: int, +) -> BudgetDecision: + """Fail closed when the estimated request cannot fit the remaining budget.""" + + if isinstance(used_tokens, bool) or not isinstance(used_tokens, int): + raise TypeError("used_tokens must be an integer") + if used_tokens < 0: + raise ValueError("used_tokens must not be negative") + + estimated_input = estimate_request_tokens(request) + remaining = ( + None + if policy.session_tokens is None + else max(0, policy.session_tokens - used_tokens) + ) + output_limit = policy.max_output_tokens + if remaining is not None: + available_output = remaining - estimated_input + if available_output < 1: + return BudgetDecision( + allowed=False, + estimated_input_tokens=estimated_input, + remaining_tokens=remaining, + max_output_tokens=None, + reason=( + "Token budget exhausted before the model request: " + f"estimated input={estimated_input}, remaining={remaining}." + ), + ) + output_limit = ( + available_output + if output_limit is None + else min(output_limit, available_output) + ) + + return BudgetDecision( + allowed=True, + estimated_input_tokens=estimated_input, + remaining_tokens=remaining, + max_output_tokens=output_limit, + ) + + +def format_budget(policy: TokenBudgetPolicy, *, used_tokens: int) -> str: + """Render stable, non-monetary budget status without claiming exact pricing.""" + + output = ( + "unbounded" + if policy.max_output_tokens is None + else str(policy.max_output_tokens) + ) + if policy.session_tokens is None: + return ( + f"token-budget=disabled used={used_tokens} " + f"max-output={output}" + ) + remaining = max(0, policy.session_tokens - used_tokens) + return ( + f"token-budget={policy.session_tokens} used={used_tokens} " + f"remaining={remaining} max-output={output}" + ) + + +__all__ = [ + "BudgetDecision", + "TokenBudgetPolicy", + "estimate_request_tokens", + "evaluate_budget", + "format_budget", +] diff --git a/src/minicode_rebuild/models/openai_compatible.py b/src/minicode_rebuild/models/openai_compatible.py index 0864854..9b7be6a 100644 --- a/src/minicode_rebuild/models/openai_compatible.py +++ b/src/minicode_rebuild/models/openai_compatible.py @@ -295,6 +295,8 @@ def complete(self, request: ModelRequest) -> ModelResponse: _serialize_tool(tool) for tool in request.tools ] + if request.max_output_tokens is not None: + payload["max_tokens"] = request.max_output_tokens response = self._transport.post_json( self._settings.chat_completions_url, diff --git a/tests/test_agent_loop.py b/tests/test_agent_loop.py index 30a237b..77a56bb 100644 --- a/tests/test_agent_loop.py +++ b/tests/test_agent_loop.py @@ -10,10 +10,12 @@ from minicode_rebuild.core import ( Message, MessageRole, + ModelRequest, ModelResponse, TokenUsage, ToolCall, ) +from minicode_rebuild.cost import TokenBudgetPolicy, estimate_request_tokens from minicode_rebuild.models import MockModel from minicode_rebuild.tooling import ( ToolContext, @@ -333,6 +335,68 @@ def test_max_steps_prevents_infinite_tool_loop(tmp_path: Path) -> None: assert "2" in result.content +def test_budget_sets_provider_output_limit(tmp_path: Path) -> None: + model = MockModel( + [ModelResponse(content="Done", usage=TokenUsage(10, 2))] + ) + + result = run( + tmp_path, + model, + token_budget=TokenBudgetPolicy( + session_tokens=10_000, + max_output_tokens=77, + ), + ) + + assert result.completed is True + assert model.requests[0].max_output_tokens == 77 + + +def test_budget_blocks_without_calling_provider(tmp_path: Path) -> None: + model = MockModel([ModelResponse(content="must not run")]) + + result = run( + tmp_path, + model, + token_budget=TokenBudgetPolicy(session_tokens=1), + ) + + assert result.stop_reason is AgentStopReason.BUDGET_EXHAUSTED + assert result.steps == 0 + assert result.usage == TokenUsage() + assert model.requests == () + + +def test_budget_uses_reported_usage_before_next_tool_step(tmp_path: Path) -> None: + call = ToolCall(id="call-1", name="echo", arguments={"text": "hello"}) + model = MockModel( + [ + ModelResponse(tool_calls=(call,), usage=TokenUsage(9_000, 500)), + ModelResponse(content="must not run"), + ] + ) + registry = ToolRegistry([echo_tool()]) + initial = ModelRequest( + messages=(Message(role=MessageRole.USER, content="Help me"),), + tools=registry.model_tools(), + ) + policy = TokenBudgetPolicy( + session_tokens=estimate_request_tokens(initial) + 9_500 + ) + + result = run( + tmp_path, + model, + registry, + token_budget=policy, + ) + + assert result.stop_reason is AgentStopReason.BUDGET_EXHAUSTED + assert result.steps == 1 + assert len(model.requests) == 1 + + def test_system_prompt_and_history_precede_new_user_message(tmp_path: Path) -> None: history = ( Message(role=MessageRole.USER, content="Earlier"), diff --git a/tests/test_cli.py b/tests/test_cli.py index 71364aa..e5f525c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -150,6 +150,32 @@ def test_headless_real_model_uses_environment_and_reports_stats( assert model.requests[0].messages[-1].content == "answer once" +def test_headless_token_budget_blocks_before_provider_call(tmp_path: Path) -> None: + stdout = StringIO() + model = MockModel([ModelResponse(content="must not run")]) + + code = main( + ["--cwd", str(tmp_path), "--token-budget", "1", "answer once"], + environment={"OPENAI_API_KEY": "secret"}, + stdin=StringIO(), + stdout=stdout, + stderr=StringIO(), + model=model, + ) + + assert code == 1 + assert "Token budget exhausted" in stdout.getvalue() + assert model.requests == () + + +@pytest.mark.parametrize("option", ["--token-budget", "--max-output-tokens"]) +def test_cli_rejects_non_positive_budget_flags(option: str) -> None: + with pytest.raises(SystemExit) as raised: + main([option, "0"]) + + assert raised.value.code == 2 + + def test_headless_can_read_prompt_from_stdin(tmp_path: Path) -> None: model = MockModel([ModelResponse(content="From pipe")]) stdout = StringIO() diff --git a/tests/test_cli_runtime.py b/tests/test_cli_runtime.py index bcd4866..b76a6a9 100644 --- a/tests/test_cli_runtime.py +++ b/tests/test_cli_runtime.py @@ -13,6 +13,7 @@ from minicode_rebuild.config import RuntimeSettings from minicode_rebuild.context import ContextManager, ContextPolicy from minicode_rebuild.core import ModelResponse, TokenUsage +from minicode_rebuild.cost import TokenBudgetPolicy from minicode_rebuild.hooks import HookEvent, HookManager from minicode_rebuild.memory import MemoryStore from minicode_rebuild.models import MockModel @@ -57,6 +58,34 @@ def test_session_accumulates_turn_and_usage_stats(tmp_path: Path) -> None: ) +def test_session_budget_uses_persisted_reported_usage(tmp_path: Path) -> None: + settings = RuntimeSettings( + max_steps=3, + system_prompt="Be precise", + token_budget_policy=TokenBudgetPolicy(session_tokens=100), + ) + session = AgentSession( + model=MockModel( + [ModelResponse(content="one", usage=TokenUsage(60, 20))] + ), + tools=ToolRegistry(), + context=ToolContext(tmp_path), + settings=settings, + output=StringIO(), + ) + + first = session.run("first") + second = session.run("second") + + assert first.completed is True + assert second.stop_reason is AgentStopReason.BUDGET_EXHAUSTED + assert session.stats.input_tokens == 60 + assert session.stats.output_tokens == 20 + assert session.budget_status() == ( + "token-budget=100 used=80 remaining=20 max-output=unbounded" + ) + + def test_permission_prompt_explains_request_and_parses_choices() -> None: request = PermissionRequest( operation="write_file", diff --git a/tests/test_config.py b/tests/test_config.py index ff58dc1..a9df6e8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -94,6 +94,8 @@ def test_runtime_settings_load_cli_environment() -> None: "MINICODE_KEEP_RECENT_TURNS": "3", "MINICODE_TOOL_RESULT_TOKENS": "700", "MINICODE_SUMMARY_TOKENS": "500", + "MINICODE_SESSION_TOKEN_BUDGET": "12000", + "MINICODE_MAX_OUTPUT_TOKENS": "800", } ) @@ -104,6 +106,8 @@ def test_runtime_settings_load_cli_environment() -> None: assert settings.context_policy.keep_recent_turns == 3 assert settings.context_policy.tool_result_tokens == 700 assert settings.context_policy.summary_tokens == 500 + assert settings.token_budget_policy.session_tokens == 12000 + assert settings.token_budget_policy.max_output_tokens == 800 @pytest.mark.parametrize("value", ["zero", "0", "-1"]) @@ -128,3 +132,19 @@ def test_runtime_settings_reject_invalid_context_policy( ) -> None: with pytest.raises(ModelConfigurationError, match=variable): RuntimeSettings.from_env({variable: value}) + + +@pytest.mark.parametrize( + ("variable", "value"), + [ + ("MINICODE_SESSION_TOKEN_BUDGET", "0"), + ("MINICODE_SESSION_TOKEN_BUDGET", "bad"), + ("MINICODE_MAX_OUTPUT_TOKENS", "-1"), + ], +) +def test_runtime_settings_reject_invalid_token_budget( + variable: str, + value: str, +) -> None: + with pytest.raises(ModelConfigurationError, match=variable): + RuntimeSettings.from_env({variable: value}) diff --git a/tests/test_core.py b/tests/test_core.py index a6356d6..9f8f07b 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -59,3 +59,12 @@ def test_model_tool_rejects_invalid_function_name(name: str) -> None: def test_model_request_requires_at_least_one_message() -> None: with pytest.raises(ValueError, match="message"): ModelRequest(messages=()) + + +@pytest.mark.parametrize("value", [0, -1, True]) +def test_model_request_rejects_invalid_output_limit(value: int) -> None: + with pytest.raises(ValueError, match="max_output_tokens"): + ModelRequest( + messages=(Message(role=MessageRole.USER, content="hello"),), + max_output_tokens=value, + ) diff --git a/tests/test_cost.py b/tests/test_cost.py new file mode 100644 index 0000000..c5566bb --- /dev/null +++ b/tests/test_cost.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import pytest + +from minicode_rebuild.core import Message, MessageRole, ModelRequest, ModelTool +from minicode_rebuild.cost import ( + TokenBudgetPolicy, + estimate_request_tokens, + evaluate_budget, + format_budget, +) + + +def request() -> ModelRequest: + return ModelRequest( + messages=(Message(role=MessageRole.USER, content="inspect README"),), + tools=( + ModelTool( + name="read_file", + description="Read one file.", + parameters={ + "type": "object", + "properties": {"path": {"type": "string"}}, + }, + ), + ), + ) + + +def test_request_estimate_includes_tool_protocol() -> None: + without_tools = ModelRequest(messages=request().messages) + + assert estimate_request_tokens(request()) > estimate_request_tokens(without_tools) + + +def test_budget_caps_output_to_remaining_session_allowance() -> None: + model_request = request() + estimate = estimate_request_tokens(model_request) + policy = TokenBudgetPolicy( + session_tokens=estimate + 25, + max_output_tokens=100, + ) + + decision = evaluate_budget(model_request, policy, used_tokens=5) + + assert decision.allowed is True + assert decision.remaining_tokens == estimate + 20 + assert decision.max_output_tokens == 20 + + +def test_budget_fails_closed_before_an_oversized_request() -> None: + model_request = request() + estimate = estimate_request_tokens(model_request) + + decision = evaluate_budget( + model_request, + TokenBudgetPolicy(session_tokens=estimate), + used_tokens=0, + ) + + assert decision.allowed is False + assert decision.max_output_tokens is None + assert "estimated input" in (decision.reason or "") + + +@pytest.mark.parametrize("value", [0, -1, True, "10"]) +def test_policy_rejects_invalid_limits(value: object) -> None: + with pytest.raises((TypeError, ValueError)): + TokenBudgetPolicy(session_tokens=value) # type: ignore[arg-type] + + +def test_budget_status_distinguishes_disabled_and_bounded_sessions() -> None: + assert format_budget(TokenBudgetPolicy(), used_tokens=3) == ( + "token-budget=disabled used=3 max-output=unbounded" + ) + assert format_budget( + TokenBudgetPolicy(session_tokens=10, max_output_tokens=4), + used_tokens=7, + ) == "token-budget=10 used=7 remaining=3 max-output=4" diff --git a/tests/test_openai_compatible.py b/tests/test_openai_compatible.py index eca966f..1b574cb 100644 --- a/tests/test_openai_compatible.py +++ b/tests/test_openai_compatible.py @@ -108,6 +108,33 @@ def test_adapter_normalizes_text_response_and_usage() -> None: assert call["timeout"] == 30 +def test_adapter_sends_normalized_output_limit() -> None: + transport = FakeTransport( + [ + json_response( + { + "choices": [ + { + "message": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + } + ] + } + ) + ] + ) + adapter = OpenAICompatibleAdapter(settings(), transport=transport) + + adapter.complete( + ModelRequest( + messages=(Message(role=MessageRole.USER, content="Hi"),), + max_output_tokens=321, + ) + ) + + assert transport.calls[0]["payload"]["max_tokens"] == 321 + + def test_adapter_serializes_tools_and_normalizes_tool_calls() -> None: transport = FakeTransport( [