From bc1d050a0291e93e5759617238bf6178ae914f1d Mon Sep 17 00:00:00 2001 From: Rahul Raj Date: Wed, 7 Oct 2026 09:21:23 +0200 Subject: [PATCH] fix: harden agent hooks, MCP, sessions and token accounting --- CHANGELOG.md | 68 +++++++++ docs/agent-review-2026-10-06.md | 55 +++++++ docs/guides/context-management.md | 4 +- docs/guides/hooks.md | 56 +++++++- docs/reliability-hardening.md | 135 ++++++++++++++++++ .../77_super_agent_live_validation.ipynb | 2 +- shipit_agent/agent.py | 14 +- shipit_agent/async_runtime.py | 3 + shipit_agent/chat_session.py | 36 +++-- shipit_agent/hooks.py | 37 ++++- shipit_agent/llms/anthropic_adapter.py | 9 +- shipit_agent/llms/litellm_adapter.py | 13 +- shipit_agent/llms/openai_adapter.py | 7 +- shipit_agent/llms/usage.py | 33 +++++ shipit_agent/mcp.py | 82 ++++++++--- shipit_agent/mcp_resilience.py | 21 ++- shipit_agent/mcp_schema_cache.py | 8 ++ shipit_agent/permissions.py | 11 +- shipit_agent/runtime_core.py | 8 +- shipit_agent/session_manager.py | 16 ++- shipit_agent/stores/session.py | 31 +++- shipit_agent/tools/_shared.py | 2 +- .../tools/edit_file/edit_file_tool.py | 5 + .../tools/recall_result/recall_result_tool.py | 23 ++- tests/test_app_build_workflow.py | 56 ++++++++ tests/test_async_chat_session.py | 132 +++++++++++++++++ tests/test_custom_hook_contracts.py | 79 ++++++++++ tests/test_file_session_identity.py | 42 ++++++ tests/test_mcp_protocol_boundaries.py | 99 +++++++++++++ tests/test_mcp_resilience.py | 12 +- tests/test_mcp_result_integrity.py | 55 +++++++ tests/test_mcp_schema_cache.py | 18 +++ tests/test_result_recall.py | 24 ++++ tests/test_session_manager.py | 38 +++++ tests/test_tool_output_budget.py | 17 +++ tests/test_tool_workspace_boundaries.py | 32 +++++ tests/test_usage_normalization.py | 103 +++++++++++++ tests/test_workspace.py | 2 +- 38 files changed, 1301 insertions(+), 87 deletions(-) create mode 100644 docs/agent-review-2026-10-06.md create mode 100644 docs/reliability-hardening.md create mode 100644 tests/test_app_build_workflow.py create mode 100644 tests/test_async_chat_session.py create mode 100644 tests/test_custom_hook_contracts.py create mode 100644 tests/test_file_session_identity.py create mode 100644 tests/test_mcp_protocol_boundaries.py create mode 100644 tests/test_mcp_result_integrity.py create mode 100644 tests/test_tool_workspace_boundaries.py create mode 100644 tests/test_usage_normalization.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d35876..b8ab869 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,74 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- HTTP MCP transports validate response IDs and object results; streamable HTTP + ignores unrelated SSE responses and adopts session headers only after a valid + response. Invalid JSON becomes an explicit MCP error. +- MCP pagination validates page shapes and opaque cursors, with a configurable + `max_discovery_pages` limit (100 by default) to bound endless discovery. +- Custom before-tool hooks validate successive argument rewrites, preserve + approval requirements, and pass rewritten arguments to permission checks. +- Post-tool output replacement clears stale model-facing excerpts; synchronous + hook dispatch rejects awaitables instead of silently ignoring async policies. +- MCP structured-only results are rendered as compact JSON instead of being + described as empty; existing text responses are not duplicated. Execution + status overrides static metadata, and only boolean `true` is a read-only hint. +- Malformed schema-cache roots/descriptors trigger live discovery instead of + crashing warm startup. +- Repeated-result suppression compares both raw evidence and model-facing + views, so a new relevant excerpt of unchanged evidence is not hidden. +- Shared core-tool workspaces reject prefix-sibling and symlink escapes using + path containment rather than string prefixes. +- Exact file edits reject empty `old_text`, preventing accidental insertion + between every character when `replace_all` is enabled. +- Session-manager chats use the manager's configured store without mutating + the original agent. Historical forks discard derived summaries/facts and + deep-copy retained messages so later evidence and edits cannot leak across. +- Closing agent/chat streams explicitly closes nested runtime iterators; + async cleanup is awaited instead of depending on generator finalization. +- Async streams propagate worker exceptions after draining emitted events; + provider failures no longer look like a successful end of stream. + +- OpenAI-compatible, LiteLLM (including streaming), and Anthropic usage + normalization no longer invents zero counters when providers omit input or + output usage. Budget enforcement can now recognize incomplete accounting. +- MCP resilience retries discovery/read requests, but no longer automatically + replays `tools/call` after ambiguous transport failures. A lost response may + follow a successful write; silently retrying could duplicate side effects. +- File session storage uses bounded, collision-resistant filenames, including + on case-insensitive filesystems. Legacy files remain readable only when their + embedded session identity matches; listing deduplicates migrated records. + +- A text tool call such as `{"path": …, "content": …}` fit several file tools + (each declares `path`) and was dropped as ambiguous, so the file was never + written. When several tools accept a nameless object, the one tool that + declares every key it carries is chosen; a genuine tie is still left alone. + +### Added + +- `recall_tool_result` accepts an optional case-insensitive literal `query`, + returning an exact bounded excerpt around a match instead of requiring + sequential paging through large historical results. +- `AgentChatSession.asend()`, `astream()`, and `astream_packets()` use the + native async runtime while retaining per-chat history, budgets, and callbacks. +- Regression coverage for custom hooks across all four agent entrypoints, + concurrent 20-turn scripted sessions, real local file-edit/test workflows, + provider usage normalization, and MCP protocol/result/cache boundaries. +- Hook integration guidance and a focused runtime reliability review documenting + implemented safeguards and remaining production-validation gaps. + +### Validation and limitations + +- Offline suite: **4,706 passed, 32 skipped**, with one pytest collection + warning (`TestCase` in `shipit_agent/deep/benchmark.py`). Live tests excluded. +- No live Bedrock quality or billed-token savings claim is made by this update. + Cross-worker durable budgets and production OAuth/MCP lifecycle validation + remain outstanding. Hook callbacks are synchronous, including in async runs; + post-tool hooks cannot retract output already streamed to clients. +- No Google Workspace extension is included in this update. + ## [2.2.0] - 2026-10-02 ### Fixed diff --git a/docs/agent-review-2026-10-06.md b/docs/agent-review-2026-10-06.md new file mode 100644 index 0000000..611bf27 --- /dev/null +++ b/docs/agent-review-2026-10-06.md @@ -0,0 +1,55 @@ +# Agent runtime review — 2026-10-06 + +## Scope and evidence + +This is a focused review of the main Agent runtime, sync/async streaming, +session management, output budgeting, MCP result handling and schema caching, +and core file tools. It is not a completed line-by-line audit of every module, +integration or optional dependency. Local regression suites and scripted-model +workflows do not establish production quality across all providers. + +## Research applied + +- Anthropic's [tool engineering guidance](https://www.anthropic.com/engineering/writing-tools-for-agents) + emphasizes relevant responses, discoverable tools, and evaluations measuring + correctness alongside calls, errors, latency and token consumption. Shipit + already has deferred discovery and bounded model-facing tool output; preserve + canonical evidence while improving those boundaries rather than indiscriminately + adding more overlapping tools. +- The MCP project's [annotation guidance](https://blog.modelcontextprotocol.io/posts/2026-03-16-tool-annotations/) + distinguishes behavioral hints from security guarantees. A read-only annotation + is not authorization or proof of safe retry. The host must enforce permissions. + +## Concrete findings addressed + +1. Structured-only MCP results were kept in metadata but described as empty in + model-visible text. They now render as compact JSON if no other text exists. +2. Static tool metadata could override current execution success/failure. Current + execution status now wins; non-boolean read-only hints are not accepted. +3. Result deduplication compared raw output but ignored changing semantic views. + Both must now match before a repeated result is suppressed. +4. Syntactically valid but malformed schema-cache data could crash warm discovery. + Invalid descriptors now produce a cache miss and live discovery. + +See [reliability hardening](reliability-hardening.md) for the preceding session, +stream lifecycle, token accounting, workspace, and targeted-recall changes. + +## Highest-value remaining work + +- **Durable budgets:** current session usage is in-memory. Company-wide or + cross-worker spend limits need durable accounting, atomic reservations and + reconciliation, not a prompt instruction or a per-agent counter. +- **Real-provider evaluations:** run held-out tasks on actual configured models, + compare success rates, uncached/cache input, output, retries, latency and cost. + An offline payload reduction is not a billed-token reduction. +- **Provider/cache behavior:** test prefix stability versus dynamic schemas and + compaction. Caching must be reported by the provider, not inferred from reuse. +- **Production MCP lifecycle:** validate authentication renewal, server-side + tool changes, tenant scoping, and reconnect behavior with authorized test + connectors. Annotation hints are not substitutes for these checks. +- **App-building evaluations:** expand beyond the tested small WSGI repair loop + to multi-file projects and real-model tasks with independent verification. + +More tools, more retries, or more subagents can increase cost. Optimization must +hold correctness and evidence availability constant, and measure total cost per +successful task rather than only the size of one request. diff --git a/docs/guides/context-management.md b/docs/guides/context-management.md index f354534..6550040 100644 --- a/docs/guides/context-management.md +++ b/docs/guides/context-management.md @@ -55,7 +55,9 @@ override these defaults. Optimized project agents also default to durable project-local stores: -- `.shipit/sessions/.json` keeps canonical multi-turn history. +- `.shipit/sessions/~.json` keeps canonical multi-turn history + without filename collisions. Use the session store API to resolve IDs; legacy + files remain readable after checking their embedded identity. - `.shipit/memory.json` keeps facts explicitly persisted by tools. Recreate the agent and reuse the same session ID to resume after a process diff --git a/docs/guides/hooks.md b/docs/guides/hooks.md index d7eac99..709ffb5 100644 --- a/docs/guides/hooks.md +++ b/docs/guides/hooks.md @@ -28,7 +28,7 @@ def log_tool_start(name, arguments): def log_tool_end(name, result): print(f"{name} returned {len(result.output)} chars") -agent = Agent.with_builtins( +agent = Agent( llm=OpenAIChatLLM(model="gpt-4o-mini"), hooks=hooks, ) @@ -44,6 +44,7 @@ result = agent.run("What is the weather in Tokyo?") | `after_llm` | `fn(response: LLMResponse)` | After each LLM completion returns | | `before_tool` | `fn(name: str, arguments: dict)` | Before a tool is executed | | `after_tool` | `fn(name: str, result: ToolResult)` | After a tool returns (success or error) | +| `user_prompt` | `fn(prompt: str)` | Rewrite or deny the incoming prompt before the model receives it | | `stop` | `fn(answer: str)` | When the agent is about to finish; return a reason to keep it working | ## Registration @@ -166,6 +167,59 @@ profile = ( `AgentHooks` works identically with `AsyncAgentRuntime`. The hook callbacks themselves are synchronous — the async runtime calls them inline between awaits. +This includes ordinary `Agent.run`, `stream`, `arun`, and `astream`, not only +the built-in-tools factory. Keep callbacks fast: slow network calls or sleeps +inside a callback block the async event loop. Returning an awaitable raises +`TypeError`; an `async def` policy must not be silently treated as approval. + +## Custom policies with ordinary Agent + +```python +from shipit_agent import Agent, AgentHooks + +hooks = AgentHooks() + +@hooks.on_user_prompt +def check_request(prompt): + if len(prompt) > 100_000: + return {"decision": "deny", "reason": "Request exceeds this application's limit."} + +@hooks.on_before_tool_matching("delete_*|send_*") +def require_review(name, arguments): + return {"decision": "ask", "reason": "This action needs review."} + +@hooks.on_after_llm +def record_usage(response): + usage_logger(response.usage) # your application-owned logging function + +agent = Agent(llm=llm, tools=tools, hooks=hooks) +``` + +Configure your application's approval handler when using `ask`; a hook does +not grant approval by itself. No provider-specific hook implementation is needed. + +Before-tool hooks run in registration order. Return a decision with +`updated_arguments` to rewrite the call. Later hooks and the permission engine +see those rewritten arguments. Deny wins over ask, and ask wins over allow; +later rewrites do not erase an earlier approval requirement. + +### Output transformations and boundaries + +Replacing output through a string, output/text dictionary, or `ToolOutput` +clears a stale `model_text` excerpt. A dictionary or `ToolOutput` can explicitly +supply a new `model_text` if a separate compact model view is desired. + +Post-tool hooks run **after execution**. They cannot undo side effects or retract +raw output deltas already streamed to subscribers. Do not rely on them alone to +keep secrets out of logs/UI: redact at the tool source or buffer/filter events +before forwarding them. Hook code is trusted application Python, not sandboxed. +Create separately scoped callback state for different users; cloned agents can +share callback objects. Model hooks are not a durable billing ledger. + +Ordinary callback exceptions propagate. Stop-hook exceptions are the documented +exception: they emit `stop_hook_error` and allow finishing. Stop hooks are quality +checks, not a security boundary, and their continuations can add token usage. + ## Match and transform tool output Matcher hooks accept glob patterns separated by `|`. Post-tool hooks may diff --git a/docs/reliability-hardening.md b/docs/reliability-hardening.md new file mode 100644 index 0000000..eed9d8d --- /dev/null +++ b/docs/reliability-hardening.md @@ -0,0 +1,135 @@ +# Runtime reliability hardening + +## Provider usage and budgets + +OpenAI-compatible, LiteLLM, and Anthropic adapters preserve missing input/output +token counters as unknown. An explicit zero is valid; a missing, negative, +boolean, or non-integer value is not treated as free work. A reported total alone +cannot determine the input/output split or cache semantics. + +Configured task/session budgets use this distinction to stop subsequent model +requests when accounting is unavailable. These remain soft, between-step limits, +not a guarantee that an in-flight request cannot exceed the limit. Cache savings +must be measured from provider-reported counters; these changes do not establish +a live cost reduction. + +## MCP replay safety + +The resilience wrapper retries allowlisted protocol discovery/read requests. +It does not replay `tools/call` or unknown methods automatically. A timeout or +disconnected socket does not prove the server failed to perform an action. +Applications needing retries for writes should use server-supported idempotency +or reconcile execution status first. This transport policy does not itself stop +a model from requesting the action again in a later step. + +## File session identity + +New session files use SHA-256-derived filenames, avoiding separator, filename +length, and case-insensitive filesystem collisions. Existing files can still be +read after verifying their embedded session ID. Subsequent saves use the new +filename; listings deduplicate legacy/current copies. Old files are not deleted. + +This prevents filename collisions, not unauthorized access: the host application +must still authenticate users, authorize session access, and provide process-safe +coordination for multiple workers. In-process session locks are not distributed +locks. Do not treat the file store as a multi-tenant security boundary. + +## Async chat lifecycle and forks + +`AgentChatSession.asend`, `astream`, and `astream_packets` share the same chat +state as synchronous sends. Event/packet callbacks remain synchronous and should +be short and non-blocking. When abandoning a stream, explicitly close it (for +example with `contextlib.aclosing`); merely breaking an async loop is not a +guarantee of immediate cleanup. Closing these wrappers now awaits nested runtime +cleanup and honors the configured runtime cancellation policy. +Worker exceptions propagate after buffered events are drained; they no longer +silently end an async stream as though it succeeded. + +```python +from contextlib import aclosing + +chat = agent.chat_session(session_id=authorized_session_id) +result = await chat.asend("Review the evidence") +async with aclosing(chat.astream_packets("Compare with the earlier result", transport="sse")) as packets: + async for packet in packets: + await send_to_client(packet) +``` + +The host must authenticate and authorize `authorized_session_id` before creating +the chat. A disconnected client should close the iterator. Do not convert an +exception from the stream into a completed assistant message. + +`SessionManager` binds returned chats to its own store without changing the +supplied agent. Forks deep-copy messages and discard derived facts, compaction +summaries, and discovery checkpoints; these can otherwise reference messages +after a historical fork point. Custom metadata is copied independently. + +Session token budgets still live in chat/agent runtime state and are not a +durable company billing ledger. Reconstructing the agent can reset that budget; +production cross-worker spending limits need durable accounting and coordination. + +## Validation boundaries + +### Targeted evidence recall and app-building tools + +Historical `recall_tool_result` accepts `query` alongside `call_id`, `offset`, +and `limit`. It searches literal text case-insensitively (not model-generated +regex), returning exact source text around the first match at or after offset. +Use a small `limit` when only a fact or error is needed. No-match responses are +explicit; normal paging remains available. The source result is unchanged. + +The regression fixture retrieves a relevant excerpt under 750 characters from +over 400,000 characters of historical output. This is returned-text reduction, +not a claim about real provider token billing or task success. + +Core workspace path resolution rejects similarly named sibling directories and +symlink escapes. This is a file-tool boundary, not an OS sandbox for subprocesses. +Exact edits reject empty search text, including `replace_all` requests. + +An app-building regression uses real file tools and Python subprocesses to write +a WSGI greeting application, observe a failing assertion, read/edit the source, +and verify the corrected response. It runs through both `Agent.run` and +`Agent.stream`; model choices are scripted. It does not install dependencies, +start a public server, or establish autonomous real-model coding quality. + +Regression coverage includes missing usage, cache-only usage, sync/stream budget +enforcement, ambiguous MCP writes, and colliding session IDs. Existing suites +cover large-catalog discovery, HTTP failure isolation, stream events, routing, +and session behavior. Fixture tests do not establish live Gemma/Bedrock quality, +production connector reliability, or comparative savings against another agent. + +The concurrent chat regression runs two isolated twenty-turn chats, alternating +async sends and streams, with forty unique local tool executions. It checks prior +user turns on every model call, exact terminal output, retained tool results, and +separate usage totals. Its model and usage counts are scripted fixtures, not a +live quality or cost benchmark. + +### Offline discovery measurement + +`PYTHONPATH=. .venv/bin/python scripts/benchmark_tool_discovery.py` exercised +twenty successful fixture lookups per configuration. Automatic discovery versus +eager loading reduced serialized request characters by 10.7% with 30 tools, +56.6% with 100 tools, and 89.1% with 500 tools. With 10 tools it retained eager +loading (no reduction). Discovery added one model request in each deferred case. +These measure existing discovery behavior, not savings introduced by this patch. +They are payload measurements, not tokenizer counts, cache hits, billed cost, or +real-model task-success measurements. +## MCP response and pagination boundaries + +HTTP MCP responses must match the outgoing request ID and contain an object +result. Streamable HTTP skips unrelated SSE responses and accepts session +headers only after response validation. This prevents unrelated content from +being treated as tool evidence; it does not change SSE into incremental reading. + +`RemoteMCPServer(..., max_discovery_pages=100)` bounds each paginated listing. +Increase the positive integer limit for servers with unusually small pages. +Cursors remain opaque and are forwarded unchanged, following the +[MCP pagination specification](https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/pagination). +Malformed pages, repeated cursors, and exhausted page budgets raise `MCPError` +inside discovery. Failed tool discovery does not publish a partial catalog. +The existing optional resources/prompts listing APIs still translate MCP errors +to empty lists; they do not yet distinguish unsupported methods from failures. + +Regression coverage includes HTTP/SSE response identity, malformed payloads, +session-header validation, cursor preservation, and bounded discovery. These +are local transport fixtures, not live-provider or production-server validation. diff --git a/notebooks/77_super_agent_live_validation.ipynb b/notebooks/77_super_agent_live_validation.ipynb index bdfd12c..9f11471 100644 --- a/notebooks/77_super_agent_live_validation.ipynb +++ b/notebooks/77_super_agent_live_validation.ipynb @@ -203,7 +203,7 @@ "first.chat_session(session_id=\"main\").send(\"remember this checkpoint\")\n", "second = Agent.for_project(llm=SimpleEchoLLM(), project_root=workspace, optimized=True)\n", "assert second.chat_session(session_id=\"main\").history()\n", - "assert (workspace / \".shipit\" / \"sessions\" / \"main.json\").exists()\n", + "assert second.session_store.load(\"main\") is not None\n", "\n", "# Screenshot/media references preserve text-image-text ordering.\n", "parsed = MediaParser(allowlist_domains=[\"cdn.example.com\"]).parse(\n", diff --git a/shipit_agent/agent.py b/shipit_agent/agent.py index b0dd022..9f97685 100644 --- a/shipit_agent/agent.py +++ b/shipit_agent/agent.py @@ -52,6 +52,8 @@ from __future__ import annotations +from contextlib import aclosing, closing + import asyncio import logging import threading @@ -1341,8 +1343,9 @@ def stream( completed = False try: with self._run_lock: - for event in runtime.stream(user_prompt, user_content=user_content): - yield event + with closing(runtime.stream(user_prompt, user_content=user_content)) as events: + for event in events: + yield event completed = True finally: self._active_runtime = None @@ -1395,10 +1398,11 @@ async def astream( sources: list[Any] = [] completed = False try: - async for event in runtime.stream( + async with aclosing(runtime.stream( user_prompt, user_content=user_content - ): - yield event + )) as events: + async for event in events: + yield event completed = True finally: self._active_runtime = None diff --git a/shipit_agent/async_runtime.py b/shipit_agent/async_runtime.py index fcb07cd..45ed33c 100644 --- a/shipit_agent/async_runtime.py +++ b/shipit_agent/async_runtime.py @@ -1995,6 +1995,9 @@ async def _worker() -> None: await wake.wait() continue yield event_buffer.popleft() + # A failed worker is not a successful end-of-stream. Drain its + # emitted events, then propagate the original exception to callers. + await task finally: self._event_subscriber = None if not task.done() and self.cancel_on_stream_close: diff --git a/shipit_agent/chat_session.py b/shipit_agent/chat_session.py index ac6e6df..a111774 100644 --- a/shipit_agent/chat_session.py +++ b/shipit_agent/chat_session.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from contextlib import aclosing, closing from typing import TYPE_CHECKING, Any, Callable from shipit_agent.models import AgentEvent, AgentResult, Message @@ -67,17 +68,36 @@ def send(self, user_prompt: str) -> AgentResult: return result def stream(self, user_prompt: str): - for event in self._session_agent().stream(user_prompt): + with closing(self._session_agent().stream(user_prompt)) as events: + for event in events: + self._emit_event(event) + yield event + + async def asend(self, user_prompt: str) -> AgentResult: + """Run a turn through the native async runtime with shared chat state.""" + result = await self._session_agent().arun(user_prompt) + for event in result.events: self._emit_event(event) - yield event + return result + + async def astream(self, user_prompt: str): + """Stream a turn; closing this iterator also closes its producer.""" + async with aclosing(self._session_agent().astream(user_prompt)) as events: + async for event in events: + self._emit_event(event) + yield event def stream_packets(self, user_prompt: str, *, transport: str = "websocket"): - if transport == "sse": - for event in self.stream(user_prompt): - yield sse_event_packet(event) - return - for event in self.stream(user_prompt): - yield websocket_event_packet(event) + packet = sse_event_packet if transport == "sse" else websocket_event_packet + with closing(self.stream(user_prompt)) as events: + for event in events: + yield packet(event) + + async def astream_packets(self, user_prompt: str, *, transport: str = "websocket"): + packet = sse_event_packet if transport == "sse" else websocket_event_packet + async with aclosing(self.astream(user_prompt)) as events: + async for event in events: + yield packet(event) def send_result_packet( self, user_prompt: str, *, transport: str = "websocket" diff --git a/shipit_agent/hooks.py b/shipit_agent/hooks.py index f61c742..195dc36 100644 --- a/shipit_agent/hooks.py +++ b/shipit_agent/hooks.py @@ -2,11 +2,22 @@ from dataclasses import dataclass, field import fnmatch +import inspect from typing import Any, Callable from shipit_agent.permissions import PermissionDecision, PermissionResult +def _invoke(fn: Callable[..., Any], *args: Any) -> Any: + """Never silently treat an unawaited async policy as approval.""" + value = fn(*args) + if inspect.isawaitable(value): + if inspect.iscoroutine(value): + value.close() + raise TypeError("AgentHooks callbacks must be synchronous; async callbacks are not supported") + return value + + def _coerce_hook_decision(value: Any) -> PermissionResult | None: """Normalize whatever a before-tool hook returned into a PermissionResult. @@ -141,11 +152,11 @@ def on_stop(self, fn: Callable[..., Any]) -> Callable[..., Any]: # ------------------------------------------------------------------ def run_before_llm(self, messages: list[Any], tools: list[Any]) -> None: for fn in self.before_llm: - fn(messages, tools) + _invoke(fn, messages, tools) def run_after_llm(self, response: Any) -> None: for fn in self.after_llm: - fn(response) + _invoke(fn, response) def run_before_tool( self, name: str, arguments: dict[str, Any] @@ -157,37 +168,51 @@ def run_before_tool( when every hook is observe-only (backward compatible). """ outcome: PermissionResult | None = None + current = dict(arguments) + rewritten = False for fn in self.before_tool: - decision = _coerce_hook_decision(fn(name, arguments)) + decision = _coerce_hook_decision(_invoke(fn, name, current)) if decision is None: continue if decision.denied: return decision # short-circuit on the first hard deny + if decision.updated_arguments is not None: + if not isinstance(decision.updated_arguments, dict): + raise TypeError("updated_arguments must be a dictionary") + current = dict(decision.updated_arguments) + rewritten = True if decision.needs_approval: outcome = decision elif outcome is None or not outcome.needs_approval: # ALLOW (possibly with a rewrite) — keep unless an ASK stands. outcome = decision + if outcome is not None and rewritten: + outcome = PermissionResult(outcome.decision, outcome.reason, current) return outcome def run_after_tool(self, name: str, result: Any) -> Any: """Run post-tool hooks, allowing each to replace or sanitize output.""" current = result for fn in self.after_tool: - value = fn(name, current) + value = _invoke(fn, name, current) if value is None: continue if isinstance(value, str) and hasattr(current, "output"): current.output = value + current.model_text = None continue if isinstance(value, dict) and hasattr(current, "output"): if "output" in value or "text" in value: current.output = str(value.get("output", value.get("text", ""))) + current.model_text = None + if "model_text" in value: + current.model_text = value["model_text"] if isinstance(value.get("metadata"), dict): current.metadata.update(value["metadata"]) continue if hasattr(value, "text") and hasattr(current, "output"): current.output = str(value.text) + current.model_text = getattr(value, "model_text", None) current.metadata.update(dict(getattr(value, "metadata", {}) or {})) continue if hasattr(value, "output"): @@ -197,7 +222,7 @@ def run_after_tool(self, name: str, result: Any) -> Any: def run_stop(self, answer: str) -> str | None: """The first reason a stop hook gives to keep going, or ``None``.""" for fn in self.stop: - value = fn(answer) + value = _invoke(fn, answer) if isinstance(value, str) and value.strip(): return value.strip() if isinstance(value, dict) and str(value.get("decision", "")).lower() in ( @@ -216,7 +241,7 @@ def run_user_prompt(self, prompt: str) -> tuple[str, PermissionResult | None]: """ current = prompt for fn in self.user_prompt: - value = fn(current) + value = _invoke(fn, current) if isinstance(value, str): current = value continue diff --git a/shipit_agent/llms/anthropic_adapter.py b/shipit_agent/llms/anthropic_adapter.py index 96e49d0..6c21bdc 100644 --- a/shipit_agent/llms/anthropic_adapter.py +++ b/shipit_agent/llms/anthropic_adapter.py @@ -7,6 +7,7 @@ from shipit_agent.llms import citations as _citations from shipit_agent.llms import server_tools as _server_tools from shipit_agent.llms.base import LLMResponse, coerce_message +from shipit_agent.llms.usage import normalize_token_usage from shipit_agent.models import Message, ToolCall # Beta header for interleaved thinking (verified in @@ -566,11 +567,9 @@ def complete( usage: dict[str, int] = {} if hasattr(response, "usage") and response.usage: - usage = { - "prompt_tokens": getattr(response.usage, "input_tokens", 0) or 0, - "completion_tokens": getattr(response.usage, "output_tokens", 0) or 0, - } - usage["total_tokens"] = usage["prompt_tokens"] + usage["completion_tokens"] + usage = normalize_token_usage( + response.usage, input_key="input_tokens", output_key="output_tokens" + ) # Prompt-caching usage. These keys match exactly what # CostTracker.as_hooks() looks for, so cache reads bill at the # cheaper rate automatically. Wrapped defensively: older SDKs diff --git a/shipit_agent/llms/litellm_adapter.py b/shipit_agent/llms/litellm_adapter.py index d4e99b0..58105b8 100644 --- a/shipit_agent/llms/litellm_adapter.py +++ b/shipit_agent/llms/litellm_adapter.py @@ -8,6 +8,7 @@ from typing import Any, Callable from shipit_agent.llms.base import LLMResponse +from shipit_agent.llms.usage import normalize_token_usage from shipit_agent.llms.bedrock_token import ( BedrockTokenError, existing_bearer_token, @@ -499,11 +500,7 @@ def complete( usage: dict[str, int] = {} if hasattr(response, "usage") and response.usage: u = response.usage - usage = { - "prompt_tokens": getattr(u, "prompt_tokens", 0) or 0, - "completion_tokens": getattr(u, "completion_tokens", 0) or 0, - "total_tokens": getattr(u, "total_tokens", 0) or 0, - } + usage = normalize_token_usage(u) usage.update(_extract_cache_usage(u)) return LLMResponse( @@ -677,11 +674,7 @@ def _stream_completion( chunk_usage = getattr(chunk, "usage", None) if chunk_usage: - usage = { - "prompt_tokens": getattr(chunk_usage, "prompt_tokens", 0) or 0, - "completion_tokens": getattr(chunk_usage, "completion_tokens", 0) or 0, - "total_tokens": getattr(chunk_usage, "total_tokens", 0) or 0, - } + usage = normalize_token_usage(chunk_usage) usage.update(_extract_cache_usage(chunk_usage)) except Exception as exc: exc_name = type(exc).__name__ diff --git a/shipit_agent/llms/openai_adapter.py b/shipit_agent/llms/openai_adapter.py index c5ff2d9..65fdd4b 100644 --- a/shipit_agent/llms/openai_adapter.py +++ b/shipit_agent/llms/openai_adapter.py @@ -4,6 +4,7 @@ from typing import Any from shipit_agent.llms.base import LLMResponse +from shipit_agent.llms.usage import normalize_token_usage from shipit_agent.models import Message, ToolCall from shipit_agent.llms.litellm_adapter import ( _extract_reasoning, @@ -92,11 +93,7 @@ def _usage_dict(usage_obj: Any) -> dict[str, int]: """ if not usage_obj: return {} - out: dict[str, int] = { - "prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0, - "completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0, - "total_tokens": getattr(usage_obj, "total_tokens", 0) or 0, - } + out: dict[str, int] = normalize_token_usage(usage_obj) try: details = getattr(usage_obj, "prompt_tokens_details", None) cached = getattr(details, "cached_tokens", None) if details else None diff --git a/shipit_agent/llms/usage.py b/shipit_agent/llms/usage.py index 23dea21..76f2fe4 100644 --- a/shipit_agent/llms/usage.py +++ b/shipit_agent/llms/usage.py @@ -1,6 +1,39 @@ """Normalize cache accounting without changing provider-reported counters.""" from __future__ import annotations +from collections.abc import Mapping +from typing import Any + + +def normalize_token_usage( + usage: Any, + *, + input_key: str = "prompt_tokens", + output_key: str = "completion_tokens", +) -> dict[str, int]: + """Keep valid provider counters; missing usage is unknown, never zero. + + Accept SDK objects and mapping responses. Do not infer either input or + output from totals: cache semantics differ between providers. + """ + result: dict[str, int] = {} + for target, source in ( + ("prompt_tokens", input_key), + ("completion_tokens", output_key), + ("total_tokens", "total_tokens"), + ): + value = ( + usage.get(source) if isinstance(usage, Mapping) + else getattr(usage, source, None) + ) + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: + result[target] = value + if "total_tokens" not in result and all( + key in result for key in ("prompt_tokens", "completion_tokens") + ): + result["total_tokens"] = result["prompt_tokens"] + result["completion_tokens"] + return result + def has_complete_token_usage(response) -> bool: """Whether normalized input/output counters can enforce a task budget. diff --git a/shipit_agent/mcp.py b/shipit_agent/mcp.py index 6ad1a69..619b111 100644 --- a/shipit_agent/mcp.py +++ b/shipit_agent/mcp.py @@ -55,6 +55,20 @@ class MCPError(RuntimeError): pass +def _http_result(parsed: Any, request_id: int) -> dict[str, Any]: + if ( + not isinstance(parsed, dict) + or type(parsed.get("id")) is not int + or parsed["id"] != request_id + ): + raise MCPError("MCP response id does not match the request") + if "error" in parsed: + raise MCPError(str(parsed["error"])) + if not isinstance(parsed.get("result"), dict): + raise MCPError("MCP response must contain an object result") + return parsed["result"] + + def _mcp_result_summary(result: dict[str, Any], rendered: str) -> str: """Use only a summary the MCP server explicitly supplied. @@ -131,7 +145,7 @@ class MCPRemoteTool: @property def read_only(self) -> bool: - return bool(self.annotations.get("readOnlyHint", False)) + return self.annotations.get("readOnlyHint") is True prompt: str = "Use this MCP tool when the remote server provides the best capability for the task." prompt_instructions: str = ( @@ -181,16 +195,21 @@ def run(self, context: ToolContext, **kwargs: Any) -> ToolOutput: f"failed: {exc}" ), metadata={ + **self.metadata, "server": self.server_name, "ok": False, + "is_error": True, "error": str(exc), - **self.metadata, }, ) content = list(result.get("content", [])) text_parts = _render_mcp_content(content) is_error = bool(result.get("isError", False)) rendered = "\n".join(part for part in text_parts if part).strip() + if not rendered and isinstance(result.get("structuredContent"), dict): + # Structured-only results are valid evidence, not empty responses. + # Do not duplicate the JSON when the server already supplied text. + rendered = json.dumps(result["structuredContent"], ensure_ascii=False, separators=(",", ":")) if not rendered: rendered = ( f"MCP tool '{self.name}' returned an error without details." @@ -223,10 +242,12 @@ def run(self, context: ToolContext, **kwargs: Any) -> ToolOutput: return ToolOutput( text=rendered, metadata={ + **self.metadata, **vision_metadata, "server": self.server_name, "ok": not is_error, "is_error": is_error, + "error": rendered if is_error else None, "structured_content": result.get("structuredContent"), "content_blocks": content, "output_schema": dict(self.output_schema), @@ -238,7 +259,6 @@ def run(self, context: ToolContext, **kwargs: Any) -> ToolOutput: else {} ), **({"summary": summary} if summary else {}), - **self.metadata, }, ) @@ -497,10 +517,11 @@ def __init__( def request( self, method: str, params: dict[str, Any] | None = None ) -> dict[str, Any]: + request_id = next(self._id_counter) payload = json.dumps( { "jsonrpc": "2.0", - "id": next(self._id_counter), + "id": request_id, "method": method, "params": params or {}, } @@ -513,10 +534,11 @@ def request( ) with request.urlopen(req, timeout=self.timeout) as response: # nosec B310 body = response.read().decode("utf-8") - parsed = json.loads(body) if body else {} - if "error" in parsed: - raise MCPError(str(parsed["error"])) - return dict(parsed.get("result", {})) + try: + parsed = json.loads(body) + except (ValueError, TypeError) as exc: + raise MCPError("Invalid JSON in MCP response") from exc + return _http_result(parsed, request_id) def close(self) -> None: return None @@ -556,10 +578,11 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: def request( self, method: str, params: dict[str, Any] | None = None ) -> dict[str, Any]: + request_id = next(self._id_counter) payload = json.dumps( { "jsonrpc": "2.0", - "id": next(self._id_counter), + "id": request_id, "method": method, "params": params or {}, } @@ -576,17 +599,19 @@ def request( ) with request.urlopen(req, timeout=self.timeout) as response: # nosec B310 session_id = response.headers.get("Mcp-Session-Id") - if session_id: - self._session_id = session_id content_type = (response.headers.get("Content-Type") or "").lower() body = response.read().decode("utf-8") if "text/event-stream" in content_type: - parsed = self._parse_sse(body) + parsed = self._parse_sse(body, request_id=request_id) else: - parsed = json.loads(body) if body else {} - if "error" in parsed: - raise MCPError(str(parsed["error"])) - return dict(parsed.get("result", {})) + try: + parsed = json.loads(body) + except (ValueError, TypeError) as exc: + raise MCPError("Invalid JSON in MCP response") from exc + result = _http_result(parsed, request_id) + if session_id: + self._session_id = session_id + return result def notify(self, method: str, params: dict[str, Any] | None = None) -> None: """Send a session-affine JSON-RPC notification.""" @@ -628,7 +653,7 @@ def close(self) -> None: return @staticmethod - def _parse_sse(body: str) -> dict[str, Any]: + def _parse_sse(body: str, *, request_id: int | None = None) -> dict[str, Any]: """Extract the JSON-RPC response from SSE `data:` lines.""" # HTTP servers commonly use CRLF framing. Splitting the raw body on # ``\n\n`` merges every CRLF event into one invalid JSON blob, which @@ -648,6 +673,10 @@ def _parse_sse(body: str) -> dict[str, Any]: continue # The response to our request carries an id (notifications don't). if isinstance(parsed, dict) and ("result" in parsed or "error" in parsed): + if request_id is not None and ( + type(parsed.get("id")) is not int or parsed["id"] != request_id + ): + continue return parsed raise MCPError("No JSON-RPC response found in SSE stream.") @@ -698,6 +727,8 @@ class RemoteMCPServer(MCPServer): cache: bool = False #: Warm-start freshness window (seconds). A cache older than this is a miss. cache_ttl: float = 24 * 60 * 60 + #: Bound pagination without silently treating a partial catalog as complete. + max_discovery_pages: int = 100 _discovered: bool = False _initialized: bool = False _handshook_generation: int = 0 @@ -752,17 +783,24 @@ def _paginated_items(self, method: str, key: str) -> list[dict[str, Any]]: """Collect every MCP cursor page and reject looping server cursors.""" if self.transport is None: raise MCPError("RemoteMCPServer requires a transport.") + if type(self.max_discovery_pages) is not int or self.max_discovery_pages < 1: + raise ValueError("max_discovery_pages must be a positive integer") items: list[dict[str, Any]] = [] cursor: str | None = None seen: set[str] = set() - while True: + for _ in range(self.max_discovery_pages): params = {"cursor": cursor} if cursor is not None else {} result = self.transport.request(method, params) - items.extend(item for item in result.get(key, []) if isinstance(item, dict)) + if not isinstance(result, dict) or not isinstance(result.get(key), list): + raise MCPError(f"Invalid {method} response: expected a {key} list") + if any(not isinstance(item, dict) for item in result[key]): + raise MCPError(f"Invalid {method} response: list items must be objects") + items.extend(result[key]) next_cursor = result.get("nextCursor") if next_cursor in (None, ""): return items - next_cursor = str(next_cursor) + if not isinstance(next_cursor, str): + raise MCPError(f"Invalid {method} response: nextCursor must be a string") if next_cursor in seen: raise MCPError( f"MCP server '{self.name}' repeated cursor {next_cursor!r} " @@ -770,6 +808,10 @@ def _paginated_items(self, method: str, key: str) -> list[dict[str, Any]]: ) seen.add(next_cursor) cursor = next_cursor + raise MCPError( + f"MCP server '{self.name}' exceeded " + f"max_discovery_pages={self.max_discovery_pages} during {method}" + ) def invalidate_tools_cache(self) -> None: """Force live tool discovery on the next registry construction.""" diff --git a/shipit_agent/mcp_resilience.py b/shipit_agent/mcp_resilience.py index 9794b89..d59abca 100644 --- a/shipit_agent/mcp_resilience.py +++ b/shipit_agent/mcp_resilience.py @@ -16,9 +16,9 @@ special-cased to a long, fixed cooldown — hammering a server that just told us to slow down is how a soft limit becomes a ban. * **Bounded retry** — a single transient failure (a subprocess that died, a - dropped connection) is retried once after a short jittered delay, because the - persistent transports reconnect on the next call. Rate-limit errors are never - retried. + dropped connection) during discovery is retried once after a short jittered + delay. Tool execution is not replayed: a lost response does not prove that + the server did not already perform the action. Rate limits are never retried. ``ResilientMCPTransport`` wraps anything with ``.request(method, params)`` / ``.close()`` — the transport Protocol — so it composes with every transport @@ -50,6 +50,13 @@ "quota exceeded", ) +# Protocol discovery/read operations only. In particular, never automatically +# replay tools/call, even after a timeout: its side effects may have committed. +_RETRY_SAFE_METHODS = frozenset({ + "ping", "tools/list", "resources/list", "resources/templates/list", + "resources/read", "prompts/list", "prompts/get", +}) + def is_rate_limited(exc: BaseException) -> bool: """Best-effort: does this error mean 'you are being rate limited'?""" @@ -156,8 +163,12 @@ def request( return result except (MCPError, OSError, TimeoutError) as exc: rate_limited = is_rate_limited(exc) - # Never retry a rate-limit; back off long and re-raise. - if rate_limited or attempt >= self.max_retries: + # Never replay ambiguous actions or retry a rate-limit. + if ( + rate_limited + or method not in _RETRY_SAFE_METHODS + or attempt >= self.max_retries + ): self.breaker.record_failure(rate_limited=rate_limited) raise attempt += 1 diff --git a/shipit_agent/mcp_schema_cache.py b/shipit_agent/mcp_schema_cache.py index d2b755a..12deac1 100644 --- a/shipit_agent/mcp_schema_cache.py +++ b/shipit_agent/mcp_schema_cache.py @@ -144,10 +144,18 @@ def load( return None try: payload = json.loads(raw) + if not isinstance(payload, dict): + return None if payload.get("schema_version") != SCHEMA_VERSION: return None if not isinstance(payload.get("tools"), list): return None + for descriptor in payload["tools"]: + if not isinstance(descriptor, dict) or not isinstance(descriptor.get("name"), str) or not descriptor["name"]: + return None + if any(key in descriptor and not isinstance(descriptor[key], dict) + for key in ("input_schema", "output_schema", "annotations", "execution")): + return None age = time.time() - float(payload.get("saved_at", 0)) if age < 0 or age > ttl: return None diff --git a/shipit_agent/permissions.py b/shipit_agent/permissions.py index e29c830..aa3720c 100644 --- a/shipit_agent/permissions.py +++ b/shipit_agent/permissions.py @@ -273,24 +273,27 @@ def authorize_tool( import. """ results: list[PermissionResult] = [] + effective_arguments = arguments if hooks is not None: hook_decision = hooks.run_before_tool(name, arguments) if hook_decision is not None: results.append(hook_decision) + if hook_decision.updated_arguments is not None: + effective_arguments = hook_decision.updated_arguments if permissions is not None: - results.append(permissions.check(name, arguments, tool)) + results.append(permissions.check(name, effective_arguments, tool)) if not results: return None for result in results: if result.denied: return result - for result in results: - if result.needs_approval: - return result updated: dict[str, Any] | None = None for result in results: if result.updated_arguments is not None: updated = result.updated_arguments + for result in results: + if result.needs_approval: + return PermissionResult(result.decision, result.reason, updated) return PermissionResult(PermissionDecision.ALLOW, updated_arguments=updated) diff --git a/shipit_agent/runtime_core.py b/shipit_agent/runtime_core.py index 9ba6218..253a15b 100644 --- a/shipit_agent/runtime_core.py +++ b/shipit_agent/runtime_core.py @@ -1124,7 +1124,11 @@ def model_visible_tool_output( if seen is not None and len(output) >= _REPEAT_MIN_CHARS: key = (name, _arguments_key(arguments)) - if seen.get(key) == canonical_digest: + # A tool can select a new relevant excerpt from unchanged raw + # evidence. Suppress only if both the evidence and view match. + view_digest = hashlib.sha256(output.encode("utf-8", "replace")).hexdigest() + identity = canonical_digest + ":" + view_digest + if seen.get(key) == identity: metadata = getattr(tool_result, "metadata", None) if isinstance(metadata, dict): metadata.update( @@ -1148,7 +1152,7 @@ def model_visible_tool_output( else "]" ) ) - seen[key] = canonical_digest + seen[key] = identity if has_semantic_output: metadata = getattr(tool_result, "metadata", None) diff --git a/shipit_agent/session_manager.py b/shipit_agent/session_manager.py index 9aa84b0..c06d5f7 100644 --- a/shipit_agent/session_manager.py +++ b/shipit_agent/session_manager.py @@ -1,6 +1,7 @@ from __future__ import annotations import uuid +from copy import deepcopy from dataclasses import dataclass from typing import Any @@ -56,9 +57,13 @@ def fork(self, agent: Any, session_id: str, *, from_message: int = -1) -> Any: fork_name = record.metadata.get("name", session_id) new_record = SessionRecord( session_id=new_id, - messages=messages, + messages=deepcopy(messages), metadata={ - **record.metadata, + # Derived state may describe turns *after* the fork point. + # Rebuild it from the fork's retained transcript instead. + **deepcopy({key: value for key, value in record.metadata.items() + if key not in {"compaction_checkpoint", "verified_facts", + "tool_discovery"}}), "name": f"Fork of {fork_name}", "forked_from": session_id, "forked_at_message": from_message, @@ -67,8 +72,11 @@ def fork(self, agent: Any, session_id: str, *, from_message: int = -1) -> Any: self.session_store.save(new_record) return self._chat_session(agent, new_id) - @staticmethod - def _chat_session(agent: Any, session_id: str) -> Any: + def _chat_session(self, agent: Any, session_id: str) -> Any: + if getattr(agent, "session_store", None) is not self.session_store: + if not callable(getattr(agent, "clone", None)): + raise ValueError("Agent must use the manager's session store or support clone()") + agent = agent.clone(session_store=self.session_store) try: return agent.chat_session(session_id=session_id) except TypeError: diff --git a/shipit_agent/stores/session.py b/shipit_agent/stores/session.py index 3983539..75795f2 100644 --- a/shipit_agent/stores/session.py +++ b/shipit_agent/stores/session.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import json import os import tempfile @@ -45,14 +46,29 @@ def __init__(self, root_dir: str | Path) -> None: self.root_dir.mkdir(parents=True, exist_ok=True) def _path_for(self, session_id: str) -> Path: - safe = session_id.replace("/", "_") - return self.root_dir / f"{safe}.json" + # Hash every ID: separator replacement collides ("a/b" == "a_b"), + # and even simple IDs collide on case-insensitive filesystems. + name = "~" + hashlib.sha256(session_id.encode("utf-8")).hexdigest() + return self.root_dir / f"{name}.json" def load(self, session_id: str) -> SessionRecord | None: path = self._path_for(session_id) if not path.exists(): - return None + # Read legacy separator-normalized files only when their embedded + # identity matches. Never return another chat after a collision. + legacy_name = session_id.replace("/", "_") + if "\\" in legacy_name or len(legacy_name) > 128: + return None + path = self.root_dir / f"{legacy_name}.json" + if not path.exists(): + return None raw = json.loads(path.read_text(encoding="utf-8")) + if raw["session_id"] != session_id: + return None + return self._decode(raw) + + @staticmethod + def _decode(raw: dict) -> SessionRecord: return SessionRecord( session_id=raw["session_id"], messages=[ @@ -73,10 +89,15 @@ def save(self, record: SessionRecord) -> None: def list_all(self) -> list[SessionRecord]: records: list[SessionRecord] = [] + seen: set[str] = set() for path in sorted(self.root_dir.glob("*.json")): - record = self.load(path.stem) - if record is not None: + raw = json.loads(path.read_text(encoding="utf-8")) + # Resolve by stored identity, not encoded filename; prefer the + # current file if both a legacy and migrated copy exist. + record = self.load(raw["session_id"]) + if record is not None and record.session_id not in seen: records.append(record) + seen.add(record.session_id) return records diff --git a/shipit_agent/tools/_shared.py b/shipit_agent/tools/_shared.py index f070102..b45a803 100644 --- a/shipit_agent/tools/_shared.py +++ b/shipit_agent/tools/_shared.py @@ -69,7 +69,7 @@ def resolve(self, path: str | Path) -> Path: walk out — checking the string first catches neither. """ target = (self.root / Path(path).expanduser()).resolve() - if not str(target).startswith(str(self.root)): + if not target.is_relative_to(self.root): raise ValueError(f"{path} is outside the workspace ({self.root})") return target diff --git a/shipit_agent/tools/edit_file/edit_file_tool.py b/shipit_agent/tools/edit_file/edit_file_tool.py index e740ef5..ebff8dc 100644 --- a/shipit_agent/tools/edit_file/edit_file_tool.py +++ b/shipit_agent/tools/edit_file/edit_file_tool.py @@ -118,6 +118,11 @@ def run(self, context: ToolContext, **kwargs) -> ToolOutput: old_text = str(kwargs.get("old_text", "")) new_text = str(kwargs.get("new_text", "")) + if not old_text: + return ToolOutput( + text="Edit failed: old_text must be a non-empty exact block from the file.", + metadata={"is_error": True, "error": "empty_old_text", "path": str(path)}, + ) replace_all = bool(kwargs.get("replace_all", False)) # Read raw bytes and decode strictly. Reading with errors="replace" diff --git a/shipit_agent/tools/recall_result/recall_result_tool.py b/shipit_agent/tools/recall_result/recall_result_tool.py index f1242fa..f20bdc7 100644 --- a/shipit_agent/tools/recall_result/recall_result_tool.py +++ b/shipit_agent/tools/recall_result/recall_result_tool.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass +import re from typing import Any, Mapping from shipit_agent.tools.base import ToolOutput @@ -27,7 +28,8 @@ class RecallToolResult: name = "recall_tool_result" description = ( "Recall an exact bounded slice of a large tool result from an earlier " - "turn using the call_id shown in conversation history." + "turn using the call_id shown in conversation history. Optionally search " + "for literal text to retrieve relevant evidence without paging the whole result." ) prompt_instructions = ( "Use only when the prior assistant answer is insufficient and exact " @@ -57,6 +59,10 @@ def schema(self) -> dict[str, Any]: "default": 0, "description": "Character offset into the result.", }, + "query": { + "type": "string", + "description": "Optional case-insensitive literal text to find at or after offset. Not a regular expression.", + }, "limit": { "type": "integer", "minimum": 256, @@ -78,6 +84,7 @@ def run( call_id: str, offset: int = 0, limit: int = 8000, + query: str = "", **_: Any, ) -> ToolOutput: record = self._results.get(str(call_id)) @@ -90,8 +97,19 @@ def run( ), metadata={"is_error": True, "call_id": call_id}, ) - start = max(0, int(offset or 0)) + start = min(len(record.output), max(0, int(offset or 0))) size = min(20_000, max(256, int(limit or 8000))) + match_offset = None + if query: + match = re.compile(re.escape(query), re.IGNORECASE).search(record.output, start) + if match is None: + return ToolOutput( + text=f"No literal match for {query!r} at or after offset={start} in call_id={record.call_id!r}.", + metadata={"recalled": True, "source_call_id": record.call_id, + "matched": False, "total_chars": len(record.output)}, + ) + match_offset = match.start() + start = max(start, match_offset - size // 4) end = min(len(record.output), start + size) excerpt = record.output[start:end] more = end < len(record.output) @@ -110,6 +128,7 @@ def run( "offset": start, "next_offset": end if more else None, "total_chars": len(record.output), + **({"matched": True, "match_offset": match_offset} if query else {}), }, ) diff --git a/tests/test_app_build_workflow.py b/tests/test_app_build_workflow.py new file mode 100644 index 0000000..f922dcb --- /dev/null +++ b/tests/test_app_build_workflow.py @@ -0,0 +1,56 @@ +"""Real local tools/subprocesses, scripted model: build, fail, patch, verify.""" +import pytest + +from shipit_agent import Agent +from shipit_agent.llms.base import LLMResponse +from shipit_agent.models import ToolCall +from shipit_agent.tools.file_write import FileWriteTool +from shipit_agent.tools.file_read import FileReadTool +from shipit_agent.tools.edit_file import EditFileTool +from shipit_agent.tools.code_execution import CodeExecutionTool + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_build_and_repair_small_wsgi_application(tmp_path, streaming): + app = ('def application(environ, start_response):\n' + ' start_response("500 Internal Server Error", [("Content-Type", "text/plain")])\n' + ' return [b"Hello Shipit"]\n') + check = ('import runpy\n' + 'app = runpy.run_path("app.py")["application"]\n' + 'statuses = []\n' + 'body = app({}, lambda status, headers: statuses.append(status))\n' + 'assert statuses == ["200 OK"], statuses\n' + 'assert body == [b"Hello Shipit"]\n' + 'print("APP_CHECK_PASSED")\n') + steps = [ + ("write_file", {"path": "app.py", "content": app}), + ("read_file", {"path": "app.py"}), + ("run_code", {"language": "python", "code": check}), + ("edit_file", {"path": "app.py", "old_text": "500 Internal Server Error", "new_text": "200 OK"}), + ("run_code", {"language": "python", "code": check}), + ] + class Model: + calls = 0 + def complete(self, *, messages, **kwargs): + index = self.calls + self.calls += 1 + if index == 3: + assert any(m.role == "tool" and "AssertionError" in m.content for m in messages) + if index == len(steps): + assert any(m.role == "tool" and "APP_CHECK_PASSED" in m.content for m in messages) + return LLMResponse(content="Created and verified the application.") + name, arguments = steps[index] + return LLMResponse(tool_calls=[ToolCall(name=name, arguments=arguments, id=f"step-{index}")]) + model = Model() + agent = Agent(llm=model, tools=[ + FileWriteTool(root_dir=tmp_path), FileReadTool(root_dir=tmp_path), + EditFileTool(root_dir=tmp_path), CodeExecutionTool(workspace_root=tmp_path), + ], auto_use_skills=False, auto_project_memory=False, auto_project_skills=False) + if streaming: + events = list(agent.stream("Build and test a small greeting application.")) + assert any(e.type == "run_completed" for e in events) + else: + assert agent.run("Build and test a small greeting application.").output == "Created and verified the application." + assert model.calls == 6 + assert "200 OK" in (tmp_path / "app.py").read_text() + assert "500 Internal Server Error" not in (tmp_path / "app.py").read_text() diff --git a/tests/test_async_chat_session.py b/tests/test_async_chat_session.py new file mode 100644 index 0000000..7b15a26 --- /dev/null +++ b/tests/test_async_chat_session.py @@ -0,0 +1,132 @@ +import asyncio + +import pytest + +from shipit_agent import Agent +from shipit_agent.chat_session import AgentChatSession +from shipit_agent.llms.base import LLMResponse +from shipit_agent.models import AgentEvent + + +class Model: + def complete(self, **kwargs): + return LLMResponse(content="answer", usage={"prompt_tokens": 10, "completion_tokens": 2}) + + +def test_async_chat_preserves_history_callbacks_and_budget(): + async def scenario(): + chat = Agent(llm=Model(), max_session_tokens=20, auto_use_skills=False, + auto_project_memory=False, auto_project_skills=False).chat_session(session_id="async") + seen = [] + chat.add_event_callback(seen.append) + await chat.asend("first") + events = [event async for event in chat.astream("second")] + assert events and seen + assert [m.content for m in chat.history() if m.role == "user"] == ["first", "second"] + result = await chat.asend("third") + assert result.metadata["run_summary"]["incomplete_reason"] == "session_token_budget" + asyncio.run(scenario()) + + +@pytest.mark.parametrize("transport", ["websocket", "sse"]) +def test_closing_async_packet_stream_closes_underlying_producer(monkeypatch, transport): + closed = [] + class Producer: + async def astream(self, prompt): + try: + yield AgentEvent(type="run_started", message="started") + yield AgentEvent(type="final_answer", message="done") + finally: + closed.append(True) + monkeypatch.setattr(AgentChatSession, "_session_agent", lambda self: Producer()) + async def scenario(): + chat = Agent(llm=Model()).chat_session(session_id="close") + stream = chat.astream_packets("hi", transport=transport) + packet = await anext(stream) + assert isinstance(packet, str if transport == "sse" else dict) + await stream.aclose() + assert closed == [True] + asyncio.run(scenario()) + + +def test_real_agent_wrapper_awaits_runtime_stream_cleanup(monkeypatch): + from shipit_agent.async_runtime import AsyncAgentRuntime + closed = [] + async def stream(self, *args, **kwargs): + try: + yield AgentEvent(type="run_started", message="started") + await asyncio.sleep(60) + finally: + await asyncio.sleep(0) + closed.append(True) + monkeypatch.setattr(AsyncAgentRuntime, "stream", stream) + async def scenario(): + chat = Agent(llm=Model(), auto_use_skills=False, + auto_project_skills=False, auto_project_memory=False).chat_session(session_id="close-runtime") + events = chat.astream_packets("hi") + await anext(events) + await events.aclose() + assert closed == [True] # no GC or event-loop shutdown required + asyncio.run(scenario()) + + +def test_two_concurrent_twenty_turn_chats_keep_tools_and_history_isolated(): + from shipit_agent import FunctionTool + from shipit_agent.models import ToolCall + executed = [] + + def lookup(query: str) -> str: + executed.append(query) + return "evidence:" + query + + class ToolModel: + def complete(self, *, messages, **kwargs): + prompts = [line for m in messages if m.role == "user" + for line in m.content.splitlines() if line.startswith("tenant-")] + latest = prompts[-1] + tenant, turn = latest.split(":") + assert all(p.startswith(tenant + ":") for p in prompts) + assert len(prompts) == int(turn) + 1, [m.content for m in messages if m.role == "user"] + usage = {"prompt_tokens": 20, "completion_tokens": 5} + if any(m.role == "tool" and "evidence:" + latest in m.content for m in messages): + return LLMResponse(content="Verified " + latest, usage=usage) + return LLMResponse(content="", usage=usage, tool_calls=[ToolCall( + name="lookup", arguments={"query": latest}, id="call-" + latest)]) + + async def scenario(): + agent = Agent(llm=ToolModel(), tools=[FunctionTool.from_callable(lookup)], + auto_use_skills=False, auto_project_memory=False, + auto_project_skills=False, max_session_tokens=2000) + async def converse(tenant): + chat = agent.chat_session(session_id=tenant) + for turn in range(20): + prompt = f"{tenant}:{turn}" + if turn % 2: + events = [e async for e in chat.astream(prompt)] + completed = [e for e in events if e.type == "run_completed"] + assert len(completed) == 1 + assert completed[0].payload["output"] == "Verified " + prompt + else: + result = await chat.asend(prompt) + assert result.output == "Verified " + prompt + assert len([m for m in chat.history() if m.role == "tool"]) == 20 + assert chat._runtime_state["session_tokens"] == 1000 + await asyncio.gather(converse("tenant-a"), converse("tenant-b")) + asyncio.run(scenario()) + assert len(executed) == len(set(executed)) == 40 + + +def test_async_stream_propagates_worker_error_instead_of_silent_success(): + class BrokenModel: + def complete(self, **kwargs): + raise ValueError("fixture provider failure") + async def scenario(): + chat = Agent(llm=BrokenModel(), auto_use_skills=False, + auto_project_memory=False, auto_project_skills=False).chat_session(session_id="broken") + seen = [] + with pytest.raises(ValueError, match="fixture provider failure"): + async for event in chat.astream("hi"): + seen.append(event.type) + assert "run_started" in seen + assert "final_answer" not in seen + asyncio.run(scenario()) diff --git a/tests/test_custom_hook_contracts.py b/tests/test_custom_hook_contracts.py new file mode 100644 index 0000000..a7ef89b --- /dev/null +++ b/tests/test_custom_hook_contracts.py @@ -0,0 +1,79 @@ +import asyncio + +import pytest + +from shipit_agent import Agent, AgentHooks, FunctionTool +from shipit_agent.models import ToolResult, ToolCall +from shipit_agent.llms.base import LLMResponse +from shipit_agent.tools.base import ToolOutput + + +def test_chained_hooks_validate_rewritten_arguments(): + hooks = AgentHooks() + hooks.on_before_tool(lambda n, a: {"decision": "allow", "updated_arguments": {"path": "blocked"}}) + hooks.on_before_tool(lambda n, a: False if a["path"] == "blocked" else None) + assert hooks.run_before_tool("write", {"path": "safe"}).denied + + +def test_approval_cannot_be_erased_by_later_rewrite(): + hooks = AgentHooks() + hooks.on_before_tool(lambda n, a: {"decision": "ask", "reason": "review"}) + hooks.on_before_tool(lambda n, a: {"decision": "allow", "updated_arguments": {"value": 2}}) + result = hooks.run_before_tool("write", {"value": 1}) + assert result.needs_approval and result.updated_arguments == {"value": 2} + + +def test_async_policy_is_rejected_not_silently_ignored(): + hooks = AgentHooks() + @hooks.on_before_tool_matching("write*") + async def deny(name, args): + return False + with pytest.raises(TypeError, match="synchronous"): + hooks.run_before_tool("write_file", {}) + + +@pytest.mark.parametrize("replacement", ["redacted", {"output": "redacted"}, ToolOutput(text="redacted")]) +def test_replacing_output_clears_stale_model_view(replacement): + hooks = AgentHooks(after_tool=[lambda n, r: replacement]) + result = hooks.run_after_tool("lookup", ToolResult(name="lookup", output="SECRET", model_text="SECRET")) + assert result.output == "redacted" + assert result.model_text is None + + +@pytest.mark.parametrize("mode", ["run", "stream", "arun", "astream"]) +def test_custom_hooks_work_across_agent_entrypoints(mode): + calls = [] + seen = [] + hooks = AgentHooks() + hooks.on_user_prompt(lambda p: p.replace("PRIVATE", "public")) + hooks.on_before_tool_matching("lookup")(lambda n, a: { + "decision": "allow", "updated_arguments": {"query": "scoped"}}) + hooks.on_after_tool_matching("lookup")(lambda n, r: "safe evidence") + def lookup(query: str): + calls.append(query) + return ToolOutput(text="SECRET", model_text="SECRET") + class Model: + count = 0 + def complete(self, *, messages, **kwargs): + seen.extend(messages) + self.count += 1 + if self.count == 1: + return LLMResponse(tool_calls=[ToolCall(name="lookup", arguments={"query": "original"})]) + assert any(m.role == "tool" and m.content == "safe evidence" for m in messages) + return LLMResponse(content="done") + agent = Agent(llm=Model(), tools=[FunctionTool.from_callable(lookup)], hooks=hooks, + auto_use_skills=False, auto_project_memory=False, auto_project_skills=False) + async def async_run(): + if mode == "arun": + await agent.arun("PRIVATE") + else: + async for _ in agent.astream("PRIVATE"): + pass + if mode == "run": + agent.run("PRIVATE") + elif mode == "stream": + list(agent.stream("PRIVATE")) + else: + asyncio.run(async_run()) + assert calls == ["scoped"] + assert all("SECRET" not in str(m.content) and "PRIVATE" not in str(m.content) for m in seen) diff --git a/tests/test_file_session_identity.py b/tests/test_file_session_identity.py new file mode 100644 index 0000000..8d0d36b --- /dev/null +++ b/tests/test_file_session_identity.py @@ -0,0 +1,42 @@ +import json + +import pytest + +from shipit_agent.models import Message +from shipit_agent.stores.session import FileSessionStore, SessionRecord + + +@pytest.mark.parametrize("first,second", [ + ("team/chat", "team_chat"), ("team\\chat", "team_chat"), + ("../chat", ".._chat"), ("é/chat", "é_chat"), + ("Chat", "chat"), +]) +def test_distinct_chat_ids_never_share_a_file(tmp_path, first, second): + store = FileSessionStore(tmp_path) + for ident in (first, second): + store.save(SessionRecord(ident, [Message(role="user", content=ident)])) + for ident in (first, second): + record = store.load(ident) + assert record.session_id == ident + assert record.messages[0].content == ident + assert store._path_for(ident).parent == tmp_path + assert {r.session_id for r in store.list_all()} == {first, second} + assert store._path_for(first).name.casefold() != store._path_for(second).name.casefold() + + +def test_legacy_collision_cannot_expose_another_chat(tmp_path): + (tmp_path / "team_chat.json").write_text(json.dumps({ + "session_id": "team/chat", "messages": [], "metadata": {}, + })) + store = FileSessionStore(tmp_path) + assert store.load("team_chat") is None + assert store.load("team/chat").session_id == "team/chat" + store.save(SessionRecord("team/chat")) + assert len(store.list_all()) == 1 + + +def test_long_session_ids_are_bounded(tmp_path): + store = FileSessionStore(tmp_path) + ident = "long" * 1000 + store.save(SessionRecord(ident)) + assert store.load(ident).session_id == ident diff --git a/tests/test_mcp_protocol_boundaries.py b/tests/test_mcp_protocol_boundaries.py new file mode 100644 index 0000000..cb56a11 --- /dev/null +++ b/tests/test_mcp_protocol_boundaries.py @@ -0,0 +1,99 @@ +"""Keep malformed or unrelated MCP responses out of agent evidence.""" + +import json +from io import BytesIO + +import pytest + +from shipit_agent.mcp import ( + MCPError, MCPHTTPTransport, MCPStreamableHTTPTransport, RemoteMCPServer, +) + + +def respond(monkeypatch, body, content_type="application/json"): + response = BytesIO(body.encode()) + response.headers = { + "Content-Type": content_type, "Mcp-Session-Id": "new-session", + } + monkeypatch.setattr("shipit_agent.mcp.request.urlopen", lambda *a, **k: response) + + +@pytest.mark.parametrize("transport_cls", [MCPHTTPTransport, MCPStreamableHTTPTransport]) +@pytest.mark.parametrize("response_id", [None, True, "1", 2]) +def test_rejects_unrelated_response(monkeypatch, transport_cls, response_id): + respond(monkeypatch, json.dumps({"id": response_id, "result": {"tools": []}})) + transport = transport_cls("https://example.test/mcp") + with pytest.raises(MCPError, match="response id"): + transport.request("tools/list") + if isinstance(transport, MCPStreamableHTTPTransport): + assert transport._session_id is None + + +@pytest.mark.parametrize("transport_cls", [MCPHTTPTransport, MCPStreamableHTTPTransport]) +@pytest.mark.parametrize("body", ["not json", '{"id":1,"result":[]}', '{"id":1}']) +def test_rejects_malformed_response(monkeypatch, transport_cls, body): + respond(monkeypatch, body) + with pytest.raises(MCPError): + transport_cls("https://example.test/mcp").request("tools/list") + + +def test_sse_selects_matching_response_and_updates_session(monkeypatch): + events = [ + {"method": "notifications/progress", "params": {}}, + {"id": 9, "error": {"message": "unrelated failure"}}, + {"id": 1, "result": {"tools": [{"name": "correct"}]}}, + ] + respond(monkeypatch, "".join("data: " + json.dumps(e) + "\r\n\r\n" for e in events), + "text/event-stream") + transport = MCPStreamableHTTPTransport("https://example.test/mcp") + assert transport.request("tools/list") == {"tools": [{"name": "correct"}]} + assert transport._session_id == "new-session" + + +class Pages: + def __init__(self, pages): + self.pages = iter(pages) + self.params = [] + + def request(self, method, params=None): + if method == "initialize": + return {"capabilities": {}} + self.params.append(params) + return next(self.pages) + + def notify(self, *args): + pass + + def close(self): + pass + + +def test_pagination_preserves_opaque_cursor_and_finishes_at_limit(): + transport = Pages([ + {"tools": [{"name": "a"}], "nextCursor": " opaque/+== "}, + {"tools": [{"name": "b"}]}, + ]) + server = RemoteMCPServer(name="pages", transport=transport, max_discovery_pages=2) + assert len(server.discover_tools()) == 2 + assert transport.params == [{}, {"cursor": " opaque/+== "}] + + +def test_unbounded_cursor_sequence_fails_without_partial_tool_catalog(): + transport = Pages([ + {"tools": [{"name": "a"}], "nextCursor": "one"}, + {"tools": [{"name": "b"}], "nextCursor": "two"}, + ]) + server = RemoteMCPServer(name="pages", transport=transport, max_discovery_pages=2) + with pytest.raises(MCPError, match="max_discovery_pages=2"): + server.discover_tools() + assert len(transport.params) == 2 + assert server.tools == [] + assert not server._discovered + + +@pytest.mark.parametrize("page", [{}, {"tools": {}}, {"tools": ["bad"]}, + {"tools": [], "nextCursor": 123}]) +def test_invalid_page_is_not_silently_accepted(page): + server = RemoteMCPServer(name="pages", transport=Pages([page])) + with pytest.raises(MCPError, match="Invalid tools/list"): + server.discover_tools() diff --git a/tests/test_mcp_resilience.py b/tests/test_mcp_resilience.py index 86cec3a..5d388b2 100644 --- a/tests/test_mcp_resilience.py +++ b/tests/test_mcp_resilience.py @@ -105,10 +105,20 @@ def test_fails_fast_when_breaker_open_without_calling_inner(): def test_retries_one_transient_failure_then_succeeds(): inner = FakeTransport(script=[MCPError("subprocess died"), None]) t = ResilientMCPTransport(inner, max_retries=1, base_delay=0, name="s") - assert t.request("tools/call")["ok"] is True + assert t.request("tools/list")["ok"] is True assert inner.calls == 2 # failed once, retried, succeeded +@pytest.mark.parametrize("method", ["tools/call", "custom/write"]) +@pytest.mark.parametrize("error", [MCPError("connection lost"), TimeoutError("response lost")]) +def test_ambiguous_execution_is_never_replayed(method, error): + inner = FakeTransport(script=[error, None]) + transport = ResilientMCPTransport(inner, max_retries=3, base_delay=0) + with pytest.raises(type(error)): + transport.request(method, {"name": "create_ticket"}) + assert inner.calls == 1 + + def test_never_retries_a_rate_limit(): inner = FakeTransport(script=[MCPError("429 too many requests"), None]) cb = CircuitBreaker(fail_threshold=5) diff --git a/tests/test_mcp_result_integrity.py b/tests/test_mcp_result_integrity.py new file mode 100644 index 0000000..c22d147 --- /dev/null +++ b/tests/test_mcp_result_integrity.py @@ -0,0 +1,55 @@ +import json + +import pytest + +from shipit_agent.mcp import MCPError, MCPRemoteTool + + +class Transport: + def __init__(self, result): + self.result = result + def request(self, method, params): + if isinstance(self.result, Exception): + raise self.result + return self.result + + +def tool(result, **kwargs): + return MCPRemoteTool(server_name="fixture", name="lookup", description="Lookup", + transport=Transport(result), **kwargs) + + +@pytest.mark.parametrize("value", [{}, {"items": [{"id": "case-1", "amount": 42}]}]) +def test_structured_only_result_is_visible_to_model(value): + result = tool({"structuredContent": value}).run(None) + assert json.loads(result.text) == value + assert result.metadata["ok"] is True + + +def test_text_and_structured_content_are_not_duplicated(): + result = tool({"content": [{"type": "text", "text": "42"}], + "structuredContent": {"amount": 42}}).run(None) + assert result.text == "42" + assert result.metadata["structured_content"] == {"amount": 42} + + +@pytest.mark.parametrize("result", [MCPError("disconnected"), {"isError": True}]) +def test_static_metadata_cannot_disguise_failed_execution(result): + output = tool(result, metadata={"ok": True, "is_error": False, "error": "old", "server": "wrong"}).run(None) + assert output.metadata["ok"] is False + assert output.metadata["is_error"] is True + assert output.metadata["server"] == "fixture" + + +def test_static_error_metadata_cannot_turn_success_into_failure(): + output = tool({"content": [{"type": "text", "text": "ok"}]}, + metadata={"error": "stale", "is_error": True, "ok": False}).run(None) + assert output.metadata["error"] is None + assert output.metadata["is_error"] is False + assert output.metadata["ok"] is True + + +@pytest.mark.parametrize("hint", [False, "false", "true", 1, None]) +def test_read_only_hint_must_be_literal_true(hint): + assert tool({}, annotations={"readOnlyHint": hint}).read_only is False + assert tool({}, annotations={"readOnlyHint": True}).read_only is True diff --git a/tests/test_mcp_schema_cache.py b/tests/test_mcp_schema_cache.py index b2f536e..f032683 100644 --- a/tests/test_mcp_schema_cache.py +++ b/tests/test_mcp_schema_cache.py @@ -93,6 +93,24 @@ def test_env_values_are_not_written_to_disk(cache_home): assert "hunter2" not in on_disk and "SECRET_TOKEN" not in on_disk +@pytest.mark.parametrize("payload", [[], None, "invalid", 42]) +def test_valid_json_with_wrong_root_is_cache_miss(cache_home, payload): + import json + sc.save("srv", "fp", [{"name": "a"}]) + sc._path_for("srv", "fp").write_text(json.dumps(payload)) + assert sc.load("srv", "fp") is None + + +@pytest.mark.parametrize("descriptor", [None, "tool", {}, {"name": 3}, + {"name": "a", "input_schema": "invalid"}, {"name": "a", "annotations": []}]) +def test_bad_cached_descriptor_falls_back_to_live_discovery(cache_home, descriptor): + transport = RecordingTransport() + server = _server(transport) + sc.save("fake", server._cache_fingerprint(), [descriptor]) + assert [tool.name for tool in server.discover_tools()] == ["do_thing"] + assert "tools/list" in transport.calls + + def test_fingerprint_is_stable_and_config_sensitive(): base = dict(name="s", identity="stdio:cmd", protocol_version="2025-11-25", allowed=None, blocked=set(), include_server_in_tool_names=False, env=None) diff --git a/tests/test_result_recall.py b/tests/test_result_recall.py index 23c81ec..5862887 100644 --- a/tests/test_result_recall.py +++ b/tests/test_result_recall.py @@ -41,6 +41,30 @@ def test_unknown_call_id_is_a_recoverable_result() -> None: assert "No recallable result" in output.text +def test_literal_search_retrieves_small_exact_excerpt_from_large_result(): + text = "unrelated evidence\n" * 20000 + "Invoice [A-42]: EUR 720" + "\nother" * 20000 + tool = RecallToolResult(recallable_results([ + Message(role="tool", name="search", content=text, tool_call_id="large") + ], min_chars=1000)) + output = tool.run(None, call_id="large", query="invoice [a-42]", limit=512) + assert "Invoice [A-42]: EUR 720" in output.text + assert len(output.text) < 750 + assert output.metadata["matched"] is True + assert output.metadata["match_offset"] == text.index("Invoice") + assert output.text.split("\n", 1)[1] == text[output.metadata["offset"]:output.metadata["next_offset"]] + absent = tool.run(None, call_id="large", query="not-present") + assert absent.metadata["matched"] is False + + +def test_recall_offset_beyond_end_is_clamped(): + tool = RecallToolResult(recallable_results([ + Message(role="tool", content=BIG, tool_call_id="large") + ], min_chars=1000)) + output = tool.run(None, call_id="large", offset=999999) + assert output.metadata["offset"] == len(BIG) + assert output.metadata["next_offset"] is None + + def test_runtime_installs_recall_only_when_useful() -> None: registry = ToolRegistry() name = AgentRuntime.install_result_recall( diff --git a/tests/test_session_manager.py b/tests/test_session_manager.py index a4de082..74515a8 100644 --- a/tests/test_session_manager.py +++ b/tests/test_session_manager.py @@ -91,3 +91,41 @@ def test_fork_session(): assert forked_record is not None assert len(forked_record.messages) == 2 assert forked_record.metadata["forked_from"] == chat.session_id + + +def test_manager_binds_its_store_without_mutating_agent(): + from shipit_agent import Agent + from shipit_agent.llms import SimpleEchoLLM + agent = Agent(llm=SimpleEchoLLM(), auto_use_skills=False) + original_store = agent.session_store + store = InMemorySessionStore() + manager = SessionManager(store) + chat = manager.create(agent) + chat.send("remember this") + resumed = manager.resume(agent, chat.session_id) + assert any(m.content == "remember this" for m in resumed.history()) + assert agent.session_store is original_store + assert original_store.load(chat.session_id) is None + + +def test_historical_fork_does_not_inherit_future_facts_or_mutable_messages(): + store = InMemorySessionStore() + manager = SessionManager(store) + agent = FakeAgent(store) + chat = manager.create(agent) + original = store.load(chat.session_id) + original.messages = [Message(role="user", content="early"), + Message(role="assistant", content="later secret")] + original.metadata.update(verified_facts=[{"value": "later secret"}], + compaction_checkpoint={"summary": "later secret"}, + tool_discovery={"schemas": {}}, + custom={"labels": ["keep"]}) + fork = manager.fork(agent, chat.session_id, from_message=1) + record = store.load(fork.session_id) + assert "verified_facts" not in record.metadata + assert "compaction_checkpoint" not in record.metadata + assert "tool_discovery" not in record.metadata + record.messages[0].content = "changed" + record.metadata["custom"]["labels"].append("fork only") + assert original.messages[0].content == "early" + assert original.metadata["custom"]["labels"] == ["keep"] diff --git a/tests/test_tool_output_budget.py b/tests/test_tool_output_budget.py index a0d37eb..13a1fff 100644 --- a/tests/test_tool_output_budget.py +++ b/tests/test_tool_output_budget.py @@ -86,6 +86,23 @@ def run(*, turns=4, tools=None, cap=None): class TestTheCap: + def test_changed_semantic_view_is_not_discarded_as_duplicate_raw_output(self): + class SemanticDump: + name = "dump" + description = "Read evidence" + calls = 0 + def schema(self): + return {"function": {"name": "dump", "parameters": {"type": "object", "properties": {}}}} + def run(self, context, **kwargs): + self.calls += 1 + view = ("first evidence " if self.calls == 1 else "second evidence ") * 200 + return ToolOutput(text=BIG, model_text=view) + result, _ = run(turns=3, tools=SemanticDump()) + messages = [m.content for m in result.messages if m.role == "tool"] + assert "first evidence" in messages[0] + assert "second evidence" in messages[1] + assert "No new data" in messages[2] + def test_it_is_on_by_default(self) -> None: assert Agent(llm=None).max_tool_output_chars == 16_000 diff --git a/tests/test_tool_workspace_boundaries.py b/tests/test_tool_workspace_boundaries.py new file mode 100644 index 0000000..3daaa6e --- /dev/null +++ b/tests/test_tool_workspace_boundaries.py @@ -0,0 +1,32 @@ +import pytest + +from shipit_agent.tools._shared import Workspace +from shipit_agent.tools.edit_file import EditFileTool +from shipit_agent.tools.base import ToolContext + + +def test_workspace_rejects_prefix_sibling_and_symlink_escape(tmp_path): + root = tmp_path / "project" + sibling = tmp_path / "project-backup" + root.mkdir() + sibling.mkdir() + workspace = Workspace(root=root) + for path in ("../project-backup/secret.txt", str(sibling / "secret.txt")): + with pytest.raises(ValueError, match="outside"): + workspace.resolve(path) + (root / "link").symlink_to(sibling, target_is_directory=True) + with pytest.raises(ValueError, match="outside"): + workspace.resolve("link/secret.txt") + assert workspace.resolve("src/app.py") == root / "src/app.py" + + +@pytest.mark.parametrize("replace_all", [False, True]) +def test_empty_edit_cannot_insert_between_every_character(tmp_path, replace_all): + path = tmp_path / "app.py" + path.write_text("print('hello')\n") + context = ToolContext(prompt="edit", state={"read_files": [str(path)]}) + output = EditFileTool(root_dir=tmp_path).run( + context, path="app.py", old_text="", new_text="oops", replace_all=replace_all, + ) + assert output.metadata["is_error"] is True + assert path.read_text() == "print('hello')\n" diff --git a/tests/test_usage_normalization.py b/tests/test_usage_normalization.py new file mode 100644 index 0000000..fd7b497 --- /dev/null +++ b/tests/test_usage_normalization.py @@ -0,0 +1,103 @@ +"""Provider omissions must not bypass task/session token budgets.""" +from types import SimpleNamespace + +import pytest + +from shipit_agent.llms.base import LLMResponse +from shipit_agent.llms.openai_adapter import _usage_dict +from shipit_agent.llms.usage import has_complete_token_usage, normalize_token_usage + + +@pytest.mark.parametrize("wrap", [dict, lambda **kw: SimpleNamespace(**kw)]) +@pytest.mark.parametrize("bad", [None, True, -1, "12", 1.5]) +def test_invalid_counter_stays_unknown(wrap, bad): + usage = normalize_token_usage(wrap(prompt_tokens=bad, completion_tokens=5)) + assert "prompt_tokens" not in usage + assert not has_complete_token_usage(LLMResponse(content="", usage=usage)) + + +def test_explicit_zero_is_valid(): + assert normalize_token_usage({"prompt_tokens": 0, "completion_tokens": 0}) == { + "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, + } + + +def test_anthropic_field_names_preserve_omissions(): + assert normalize_token_usage(SimpleNamespace(input_tokens=8), + input_key="input_tokens", output_key="output_tokens") == { + "prompt_tokens": 8, + } + + +def test_openai_cache_only_usage_does_not_invent_free_completion(): + usage = _usage_dict(SimpleNamespace( + prompt_tokens_details=SimpleNamespace(cached_tokens=100), total_tokens=110, + )) + assert usage == {"total_tokens": 110, "cache_read_input_tokens": 100} + assert not has_complete_token_usage(LLMResponse(content="", usage=usage)) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_native_anthropic_missing_output_remains_unknown(monkeypatch, stream): + import sys + from shipit_agent.llms.anthropic_adapter import AnthropicChatLLM + from shipit_agent.models import Message + + response = SimpleNamespace( + content=[SimpleNamespace(type="text", text="answer")], + usage=SimpleNamespace(input_tokens=15), stop_reason="end_turn", + ) + class Stream: + text_stream = ["answer"] + def __enter__(self): + return self + def __exit__(self, *args): + pass + def get_final_message(self): + return response + messages = SimpleNamespace(create=lambda **kw: response, stream=lambda **kw: Stream()) + monkeypatch.setitem(sys.modules, "anthropic", SimpleNamespace(Anthropic=lambda **kw: + SimpleNamespace(messages=messages, beta=SimpleNamespace(messages=messages)))) + result = AnthropicChatLLM(model="fixture").complete( + messages=[Message(role="user", content="hello")], + text_delta_callback=(lambda text: None) if stream else None, + ) + assert result.usage == {"prompt_tokens": 15} + assert not has_complete_token_usage(result) + + +@pytest.mark.parametrize("provider", ["openai", "litellm"]) +@pytest.mark.parametrize("stream", [False, True]) +def test_adapter_omission_blocks_next_budgeted_run(monkeypatch, provider, stream): + import sys + from shipit_agent import Agent + from shipit_agent.llms.openai_adapter import OpenAIChatLLM + from shipit_agent.llms.litellm_adapter import LiteLLMChatLLM + + calls = [] + def complete(**kwargs): + calls.append(kwargs) + usage = SimpleNamespace(total_tokens=10) + message = SimpleNamespace(content="answer", tool_calls=None) + if kwargs.get("stream"): + return iter([SimpleNamespace(choices=[SimpleNamespace( + delta=message, finish_reason="stop")], usage=usage)]) + return SimpleNamespace(choices=[SimpleNamespace(message=message, + finish_reason="stop")], usage=usage) + + if provider == "openai": + monkeypatch.setitem(sys.modules, "openai", SimpleNamespace(OpenAI=lambda **kw: + SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=complete))))) + llm = OpenAIChatLLM(model="fixture", api_key="fixture") + else: + monkeypatch.setitem(sys.modules, "litellm", SimpleNamespace(completion=complete)) + llm = LiteLLMChatLLM(model="fixture") + agent = Agent(llm=llm, max_session_tokens=100, auto_use_skills=False, + auto_project_memory=False, auto_project_skills=False) + if stream: + list(agent.stream("hello")) + else: + agent.run("hello") + result = agent.run("continue") + assert len(calls) == 1 + assert result.metadata["run_summary"]["incomplete_reason"] == "session_budget_usage_unavailable" diff --git a/tests/test_workspace.py b/tests/test_workspace.py index b20bc21..437ce68 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -247,7 +247,7 @@ def complete(self, **kwargs): assert user_messages == ["remember alpha", "now beta"] assert isinstance(restarted.session_store, FileSessionStore) assert isinstance(restarted.memory_store, FileMemoryStore) - assert (tmp_path / ".shipit" / "sessions" / "main.json").is_file() + assert restarted.session_store._path_for("main").is_file() assert (tmp_path / ".shipit" / "memory.json").is_file() def test_optimized_project_agent_preserves_explicit_stores(self, tmp_path) -> None: