Context Engineering

# script_id: day_049_context_engineering/drowning_agent
# Replays a Day 48-style 6-step run on Orbit's longest issue and prints what each step re-sends.
# Install: uv add tiktoken   Env: ORBIT_DATA (default datasets/orbit)   Services: none (offline)
import os, sqlite3
from pathlib import Path
import tiktoken

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
PRICE_IN = 0.10 / 1_000_000          # gpt-6-luna input, $/token (as of 2026-11; verify at the pricing page)
enc = tiktoken.get_encoding("o200k_base")   # approximate for GPT-6 (Day 4)

db = sqlite3.connect(f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro", uri=True)
# Pick the long-log issue by property (most text), never by a hard-coded id
issue_id, body = db.execute("SELECT id, body FROM issues ORDER BY length(body) DESC LIMIT 1").fetchone()
comments = [c for (c,) in db.execute("SELECT body FROM comments WHERE issue_id = ? ORDER BY position", (issue_id,))]
issue_tokens = len(enc.encode("\n\n".join([body, *comments])))

# What each step ADDS to the history (call item + result); step 1 starts with instructions, tools, task
steps = [("search_issues", 400), (f"get_issue({issue_id})", issue_tokens), ("search_docs", 300),
         ("get_doc_section", 1_200), ("search_issues", 400), ("(answer)", 0)]
history, billed = 1_500, 0
print(f"issue #{issue_id}: {issue_tokens:,} tokens of logs\n")
print(f"{'step':<6}{'tool called':<18}{'input sent':>11}{'billed so far':>15}")
for n, (tool, added) in enumerate(steps, 1):
    billed += history                     # every call re-sends the whole history
    print(f"{n:<6}{tool:<18}{history:>11,}{billed:>15,}")
    history += added + 150                # +150 for the call item and a reasoning item
print(f"\ninput billed: {billed:,} tokens (≈ ${billed * PRICE_IN:.4f}); the log was re-sent {len(steps) - 2} times")

issue #227: 32,867 tokens of logs

step  tool called        input sent  billed so far
1     search_issues           1,500          1,500
2     get_issue(227)          2,050          3,550
3     search_docs            35,067         38,617
4     get_doc_section        35,517         74,134
5     search_issues          36,867        111,001
6     (answer)               37,417        148,418

input billed: 148,418 tokens (≈ $0.0148); the log was re-sent 4 times

# script_id: day_049_context_engineering/context_budget
# The smallest useful context budget: count what you send, keep a reserve. Offline.
# Install: uv add tiktoken   Env: ORBIT_DATA (default datasets/orbit)   Services: none
import json, os, sqlite3
from dataclasses import dataclass
from pathlib import Path
import tiktoken

enc = tiktoken.get_encoding("o200k_base")       # approximate for GPT-6 (Day 4)

@dataclass
class ContextBudget:
    window: int                                  # YOUR working limit, not the model maximum
    output_reserve: int                          # room for reasoning + the answer

    @property
    def available_input(self) -> int:
        return self.window - self.output_reserve

    def used(self, items: list[dict]) -> int:    # serialize each item, count its tokens
        return sum(len(enc.encode(json.dumps(i, ensure_ascii=False))) for i in items)

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
db = sqlite3.connect(f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro", uri=True)
iid, body = db.execute("SELECT id, body FROM issues ORDER BY length(body) DESC LIMIT 1").fetchone()
items = [{"role": "user", "content": f"Investigate Orbit issue #{iid}."},
         {"type": "function_call", "call_id": "c1", "name": "get_issue", "arguments": json.dumps({"issue_id": iid})},
         {"type": "function_call_output", "call_id": "c1", "output": body}]
b = ContextBudget(window=64_000, output_reserve=25_000)
used = b.used(items)
print(f"used={used:,} reserve={b.output_reserve:,} available_input={b.available_input:,} fits={used <= b.available_input}")

used=32,790 reserve=25,000 available_input=39,000 fits=True

# script_id: day_049_context_engineering/live_checks
# Exact vs approximate token counts, then two ways to break an agent's context. Needs a key.
# Install: uv add "openai>=3.46,<4" tiktoken   Env: OPENAI_API_KEY, ATLAS_MODEL, ORBIT_DATA
import json, os, sqlite3
from pathlib import Path
import openai, tiktoken
from openai import OpenAI

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
client = OpenAI()
enc = tiktoken.get_encoding("o200k_base")
GET_ISSUE = {"type": "function", "name": "get_issue", "description": "Full text of one Orbit issue.",
             "strict": True, "parameters": {"type": "object", "additionalProperties": False,
             "required": ["issue_id"], "properties": {"issue_id": {"type": "integer"}}}}

db = sqlite3.connect(f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro", uri=True)
iid, body = db.execute("SELECT id, body FROM issues ORDER BY length(body) DESC LIMIT 1").fetchone()
items = [{"role": "user", "content": f"Investigate Orbit issue #{iid}."},
         {"type": "function_call", "call_id": "c1", "name": "get_issue", "arguments": json.dumps({"issue_id": iid})},
         {"type": "function_call_output", "call_id": "c1", "output": body}]
approx = sum(len(enc.encode(json.dumps(i, ensure_ascii=False))) for i in items)
exact = client.responses.input_tokens.count(model=MODEL, input=items, tools=[GET_ISSUE]).input_tokens
print(f"approx={approx:,} exact={exact:,} diff={exact - approx:+,}")

approx=32,790 exact=… diff=…      # exact needs a key; expect them to differ by a few percent

# script_id: day_049_context_engineering/live_checks
# (a) No reserve: a tiny output cap on a big input. Reasoning eats the cap before any text appears.
ask = items + [{"role": "user", "content": "Diagnose this issue."}]
r = client.responses.create(model=MODEL, input=ask, tools=[GET_ISSUE], max_output_tokens=200)
reason = r.incomplete_details.reason if r.incomplete_details else None
print(f"(a) status={r.status} reason={reason} text={r.output_text[:60]!r}")

# (b) Delete the call but keep its output: what naive "drop the oldest items" code does
orphan = [items[0], items[2], {"role": "user", "content": "Diagnose this issue."}]
try:
    client.responses.create(model=MODEL, input=orphan, tools=[GET_ISSUE])
except openai.BadRequestError as e:
    print(f"(b) HTTP {e.status_code}: {e.message[:120]}")
else:
    print("(b) accepted: no 400 on this model, so the pairing rule still applies (⚠ verify)")

(a) status=incomplete reason=max_output_tokens text=''   # typical; depends on reasoning effort
(b) HTTP 400: …                     # message text not captured here (run without a key); ⚠ verify

# script_id: day_049_context_engineering/atlas_context
# atlas/context.py: budget, clear, truncate-by-reference and compact an agent's item list.
# Install: uv add tiktoken   Env: ORBIT_DATA (demo only)   Services: none. Runs offline:
# compaction takes any summarize(text) -> str callable, so the model call lives in agent code.
from __future__ import annotations
import copy, json, os, sqlite3
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable

import tiktoken

_ENC = tiktoken.get_encoding("o200k_base")   # approximate for GPT-6; exact: responses.input_tokens.count


def as_dict(item: Any) -> dict:
    """SDK items are Pydantic models, ours are dicts. Always return a copy we may edit."""
    if hasattr(item, "model_dump"):
        return item.model_dump(exclude_none=True)
    return copy.deepcopy(dict(item))


def count_tokens(text: str) -> int:          # built on Day 4
    return len(_ENC.encode(text))


def item_tokens(item: Any) -> int:
    return count_tokens(json.dumps(as_dict(item), ensure_ascii=False))


@dataclass(frozen=True)
class ContextBudget:
    window: int = 64_000           # your working limit (cost + quality), not the model's 1.05M
    output_reserve: int = 25_000   # reasoning + answer; pass it as max_output_tokens

    def __post_init__(self) -> None:
        if not 0 < self.output_reserve < self.window:
            raise ValueError("need 0 < output_reserve < window")

    @property
    def available_input(self) -> int:
        return self.window - self.output_reserve

    def used(self, items: Iterable[Any]) -> int:
        return sum(item_tokens(i) for i in items)

    def fits(self, items: Iterable[Any]) -> bool:
        return self.used(items) <= self.available_input

# script_id: day_049_context_engineering/atlas_context
CLEARED = "[cleared: {name}({args}) returned ~{n:,} tokens; call {name} again to re-read]"


def _text(output: Any) -> str:
    return output if isinstance(output, str) else json.dumps(output, ensure_ascii=False)


def clear_tool_results(items: Iterable[Any], keep_last: int = 2, min_tokens: int = 200) -> list[dict]:
    """Replace the text of all but the last `keep_last` tool results with a placeholder.
    Every item and call_id survives, so each function_call keeps its output (no 400)."""
    out = [as_dict(i) for i in items]
    calls = {i["call_id"]: i for i in out if i.get("type") == "function_call"}
    outputs = [i for i in out if i.get("type") == "function_call_output"]
    for item in outputs[:-keep_last] if keep_last else outputs:
        text = _text(item["output"])
        if text.startswith("[cleared:") or count_tokens(text) < min_tokens:
            continue                                  # already cleared, or too small to matter
        call = calls.get(item["call_id"], {})
        item["output"] = CLEARED.format(name=call.get("name", "tool"),
                                        args=call.get("arguments", ""), n=count_tokens(text))
    return out


def pairing_ok(items: Iterable[Any]) -> bool:
    """Every function_call has exactly one output and every output has its call."""
    d = [as_dict(i) for i in items]
    calls = sorted(i["call_id"] for i in d if i.get("type") == "function_call")
    outs = sorted(i["call_id"] for i in d if i.get("type") == "function_call_output")
    return calls == outs

# script_id: day_049_context_engineering/atlas_context
class ResultStore:
    """Full tool results live here; the model sees a capped first page plus a reference."""

    def __init__(self, max_tokens: int = 1_500) -> None:
        self.max_tokens = max_tokens
        self._blobs: dict[str, list[int]] = {}       # ref -> token ids of the full text

    def cap(self, tool: str, text: str) -> str:
        toks = _ENC.encode(text)
        if len(toks) <= self.max_tokens:
            return text
        ref = f"{tool}-{len(self._blobs) + 1}"
        self._blobs[ref] = toks
        rest = len(toks) - self.max_tokens
        return (_ENC.decode(toks[: self.max_tokens]) + f"\n[truncated: {rest:,} more tokens; "
                f'use read_more(ref="{ref}", offset={self.max_tokens})]')

    def read_more(self, ref: str, offset: int) -> str:
        toks = self._blobs.get(ref)
        if toks is None:                              # an error the model can act on (Day 37)
            return f"error: unknown ref {ref!r}; use a ref from a [truncated: ...] note"
        offset = max(0, offset)                       # a model can send a negative offset
        end = offset + self.max_tokens
        tail = f'\n[more: use read_more(ref="{ref}", offset={end})]' if end < len(toks) else "\n[end]"
        return _ENC.decode(toks[offset:end]) + tail


READ_MORE_TOOL = {
    "type": "function", "name": "read_more", "strict": True,
    "description": "Read the next page of a truncated tool result. Use only a ref and offset "
                   "taken from a [truncated: ...] or [more: ...] note.",
    "parameters": {"type": "object", "additionalProperties": False, "required": ["ref", "offset"],
                   "properties": {"ref": {"type": "string"}, "offset": {"type": "integer"}}},
}

# script_id: day_049_context_engineering/atlas_context
SUMMARY_PROMPT = ("Summarize these earlier agent steps for the agent itself. Copy issue numbers, doc ids, "
                  "versions, config keys and error codes VERBATIM. Say what was found, what was ruled "
                  "out and what is still open. At most 200 words.")


def render(items: list[dict]) -> str:
    """Plain-text transcript for the summarizer. Reasoning items are opaque, so they are skipped."""
    lines = []
    for i in items:
        kind = i.get("type", "message")
        if kind == "function_call":
            lines.append(f"CALL {i['name']}({i['arguments']})")
        elif kind == "function_call_output":
            lines.append(f"RESULT: {_text(i['output'])}")
        elif kind == "message":
            c = i.get("content")
            text = c if isinstance(c, str) else " ".join(p.get("text", "") for p in c or [])
            lines.append(f"{i.get('role', '?').upper()}: {text}")
    return "\n".join(lines)


def safe_split(items: list[dict], keep_recent: int) -> int:
    """Start of the recent tail, moved left so no call is separated from its output or reasoning."""
    cut = max(1, len(items) - max(1, keep_recent))
    while cut > 1 and (items[cut].get("type") == "function_call_output"
                       or items[cut - 1].get("type") in ("reasoning", "function_call")):
        cut -= 1
    return cut


def compact(items: Iterable[Any], summarize: Callable[[str], str], keep_recent: int = 6) -> list[dict]:
    """DIY compaction: the task (items[0]) + one summary note + the recent tail."""
    out = [as_dict(i) for i in items]
    if any(i.get("type") == "compaction" for i in out):
        return out                       # server-side compaction owns this transcript; don't mix
    cut = safe_split(out, keep_recent)
    if cut <= 1:
        return out                       # nothing old enough to summarize
    note = summarize(render(out[1:cut]))
    return [out[0], {"role": "user", "content": f"[Notes on earlier steps]\n{note}"}, *out[cut:]]

# script_id: day_049_context_engineering/atlas_context
class ContextManager:
    """Keeps an agent's input under budget: clear old results first, compact only if still over."""

    def __init__(self, budget: ContextBudget | None = None, *, keep_last: int = 2,
                 max_tool_tokens: int | None = 1_500, keep_recent: int = 6,
                 summarize: Callable[[str], str] | None = None) -> None:
        self.budget = budget or ContextBudget()
        self.keep_last, self.keep_recent, self.summarize = keep_last, keep_recent, summarize
        self.store = ResultStore(max_tool_tokens) if max_tool_tokens else None

    def cap(self, tool: str, text: str) -> str:          # wrap every tool result with this
        return self.store.cap(tool, text) if self.store else text

    def tools(self) -> list[dict]:                        # extra tool specs this manager needs
        return [READ_MORE_TOOL] if self.store else []

    def clear_tool_results(self, items: Iterable[Any], keep_last: int | None = None) -> list[dict]:
        return clear_tool_results(items, self.keep_last if keep_last is None else keep_last)

    def compact(self, items: Iterable[Any]) -> list[dict]:
        if self.summarize is None:                        # no summarizer configured: clear-only mode
            return [as_dict(i) for i in items]
        return compact(items, self.summarize, self.keep_recent)

    def fit(self, items: Iterable[Any]) -> list[dict]:
        """Return a NEW list that fits when possible. Never mutates `items`."""
        out = [as_dict(i) for i in items]
        if self.budget.fits(out):
            return out
        out = self.clear_tool_results(out)
        if self.budget.fits(out):
            return out
        return self.compact(out)        # may still be over: run_agent's max_tokens stop is the backstop

# script_id: day_049_context_engineering/atlas_context
def orbit_transcript(data_dir: Path | None = None, n_docs: int = 3) -> list[dict]:
    """An offline agent transcript: a search, the long-log issue, then its component's docs."""
    data_dir = data_dir or Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
    db = sqlite3.connect(f"file:{data_dir / 'issues' / 'orbit.db'}?mode=ro", uri=True)
    iid, title, body = db.execute(
        "SELECT id, title, body FROM issues ORDER BY length(body) DESC LIMIT 1").fetchone()
    comp = db.execute("SELECT l.value FROM labels l JOIN issue_labels il ON il.label_id = l.id "
                      "WHERE il.issue_id = ? AND l.kind = 'comp'", (iid,)).fetchone()[0]
    db.close()
    docs = list(dict.fromkeys([data_dir / "docs" / f"{comp}.md", *sorted((data_dir / "docs").glob("*.md"))]))
    calls = [("search_issues", {"query": title[:40]}, "[5 short search results]"), ("get_issue", {"issue_id": iid}, body)]
    calls += [("get_doc_section", {"doc_id": p.stem, "heading": ""}, p.read_text(encoding="utf-8")) for p in docs[:n_docs]]
    items: list[dict] = [{"role": "user", "content": f"Investigate Orbit issue #{iid}: duplicate? docs? cause?"}]
    for n, (name, args, result) in enumerate(calls, 1):
        items += [{"type": "function_call", "call_id": f"call_{n}", "name": name, "arguments": json.dumps(args)},
                  {"type": "function_call_output", "call_id": f"call_{n}", "output": result}]
    return items


if __name__ == "__main__":
    import re
    items = orbit_transcript()
    fake = lambda text: "kept " + ", ".join(sorted(set(re.findall(r"ORB-\d{4}|#\d+", text))))  # no model call
    cm = ContextManager(ContextBudget(window=28_000, output_reserve=25_000), keep_recent=2, summarize=fake)
    b = cm.budget
    print(f"start:     {len(items)} items, {b.used(items):,} tokens (available_input={b.available_input:,})")
    cleared = cm.clear_tool_results(items)
    print(f"cleared:   {b.used(cleared):,} tokens, pairing_ok={pairing_ok(cleared)}, original still {b.used(items):,}")
    fitted = cm.fit(items)
    note = next((i["content"] for i in fitted if str(i.get("content", "")).startswith("[Notes")), "(none)")
    print(f"fit():     {len(fitted)} items, {b.used(fitted):,} tokens, pairing_ok={pairing_ok(fitted)}")
    print(f"note:      {note.splitlines()[-1]}")
    capped = cm.cap("get_issue", items[4]["output"])
    print(f"capped:    {count_tokens(capped):,} tokens, ends {capped.splitlines()[-1]}")
    page = cm.store.read_more(ref="get_issue-1", offset=1_500)
    print(f"read_more: {count_tokens(page):,} tokens, ends {page.splitlines()[-1]}")

start:     11 items, 37,518 tokens (available_input=3,000)
cleared:   3,484 tokens, pairing_ok=True, original still 37,518
fit():     4 items, 1,667 tokens, pairing_ok=True
note:      kept ORB-1201, ORB-1202, ORB-1501, ORB-1502, ORB-1503, ORB-1601
capped:    1,526 tokens, ends [truncated: 30,914 more tokens; use read_more(ref="get_issue-1", offset=1500)]
read_more: 1,518 tokens, ends [more: use read_more(ref="get_issue-1", offset=3000)]

# script_id: day_049_context_engineering/server_compaction
# Server-side compaction: inline in responses.create, then the standalone compact endpoint.
# Install: uv add "openai>=3.46,<4"   Env: OPENAI_API_KEY, ATLAS_MODEL, ORBIT_DATA, ATLAS_COMPACT_THRESHOLD
import os, sqlite3
from pathlib import Path
import openai
from openai import OpenAI

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
# The guide's examples use 200_000 and state no minimum (⚠ verify). Low here so a demo triggers it;
# if the API rejects it, raise ATLAS_COMPACT_THRESHOLD rather than editing code.
THRESHOLD = int(os.environ.get("ATLAS_COMPACT_THRESHOLD", "20000"))
DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
client = OpenAI()

db = sqlite3.connect(f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro", uri=True)
iid, body = db.execute("SELECT id, body FROM issues ORDER BY length(body) DESC LIMIT 1").fetchone()
items: list = []                                  # a long transcript: the log pasted in parts
for n, start in enumerate(range(0, len(body), 12_000), 1):
    items += [{"role": "user", "content": f"Log part {n} of issue #{iid}:\n{body[start:start + 12_000]}"},
              {"role": "assistant", "content": "Noted."}]
items.append({"role": "user", "content": "Which error code appears, and what do the two symptoms share?"})

try:
    resp = client.responses.create(model=MODEL, input=items, store=False,  # store=False: ZDR-friendly
                                   context_management=[{"type": "compaction", "compact_threshold": THRESHOLD}])
except openai.BadRequestError as e:
    raise SystemExit(f"rejected ({e.message[:120]}); try a larger ATLAS_COMPACT_THRESHOLD")
print("output types:", [i.type for i in resp.output], "| input tokens:", resp.usage.input_tokens)
items += resp.output
latest = max((n for n, i in enumerate(items) if getattr(i, "type", None) == "compaction"), default=None)
if latest is None:
    print(f"no compaction item: the context stayed under {THRESHOLD:,} tokens")
else:   # the guide allows dropping items before the latest compaction item, whatever the order in output
    items = items[latest:]

window = client.responses.compact(model=MODEL, input=items)   # you pick the moment; don't prune its output
nxt = client.responses.create(model=MODEL, store=False,
                              input=[*window.output, {"role": "user", "content": "One-line fix suggestion?"}])
print(len(window.output), "items in the compacted window |", nxt.output_text[:200])

output types: [… 'compaction' …, 'message'] | input tokens: …   # shape only; item order not verified
… items in the compacted window | …

# script_id: day_049_context_engineering/atlas_agent
# Compact Day 48 agent + context= hook. Install: uv add "openai>=3.46,<4" tiktoken
# Env: OPENAI_API_KEY, ATLAS_MODEL (default gpt-6-luna), ORBIT_DATA   Services: none
# built on Day 48 (run_agent, AgentConfig, AgentResult) and Days 35-36 (tools); full versions there
from __future__ import annotations
import json, os, re, sqlite3, time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from openai import OpenAI
from atlas_context import SUMMARY_PROMPT, ContextBudget, ContextManager

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
DB = f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro"          # read-only: the agent never writes
client = OpenAI()


def _q(sql: str, *args) -> list[tuple]:
    db = sqlite3.connect(DB, uri=True)
    try:
        return db.execute(sql, args).fetchall()
    finally:
        db.close()


def search_issues(query: str) -> str:
    words = re.findall(r"\w+", query.lower())[:4] or [""]
    where = " AND ".join("lower(title || ' ' || body) LIKE ?" for _ in words)
    rows = _q(f"SELECT id, title, state, duplicate_of FROM issues WHERE {where} LIMIT 5", *[f"%{w}%" for w in words])
    return json.dumps([dict(zip(("id", "title", "state", "duplicate_of"), r)) for r in rows])


def get_issue(issue_id: int) -> str:              # returns everything: this is the context hog
    row = _q("SELECT title, body FROM issues WHERE id = ?", issue_id)
    if not row:
        return f"error: no issue {issue_id}; use search_issues to find ids"
    comments = [b for (b,) in _q("SELECT body FROM comments WHERE issue_id = ? ORDER BY position", issue_id)]
    return "\n\n".join([f"# {row[0][0]}", row[0][1], *comments])

# script_id: day_049_context_engineering/atlas_agent
def search_docs(query: str) -> str:               # keyword stand-in for R2's index (Day 34)
    words = set(re.findall(r"\w+", query.lower()))
    scored = [(sum(p.read_text(encoding="utf-8").lower().count(w) for w in words), p.stem)
              for p in (DATA_DIR / "docs").glob("*.md")]
    return json.dumps([doc for n, doc in sorted(scored, reverse=True)[:3] if n])


def get_doc_section(doc_id: str, heading: str) -> str:
    path = DATA_DIR / "docs" / f"{Path(doc_id).name}.md"        # .name blocks ../ traversal
    if not path.exists():
        return f"error: no doc {doc_id!r}; use search_docs first"
    text = path.read_text(encoding="utf-8")
    parts = [p for p in re.split(r"(?m)^(?=#{1,3} )", text) if heading.lower() in p[:120].lower()]
    return parts[0] if heading and parts else text


def _spec(fn: Callable, desc: str, **props: str) -> dict:
    return {"type": "function", "name": fn.__name__, "description": desc, "strict": True,
            "parameters": {"type": "object", "additionalProperties": False, "required": list(props),
                           "properties": {k: {"type": t} for k, t in props.items()}}}


REGISTRY = {f.__name__: f for f in (search_issues, get_issue, search_docs, get_doc_section)}
TOOLS = [_spec(search_issues, "Keyword search over past Orbit issues: ids, titles, duplicate_of.", query="string"),
         _spec(get_issue, "Full text of one Orbit issue with comments. Can be very long.", issue_id="integer"),
         _spec(search_docs, "Find up to 3 Orbit doc ids for a keyword query.", query="string"),
         _spec(get_doc_section, 'Read a doc section by heading text; heading="" returns the whole page.',
               doc_id="string", heading="string")]
AGENT_SYSTEM = ("You are Atlas, investigating Orbit issues for maintainers. Use the tools; never guess. "
                "Finish with: likely duplicate (issue number or none), relevant doc ids, likely cause, "
                "what is still unknown. Copy error codes and issue numbers exactly.")
PRICE_IN, PRICE_CACHED, PRICE_OUT = 0.10e-6, 0.01e-6, 0.50e-6  # gpt-6-luna $/token, as of 2026-11; verify

# script_id: day_049_context_engineering/atlas_agent
@dataclass
class AgentConfig:                                 # Day 48 defaults
    max_steps: int = 8
    max_seconds: float = 60
    max_tokens: int = 100_000                      # cumulative input + output across steps
    max_errors: int = 3


@dataclass
class StepRecord:
    index: int
    tool_calls: list[str]
    errors: int
    input_tokens: int
    output_tokens: int
    ms: int


@dataclass
class AgentResult:
    text: str
    stop_reason: str
    steps: list[StepRecord]
    usage: dict
    cost_usd: float
    partial: bool


def run_tool(name: str, arguments: str, context: ContextManager | None) -> str:
    try:
        args = json.loads(arguments)
        if name == "read_more" and context and context.store:
            return context.store.read_more(**args)
        out = REGISTRY[name](**args)
    except Exception as e:                         # an envelope the model can read, not a crash (Day 37)
        return f"error: {type(e).__name__}: {e}"
    return context.cap(name, out) if context else out


def _cost(u) -> float:
    cached = u.input_tokens_details.cached_tokens
    return (u.input_tokens - cached) * PRICE_IN + cached * PRICE_CACHED + u.output_tokens * PRICE_OUT

# script_id: day_049_context_engineering/atlas_agent
def run_agent(task: str, *, tools: list[dict] = TOOLS, system: str = AGENT_SYSTEM,
              config: AgentConfig = AgentConfig(), context: ContextManager | None = None) -> AgentResult:
    items: list = [{"role": "user", "content": task}]
    specs = tools + (context.tools() if context else [])
    extra = {"max_output_tokens": context.budget.output_reserve} if context else {}
    steps, usage, cost, errors, t0 = [], {"input_tokens": 0, "output_tokens": 0}, 0.0, 0, time.monotonic()
    stop = "max_steps"
    for n in range(1, config.max_steps + 1):
        if context:
            items = context.fit(items)             # clear / compact BEFORE the call, on a copy
        t = time.monotonic()
        resp = client.responses.create(model=MODEL, instructions=system, input=items, tools=specs, **extra)
        usage["input_tokens"] += resp.usage.input_tokens
        usage["output_tokens"] += resp.usage.output_tokens
        cost += _cost(resp.usage)
        items += resp.output                       # keep reasoning + calls: pairing!
        calls = [i for i in resp.output if i.type == "function_call"]
        outs = [run_tool(c.name, c.arguments, context) for c in calls]
        items += [{"type": "function_call_output", "call_id": c.call_id, "output": o} for c, o in zip(calls, outs)]
        bad = sum(o.startswith("error:") for o in outs) + (resp.status == "incomplete")
        errors = errors + bad if bad else 0        # consecutive error steps
        steps.append(StepRecord(n, [c.name for c in calls], bad, resp.usage.input_tokens,
                                resp.usage.output_tokens, int((time.monotonic() - t) * 1000)))
        if not calls and resp.status == "completed":
            return AgentResult(resp.output_text, "done", steps, usage, cost, False)
        if errors >= config.max_errors:
            stop = "max_errors"; break
        if usage["input_tokens"] + usage["output_tokens"] > config.max_tokens:
            stop = "max_tokens"; break
        if time.monotonic() - t0 > config.max_seconds:
            stop = "max_seconds"; break
    # Budget stop: one wrap-up call with no tools, so the run still ends with something usable (Day 48)
    items = context.fit(items) if context else items
    wrap = client.responses.create(model=MODEL, instructions=system, tools=specs, tool_choice="none", **extra,
                                   input=[*items, {"role": "user", "content": "Stop. State what you found and what is missing."}])
    return AgentResult(wrap.output_text, stop, steps, usage, cost + _cost(wrap.usage), True)

# script_id: day_049_context_engineering/atlas_agent
def make_summarizer(model: str = MODEL, prompt: str = SUMMARY_PROMPT) -> Callable[[str], str]:
    """One cheap call that turns old steps into a note. Plug into ContextManager(summarize=...)."""
    def summarize(transcript: str) -> str:
        r = client.responses.create(model=model, instructions=prompt, input=transcript,
                                    reasoning={"effort": "low"})
        return r.output_text
    return summarize


def long_issue_id() -> int:                        # by property, never hard-coded
    return _q("SELECT id FROM issues ORDER BY length(body) DESC LIMIT 1")[0][0]


if __name__ == "__main__":
    task = f"Investigate Orbit issue #{long_issue_id()}: duplicate? what do the docs say? likely cause?"
    for label, ctx in [("baseline", None), ("context", ContextManager(summarize=make_summarizer()))]:
        r = run_agent(task, context=ctx)
        print(f"{label:<9} stop={r.stop_reason} steps={len(r.steps)} "
              f"input={r.usage['input_tokens']:,} cost=${r.cost_usd:.4f}")
        print("          input per step:", [s.input_tokens for s in r.steps])

baseline  stop=… steps=… input=… cost=$…        # shape only: run it to get your numbers
          input per step: […]                   # look for the jump after get_issue
context   stop=… steps=… input=… cost=$…
          input per step: […]                   # expect flatter: the log arrives as 1,500-token pages

# script_id: day_049_context_engineering/measure_context
# Scoreboard row D49: baseline vs clear vs clear+truncate vs +compaction. Needs a key.
# Install: uv add "openai>=3.46,<4" tiktoken   Env: OPENAI_API_KEY, ATLAS_MODEL, ORBIT_DATA
# Cost: ≈ $0.10 on gpt-6-luna for 6 tasks x 4 variants (as of 2026-11; verify at the pricing page)
import json, re
from pathlib import Path

from atlas_agent import DATA_DIR, make_summarizer, run_agent
from atlas_context import ContextBudget, ContextManager


def fallback_set(issues: dict[int, dict]) -> list[int]:
    """Day 48's properties: traceback bug, question, duplicate, vague, docs."""
    rules = [lambda r: "Traceback" in r["body"], lambda r: "type:question" in r["labels"],
             lambda r: r["duplicate_of"] is not None, lambda r: len(r["body"]) < 400,
             lambda r: "type:docs" in r["labels"]]
    return [next(i for i, r in sorted(issues.items()) if rule(r)) for rule in rules]


def load_cases() -> list[dict]:
    rows = (json.loads(line) for line in (DATA_DIR / "issues" / "issues.jsonl").open(encoding="utf-8"))
    issues = {r["id"]: r for r in rows}
    saved = Path("agent_set.json")                 # written by Day 48's pick_agent_set()
    # Day 48 saves a list of {"issue_id", "kind", "task"}; otherwise pick five by the same properties
    ids = [int(p["issue_id"]) for p in json.loads(saved.read_text())] if saved.exists() else fallback_set(issues)
    long_id = max(issues, key=lambda i: len(issues[i]["body"]))
    return [issues[i] for i in dict.fromkeys([long_id, *ids])]


def passes(issue: dict, answer: str) -> bool:
    codes = re.findall(r"ORB-\d{4}", issue["body"])
    want = [min(sorted(set(codes)), key=codes.count)] if codes else []   # rarest code; sorted() breaks ties the same way every run
    ok = all(w in answer for w in want)
    dup = issue["duplicate_of"]
    return ok and (dup is None or str(dup) in re.findall(r"\d+", answer))

# script_id: day_049_context_engineering/measure_context
def variant(name: str) -> ContextManager | None:   # a fresh manager per run: its ResultStore is per run
    b = ContextBudget(window=64_000, output_reserve=25_000)
    return {"baseline": None,
            "clear": ContextManager(b, max_tool_tokens=None),
            "clear+truncate": ContextManager(b),
            # a tight window forces compaction to actually fire on these short runs
            "+compaction": ContextManager(ContextBudget(40_000, 25_000), summarize=make_summarizer())}[name]


if __name__ == "__main__":
    cases = load_cases()
    print(f"{'variant':<16}{'steps':>6}{'input tokens':>14}{'$':>9}{'checks':>8}")
    for name in ["baseline", "clear", "clear+truncate", "+compaction"]:
        steps = tokens = passed = 0
        cost = 0.0
        with open(f"answers_{name}.jsonl", "w", encoding="utf-8") as log:
            for issue in cases:
                r = run_agent(f"Investigate Orbit issue #{issue['id']}: duplicate? docs? likely cause?",
                              context=variant(name))
                steps, tokens, cost = steps + len(r.steps), tokens + r.usage["input_tokens"], cost + r.cost_usd
                passed += passes(issue, r.text)
                log.write(json.dumps({"id": issue["id"], "stop": r.stop_reason, "answer": r.text}) + "\n")
        print(f"{name:<16}{steps:>6}{tokens:>14,}{cost:>9.4f}{passed:>5}/{len(cases)}")

variant          steps  input tokens        $  checks
baseline             …             …        …    …/6      # shape only: your numbers will differ
clear                …             …        …    …/6
clear+truncate       …             …        …    …/6
+compaction          …             …        …    …/6