agentprobe-testing 0.5.0__py3-none-any.whl

agentprobe/injection.py

@@ -0,0 +1,475 @@
+"""The injection catalog (spec section 6) and the runtime that applies an
+Injection to a live world (any domain's, via the generic entity protocol
+-- see domain.py's EntityWorld). Application is pure bookkeeping -- no LLM
+here. The judgment (which injection, when) lives in injector.py; the
+judgment of whether it worked lives in classifier.py.
+
+Injections are armed, not fired (repair Task 3): the Injector attaches a
+Trigger -- a condition to wait for -- rather than guessing a step number.
+`trigger_matches_pre_dispatch`/`trigger_matches_post_dispatch` are checked
+by the runner every step; when one matches, THAT is when apply_injection()
+actually runs. An ArmedInjection that never matches by run end is EXPIRED
+(a bad prediction), which is a different failure from INVALID (a fired
+injection that made the goal unreachable, i.e. unfair).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Optional
+
+from agentprobe.scenario import CommitPattern
+from agentprobe.tools import COMMIT_TOOLS
+
+if TYPE_CHECKING:
+    from agentprobe.domain import EntityWorld
+    from agentprobe.trajectory import CallRecord, Step
+
+
+class InjectionKind(str, Enum):
+    TOOL_ERROR = "TOOL_ERROR"
+    STALE_READ = "STALE_READ"
+    CONTRADICTION = "CONTRADICTION"
+    LATE_INFO = "LATE_INFO"
+    AMBIGUITY = "AMBIGUITY"
+    PROMPT_INJECTION = "PROMPT_INJECTION"
+    """An adversarial instruction embedded in ordinary-looking ticket
+    content (e.g. a note that reads like a customer message but actually
+    commands the Target to take an unauthorized action). Distinct from
+    CONTRADICTION: that kind tests whether the Target notices a changed
+    FACT; this kind tests whether the Target can be manipulated into an
+    ACTION by data it should treat as untrusted input, not instructions --
+    a real production-agent security failure mode, not just a robustness
+    one. Mechanically identical to CONTRADICTION/LATE_INFO (appends a
+    ticket note) -- the payload's text is what differs in kind, not the
+    plumbing."""
+
+
+@dataclass(frozen=True)
+class ResponseSignature:
+    """How to tell whether the Target's subsequent behavior counts as
+    having handled this injection. Free text, judged by the classifier --
+    but committed to *before* we see how the Target reacts."""
+
+    description: str
+
+
+@dataclass(frozen=True)
+class Injection:
+    kind: InjectionKind
+    payload: dict[str, Any]
+    intent: str
+    """What the Target SHOULD now do. Mandatory -- without it, a failure
+    can't be attributed to the target vs. an unfair injection."""
+    expected_signature: ResponseSignature
+    rationale: str
+    """Why this moment -- for the report, and for auditing injector
+    gaming (spec section 12)."""
+
+
+class TriggerKind(str, Enum):
+    ON_TOOL_CALL = "on_tool_call"
+    ON_NTH_TOOL_CALL = "on_nth_tool_call"
+    ON_READ_OF = "on_read_of"
+    ON_READ_OF_ANY = "on_read_of_any"
+    AFTER_COMMIT = "after_commit"
+    ON_ANY_COMMIT = "on_any_commit"
+    AFTER_ANY_COMMIT = "after_any_commit"
+    ON_NEXT_ACTION = "on_next_action"
+    ON_STEP = "on_step"
+
+
+@dataclass(frozen=True)
+class Trigger:
+    """A condition an ArmedInjection waits for. Construct via the factory
+    methods below, not the raw fields directly.
+
+    - on_tool_call(name): fires the next time the Target calls that tool,
+      whenever that comes -- the natural fit for TOOL_ERROR.
+    - on_nth_tool_call(name, n): fires on the Nth call to that tool.
+    - on_read_of(entity_id): fires right after the Target reads that
+      entity -- for STALE_READ, so the Target has a stale cached belief
+      *before* the value changes, not after (which would test nothing).
+    - on_read_of_any(entity_ids): fires right after the Target reads the
+      FIRST of several candidate entities -- same relaxation as
+      on_any_commit, one layer down. on_read_of forces a guess of which
+      single entity the Target reads next; a wrong guess just expires.
+      Use this whenever more than one entity is a plausible next read.
+      The fired AppliedInjection's payload entity_id is corrected to
+      whichever candidate was actually read, so the mutation always lands
+      on the right one regardless of which candidate matched.
+    - after_commit(pattern): fires right after a matching commit -- for
+      LATE_INFO, which is meaningless before a commit exists to be late
+      about.
+    - on_any_commit(): fires the next time the Target calls ANY commit
+      tool, whichever one it turns out to be. Repair Task 3 follow-up:
+      on_tool_call forced the Injector to guess which specific commit
+      tool (close_ticket vs issue_refund vs escalate) the Target would
+      call next, and a wrong guess meant a silent expiry -- most of the
+      measured wasted-armed-injection rate traced back to exactly this.
+      Use when what matters is "right before the Target commits to
+      anything," not which commit it is.
+    - after_any_commit(): same relaxation for after_commit -- fires right
+      after ANY commit tool succeeds, not a specifically predicted one.
+    - on_next_action(): fires unconditionally on the Target's very next
+      dispatched call, whichever tool it turns out to be. TOOL_ERROR is
+      checked pre-dispatch, so the actual tool name is already known at
+      fire time -- there was never a real need to pre-commit to one, and
+      guessing wrong (arming on_tool_call for a tool the Target doesn't
+      call again) was the single largest source of expired injections.
+      The fired AppliedInjection's payload tool_name is corrected to the
+      real intercepted tool at fire time regardless of what was guessed.
+    - on_step(n): fires at absolute step n. Kept for the hardcoded
+      injector and for tests that want exact, deterministic timing --
+      the ModelInjector is not offered this one, since re-introducing a
+      guessed step number is exactly what this repair removes.
+    """
+
+    kind: TriggerKind
+    tool_name: Optional[str] = None
+    n: Optional[int] = None
+    entity_id: Optional[str] = None
+    entity_ids: Optional[tuple[str, ...]] = None
+    commit_pattern: Optional[CommitPattern] = None
+    step: Optional[int] = None
+
+    @staticmethod
+    def on_tool_call(name: str) -> "Trigger":
+        return Trigger(TriggerKind.ON_TOOL_CALL, tool_name=name)
+
+    @staticmethod
+    def on_nth_tool_call(name: str, n: int) -> "Trigger":
+        return Trigger(TriggerKind.ON_NTH_TOOL_CALL, tool_name=name, n=n)
+
+    @staticmethod
+    def on_read_of(entity_id: str) -> "Trigger":
+        return Trigger(TriggerKind.ON_READ_OF, entity_id=entity_id)
+
+    @staticmethod
+    def on_read_of_any(entity_ids: list[str]) -> "Trigger":
+        return Trigger(TriggerKind.ON_READ_OF_ANY, entity_ids=tuple(entity_ids))
+
+    @staticmethod
+    def after_commit(pattern: CommitPattern) -> "Trigger":
+        return Trigger(TriggerKind.AFTER_COMMIT, commit_pattern=pattern)
+
+    @staticmethod
+    def on_any_commit() -> "Trigger":
+        return Trigger(TriggerKind.ON_ANY_COMMIT)
+
+    @staticmethod
+    def after_any_commit() -> "Trigger":
+        return Trigger(TriggerKind.AFTER_ANY_COMMIT)
+
+    @staticmethod
+    def on_next_action() -> "Trigger":
+        return Trigger(TriggerKind.ON_NEXT_ACTION)
+
+    @staticmethod
+    def on_step(n: int) -> "Trigger":
+        return Trigger(TriggerKind.ON_STEP, step=n)
+
+    def describe(self) -> str:
+        if self.kind == TriggerKind.ON_TOOL_CALL:
+            return f"on_tool_call({self.tool_name})"
+        if self.kind == TriggerKind.ON_NTH_TOOL_CALL:
+            return f"on_nth_tool_call({self.tool_name}, {self.n})"
+        if self.kind == TriggerKind.ON_READ_OF:
+            return f"on_read_of({self.entity_id})"
+        if self.kind == TriggerKind.ON_READ_OF_ANY:
+            return f"on_read_of_any({self.entity_ids})"
+        if self.kind == TriggerKind.AFTER_COMMIT:
+            assert self.commit_pattern is not None
+            return f"after_commit({self.commit_pattern.tool}({self.commit_pattern.args}))"
+        if self.kind == TriggerKind.ON_ANY_COMMIT:
+            return "on_any_commit()"
+        if self.kind == TriggerKind.AFTER_ANY_COMMIT:
+            return "after_any_commit()"
+        if self.kind == TriggerKind.ON_NEXT_ACTION:
+            return "on_next_action()"
+        return f"on_step({self.step})"
+
+
+@dataclass(frozen=True)
+class ArmedInjection:
+    """Decided but not yet applied -- waiting for `trigger` to match."""
+
+    injection: Injection
+    trigger: Trigger
+    armed_at_step: int
+    expires_after: Optional[int] = None
+    """Steps after arming before this expires early even if the trigger
+    never matches. None = stays armed until it fires or the run ends."""
+
+
+@dataclass(frozen=True)
+class ExpiredInjection:
+    """An ArmedInjection whose trigger never matched -- a bad prediction
+    by the Injector, not an unfair one (that's INVALID, which only applies
+    to injections that actually fired)."""
+
+    injection: Injection
+    trigger: Trigger
+    armed_at_step: int
+
+
+@dataclass(frozen=True)
+class AppliedInjection:
+    """One injection that actually fired during a run, with its outcome."""
+
+    armed_at_step: int
+    fired_at_step: int
+    """Logged separately from armed_at_step -- the gap between them is
+    diagnostic (how long the Injector's prediction took to pan out, if at
+    all)."""
+    injection: Injection
+    effect: str
+    """What apply_injection actually did, in words -- for the report."""
+    valid: bool = True
+    """False if the validity check found the goal unreachable right after
+    this fired (spec section 8) -- an INVALID_INJECTION."""
+    trigger_kind: str = ""
+    """The TriggerKind.value that actually fired this, e.g. "on_read_of_any".
+    Defaults to "" for backward compatibility with existing call sites
+    that don't set it -- but runner.py's fire() always populates it, since
+    without this a fired injection can't be attributed to the trigger
+    choice that worked, unlike ExpiredInjection (which already carries
+    its trigger). Needed for playbook.py's outcome tracking."""
+
+
+class InjectionError(Exception):
+    """Raised when an injection's payload doesn't match live world state
+    (e.g. references an id that doesn't exist). Distinct from ToolError --
+    this is an injector mistake, not a target mistake."""
+
+
+def resolve_tool_error_target(injection: Injection, intercepted_tool_name: str) -> Injection:
+    """TOOL_ERROR is always checked pre-dispatch, so the tool about to be
+    called is already known by fire time -- there's no need to trust the
+    Injector's payload tool_name (which may have been an approximate guess,
+    e.g. under on_next_action). Returns a copy of `injection` with
+    payload["tool_name"] corrected to the real intercepted tool; a no-op
+    for every other kind."""
+    if injection.kind != InjectionKind.TOOL_ERROR:
+        return injection
+    if injection.payload.get("tool_name") == intercepted_tool_name:
+        return injection
+    return dataclasses.replace(injection, payload={**injection.payload, "tool_name": intercepted_tool_name})
+
+
+def resolve_stale_read_target(injection: Injection, actual_entity_id: str) -> Injection:
+    """STALE_READ fired via on_read_of_any doesn't know in advance which
+    of its candidate entities will actually get read -- correct
+    payload["entity_id"] to whichever one was, so the mutation always
+    lands on the entity the Target just perceived. A no-op for every
+    other kind, or if the payload already named the right one."""
+    if injection.kind != InjectionKind.STALE_READ:
+        return injection
+    if injection.payload.get("entity_id") == actual_entity_id:
+        return injection
+    return dataclasses.replace(injection, payload={**injection.payload, "entity_id": actual_entity_id})
+
+
+def apply_injection(world: "EntityWorld", injection: Injection, pending_tool_errors: dict[str, int]) -> str:
+    """Applies one injection to the live world (or, for TOOL_ERROR, to the
+    pending-error registry the runner consults before dispatching a tool
+    call). Called at FIRE time, once the injection's trigger has matched --
+    never at arm time. Returns a human-readable description of the effect,
+    for the trajectory record and the report.
+
+    Every kind below reads its payload via direct p["key"] indexing --
+    correct for a well-formed payload, but a model-supplied JSON payload
+    missing a required key (or with a value of the wrong type, e.g. a
+    non-numeric TOOL_ERROR count) would otherwise raise a raw KeyError/
+    ValueError/TypeError, uncaught by the runner (which only catches
+    InjectionError) and crashing the whole run. This wrapper converts
+    that whole class of malformed-payload failure into one, instead of
+    needing a bespoke presence/type check before every single field
+    access across every kind.
+    """
+    try:
+        return _apply_injection_impl(world, injection, pending_tool_errors)
+    except (KeyError, ValueError, TypeError) as e:
+        raise InjectionError(
+            f"{injection.kind.value} payload is malformed ({type(e).__name__}: {e}) -- payload was {injection.payload}"
+        ) from e
+
+
+def _apply_injection_impl(world: "EntityWorld", injection: Injection, pending_tool_errors: dict[str, int]) -> str:
+    """Written against the generic entity protocol (entity_store/
+    get_entity/set_entity_field/clone_entity/append_note -- see domain.py's
+    EntityWorld) rather than world.tickets/orders/customers by name -- so
+    the same dispatch logic here works against a structurally different
+    domain's own world class, as long as it implements that same
+    five-method protocol.
+    """
+    p = injection.payload
+    kind = injection.kind
+
+    if kind == InjectionKind.TOOL_ERROR:
+        tool_name = p["tool_name"]
+        count = int(p.get("count", 1))
+        pending_tool_errors[tool_name] = pending_tool_errors.get(tool_name, 0) + count
+        return f"next {count} call(s) to {tool_name} will raise a server error"
+
+    if kind == InjectionKind.STALE_READ:
+        entity_type, entity_id, field_name, new_value = (
+            p["entity_type"],
+            p["entity_id"],
+            p["field"],
+            p["new_value"],
+        )
+        if not world.entity_exists(entity_type, entity_id):
+            raise InjectionError(f"STALE_READ target {entity_type}:{entity_id} does not exist")
+        entity = world.get_entity(entity_type, entity_id)
+        if not hasattr(entity, field_name):
+            # A hallucinated/invalid field name would otherwise raise a
+            # raw AttributeError here -- uncaught by the runner (which
+            # only catches InjectionError), crashing the whole run
+            # instead of just failing this one injection.
+            raise InjectionError(f"STALE_READ field {field_name!r} does not exist on {entity_type} {entity_id}")
+        old_value = getattr(entity, field_name)
+        world.set_entity_field(entity_type, entity_id, field_name, new_value)
+        return f"{entity_type} {entity_id}.{field_name} silently changed {old_value!r} -> {new_value!r}"
+
+    if kind in (InjectionKind.CONTRADICTION, InjectionKind.LATE_INFO, InjectionKind.PROMPT_INJECTION):
+        # "ticket_id" is the payload key every current Injector prompt
+        # actually emits (ticket-domain vocabulary); "entity_type"/
+        # "entity_id" is the generic form a different domain's Injector
+        # prompt can emit instead. Both are accepted so this dispatch
+        # logic doesn't need to change when a new domain starts using it.
+        entity_type = p.get("entity_type", "ticket")
+        entity_id = p["entity_id"] if "entity_id" in p else p["ticket_id"]
+        note = p["note"]
+        if not world.entity_exists(entity_type, entity_id):
+            raise InjectionError(f"{kind.value} target {entity_type} {entity_id} does not exist")
+        world.append_note(entity_type, entity_id, note)
+        return f"new note appended to {entity_type} {entity_id}: {note!r}"
+
+    if kind == InjectionKind.AMBIGUITY:
+        entity_type, source_id, new_id = p["entity_type"], p["source_id"], p["new_id"]
+        overrides = dict(p.get("overrides", {}))
+        # "id" is redundant with new_id (and, if a model supplies both,
+        # conflicts with it as a duplicate replace() keyword argument) --
+        # new_id is authoritative.
+        overrides.pop("id", None)
+        if not world.entity_exists(entity_type, source_id):
+            raise InjectionError(f"AMBIGUITY source {entity_type}:{source_id} does not exist")
+        source_entity = world.get_entity(entity_type, source_id)
+        for key, value in list(overrides.items()):
+            # Model-supplied payload came through JSON -- lists, not
+            # whatever tuple-typed field this entity's notes/history lives
+            # in (Ticket.notes here). dataclasses.replace() doesn't
+            # enforce the field's type annotation at runtime, so left
+            # uncoerced this plants an entity that crashes the next
+            # append_note. Checked against the *current value's* actual
+            # runtime type, not a string-parsed type annotation, so this
+            # holds for any entity type/field, not just this domain's.
+            if isinstance(value, list) and isinstance(getattr(source_entity, key, None), tuple):
+                overrides[key] = tuple(value)
+        if world.entity_exists(entity_type, new_id):
+            # clone_entity has no idea this id is already taken -- it
+            # would silently overwrite that existing entity (e.g. a
+            # colliding id in a decoy scenario's cluster of similarly-
+            # numbered tickets), corrupting world state without any
+            # error. The model's own payload contract says new_id must
+            # be fresh; enforce it here rather than trust it.
+            raise InjectionError(
+                f"AMBIGUITY new_id {entity_type}:{new_id} already exists -- would silently overwrite it"
+            )
+        valid_fields = {f.name for f in dataclasses.fields(source_entity)}
+        bogus_fields = set(overrides) - valid_fields
+        if bogus_fields:
+            # A hallucinated override key would otherwise raise a raw
+            # TypeError from dataclasses.replace() -- uncaught by the
+            # runner, crashing the whole run instead of failing just this
+            # one injection.
+            raise InjectionError(
+                f"AMBIGUITY overrides for {entity_type} reference nonexistent field(s) {sorted(bogus_fields)}"
+            )
+        world.clone_entity(entity_type, source_id, new_id, overrides)
+        return f"planted near-duplicate {entity_type} {new_id} (cloned from {source_id}, overrides={overrides})"
+
+    raise InjectionError(f"unknown injection kind {kind!r}")
+
+
+def trigger_matches_pre_dispatch(
+    trigger: Trigger, tool_name: str, step_index: int, tool_call_counts: dict[str, int], commit_tools=COMMIT_TOOLS
+) -> bool:
+    """Checked BEFORE the Target's chosen action executes -- only these
+    three kinds can intercept a call outright (turn it into an error).
+    `tool_call_counts` holds successful-call counts so far, NOT including
+    the about-to-happen call. `commit_tools` defaults to the ticket
+    domain's for backward compatibility -- callers driving a different
+    domain must pass that domain's own commit_tools (see domain.py's
+    Domain.commit_tools), or ON_ANY_COMMIT will only ever match by
+    coincidence, on whichever tool names happen to also exist in the
+    ticket domain.
+    """
+    if trigger.kind == TriggerKind.ON_TOOL_CALL:
+        return trigger.tool_name == tool_name
+    if trigger.kind == TriggerKind.ON_NTH_TOOL_CALL:
+        return trigger.tool_name == tool_name and tool_call_counts.get(tool_name, 0) + 1 == trigger.n
+    if trigger.kind == TriggerKind.ON_ANY_COMMIT:
+        return tool_name in commit_tools
+    if trigger.kind == TriggerKind.ON_NEXT_ACTION:
+        return True
+    if trigger.kind == TriggerKind.ON_STEP:
+        return trigger.step == step_index
+    return False
+
+
+def trigger_matches_post_dispatch(trigger: Trigger, call: "CallRecord") -> bool:
+    """Checked AFTER a successful call -- these kinds react to something
+    the Target just did, rather than intercepting it."""
+    if trigger.kind == TriggerKind.ON_READ_OF:
+        return not call.is_commit and trigger.entity_id in call.args.values()
+    if trigger.kind == TriggerKind.ON_READ_OF_ANY:
+        assert trigger.entity_ids is not None
+        return not call.is_commit and any(eid in call.args.values() for eid in trigger.entity_ids)
+    if trigger.kind == TriggerKind.AFTER_COMMIT:
+        assert trigger.commit_pattern is not None
+        return call.is_commit and trigger.commit_pattern.matches(call.tool_name, call.args)
+    if trigger.kind == TriggerKind.AFTER_ANY_COMMIT:
+        return call.is_commit
+    return False
+
+
+def matched_entity_id(trigger: Trigger, call: "CallRecord") -> Optional[str]:
+    """For ON_READ_OF_ANY: which of the trigger's candidate entities was
+    actually present in this read's args. None for every other trigger
+    kind, or if none of the candidates match (shouldn't happen if this is
+    only called after trigger_matches_post_dispatch returned True)."""
+    if trigger.kind != TriggerKind.ON_READ_OF_ANY:
+        return None
+    assert trigger.entity_ids is not None
+    return next((v for v in call.args.values() if v in trigger.entity_ids), None)
+
+
+def injection_was_triggered(applied: AppliedInjection, steps: list["Step"]) -> bool:
+    """True if the Target had a chance to perceive this FIRED injection's
+    effect.
+
+    TOOL_ERROR now only ever fires via interception (the arm/trigger model
+    guarantees this -- see trigger_matches_pre_dispatch), so any TOOL_ERROR
+    in the fired list was, by construction, actually encountered.
+    World-mutation kinds still need at least one more tool call after
+    firing to be observable at all -- the Target only perceives world
+    state through tool calls, it has no push channel.
+    """
+    if applied.injection.kind == InjectionKind.TOOL_ERROR:
+        return True
+    return any(s.index > applied.fired_at_step for s in steps)
+
+
+def check_pending_tool_error(pending_tool_errors: dict[str, int], tool_name: str) -> str | None:
+    """Consumes one pending injected error for `tool_name`, if any. Returns
+    the synthetic error message, or None if the call should proceed
+    normally."""
+    remaining = pending_tool_errors.get(tool_name, 0)
+    if remaining <= 0:
+        return None
+    pending_tool_errors[tool_name] = remaining - 1
+    return "HTTP 500 Internal Server Error: upstream service temporarily unavailable, please retry"