atif-sql 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (164) hide show
  1. atif_analytics/__init__.py +9 -0
  2. atif_analytics/application/__init__.py +3 -0
  3. atif_analytics/application/analyze.py +158 -0
  4. atif_analytics/application/prompts.py +750 -0
  5. atif_analytics/application/use_cases/__init__.py +3 -0
  6. atif_analytics/application/use_cases/_shared.py +264 -0
  7. atif_analytics/application/use_cases/classify.py +430 -0
  8. atif_analytics/application/use_cases/conflicts.py +448 -0
  9. atif_analytics/application/use_cases/friction.py +634 -0
  10. atif_analytics/application/use_cases/perceived.py +451 -0
  11. atif_analytics/domain/__init__.py +3 -0
  12. atif_analytics/domain/authorship.py +166 -0
  13. atif_analytics/domain/config.py +30 -0
  14. atif_analytics/domain/costs.py +50 -0
  15. atif_analytics/domain/friction.py +117 -0
  16. atif_analytics/domain/layout.py +94 -0
  17. atif_analytics/domain/models.py +385 -0
  18. atif_analytics/domain/ports.py +100 -0
  19. atif_analytics/domain/transcript.py +359 -0
  20. atif_analytics/infrastructure/__init__.py +3 -0
  21. atif_analytics/infrastructure/corpus_reader.py +546 -0
  22. atif_analytics/infrastructure/parquet_cache.py +186 -0
  23. atif_analytics/infrastructure/settings.py +114 -0
  24. atif_analytics/infrastructure/sqlite_state/__init__.py +3 -0
  25. atif_analytics/infrastructure/sqlite_state/checkpointer.py +326 -0
  26. atif_analytics/infrastructure/sqlite_state/retry_queue.py +296 -0
  27. atif_analytics/py.typed +0 -0
  28. atif_cli/__init__.py +3 -0
  29. atif_cli/__main__.py +8 -0
  30. atif_cli/app.py +2459 -0
  31. atif_cli/converter_adapter.py +151 -0
  32. atif_cli/corpus.py +409 -0
  33. atif_cli/cron.py +311 -0
  34. atif_cli/duck_errors.py +97 -0
  35. atif_cli/embed_lake.py +101 -0
  36. atif_cli/errors.py +89 -0
  37. atif_cli/lake.py +373 -0
  38. atif_cli/lake_sessions.py +270 -0
  39. atif_cli/output.py +297 -0
  40. atif_cli/py.typed +0 -0
  41. atif_converter/__init__.py +3 -0
  42. atif_converter/application/__init__.py +3 -0
  43. atif_converter/application/convert_and_audit.py +294 -0
  44. atif_converter/application/convert_codex.py +210 -0
  45. atif_converter/domain/__init__.py +9 -0
  46. atif_converter/domain/agents.py +40 -0
  47. atif_converter/domain/atif/__init__.py +55 -0
  48. atif_converter/domain/atif/agent.py +37 -0
  49. atif_converter/domain/atif/content.py +161 -0
  50. atif_converter/domain/atif/final_metrics.py +45 -0
  51. atif_converter/domain/atif/metrics.py +49 -0
  52. atif_converter/domain/atif/observation.py +21 -0
  53. atif_converter/domain/atif/observation_result.py +47 -0
  54. atif_converter/domain/atif/step.py +144 -0
  55. atif_converter/domain/atif/subagent_trajectory_ref.py +90 -0
  56. atif_converter/domain/atif/tool_call.py +36 -0
  57. atif_converter/domain/atif/trajectory.py +211 -0
  58. atif_converter/domain/atif/trajectory_validator.py +299 -0
  59. atif_converter/domain/blobs.py +510 -0
  60. atif_converter/domain/claude_code_conversion.py +1077 -0
  61. atif_converter/domain/codex_conversion.py +1036 -0
  62. atif_converter/domain/codex_edges.py +192 -0
  63. atif_converter/domain/codex_enrichment.py +419 -0
  64. atif_converter/domain/codex_fidelity.py +124 -0
  65. atif_converter/domain/codex_nested_calls.py +165 -0
  66. atif_converter/domain/edges.py +113 -0
  67. atif_converter/domain/enrichment.py +412 -0
  68. atif_converter/domain/errors.py +60 -0
  69. atif_converter/domain/fidelity.py +171 -0
  70. atif_converter/domain/model_prices.json +8538 -0
  71. atif_converter/domain/pricing.py +790 -0
  72. atif_converter/domain/result_signals.py +774 -0
  73. atif_converter/domain/schema_version.py +90 -0
  74. atif_converter/domain/session_events.py +408 -0
  75. atif_converter/infrastructure/__init__.py +3 -0
  76. atif_converter/infrastructure/census.py +89 -0
  77. atif_converter/infrastructure/claude_code_converter.py +127 -0
  78. atif_converter/infrastructure/codex_adapter.py +132 -0
  79. atif_converter/infrastructure/codex_census.py +100 -0
  80. atif_converter/infrastructure/codex_converter.py +82 -0
  81. atif_converter/infrastructure/harbor_adapter.py +187 -0
  82. atif_converter/infrastructure/raw_records.py +426 -0
  83. atif_converter/infrastructure/source_archive.py +154 -0
  84. atif_converter/py.typed +0 -0
  85. atif_corpus/__init__.py +13 -0
  86. atif_corpus/application/__init__.py +7 -0
  87. atif_corpus/application/materialize.py +1826 -0
  88. atif_corpus/domain/__init__.py +9 -0
  89. atif_corpus/domain/agents.py +33 -0
  90. atif_corpus/domain/generation.py +34 -0
  91. atif_corpus/domain/layout.py +168 -0
  92. atif_corpus/domain/ports.py +175 -0
  93. atif_corpus/domain/session_id.py +68 -0
  94. atif_corpus/domain/sessions.py +228 -0
  95. atif_corpus/domain/slug.py +72 -0
  96. atif_corpus/domain/source_layout.py +143 -0
  97. atif_corpus/domain/watermark.py +76 -0
  98. atif_corpus/infrastructure/__init__.py +7 -0
  99. atif_corpus/infrastructure/atomic.py +266 -0
  100. atif_corpus/infrastructure/compress_artifacts.py +161 -0
  101. atif_corpus/infrastructure/fake_converter.py +174 -0
  102. atif_corpus/infrastructure/scanner.py +316 -0
  103. atif_corpus/infrastructure/settings.py +130 -0
  104. atif_corpus/infrastructure/source_archive.py +202 -0
  105. atif_corpus/py.typed +0 -0
  106. atif_duck/__init__.py +43 -0
  107. atif_duck/domain/__init__.py +32 -0
  108. atif_duck/domain/artifacts.py +66 -0
  109. atif_duck/domain/authorship.py +165 -0
  110. atif_duck/domain/catalog.py +626 -0
  111. atif_duck/domain/columnar.py +149 -0
  112. atif_duck/domain/embedding_guard.py +88 -0
  113. atif_duck/domain/examples.py +265 -0
  114. atif_duck/domain/lake.py +427 -0
  115. atif_duck/domain/lake_sessions.py +97 -0
  116. atif_duck/domain/raw_readers.py +140 -0
  117. atif_duck/domain/session_id.py +69 -0
  118. atif_duck/domain/sql_literal.py +50 -0
  119. atif_duck/infrastructure/__init__.py +34 -0
  120. atif_duck/infrastructure/analytics.py +469 -0
  121. atif_duck/infrastructure/authorship.py +139 -0
  122. atif_duck/infrastructure/columnar.py +641 -0
  123. atif_duck/infrastructure/lake.py +1383 -0
  124. atif_duck/infrastructure/lake_sessions.py +194 -0
  125. atif_duck/infrastructure/lake_settings.py +63 -0
  126. atif_duck/infrastructure/lake_steps.py +225 -0
  127. atif_duck/infrastructure/projections.py +221 -0
  128. atif_duck/infrastructure/registry.py +2194 -0
  129. atif_duck/infrastructure/stored_artifacts.py +115 -0
  130. atif_duck/py.typed +0 -0
  131. atif_embed/__init__.py +12 -0
  132. atif_embed/application/__init__.py +3 -0
  133. atif_embed/application/embed.py +293 -0
  134. atif_embed/application/prune.py +80 -0
  135. atif_embed/domain/__init__.py +3 -0
  136. atif_embed/domain/discovery.py +70 -0
  137. atif_embed/domain/embedding_guard.py +83 -0
  138. atif_embed/domain/errors.py +94 -0
  139. atif_embed/domain/ports.py +223 -0
  140. atif_embed/domain/sql_literal.py +49 -0
  141. atif_embed/domain/text_stamp.py +71 -0
  142. atif_embed/infrastructure/__init__.py +3 -0
  143. atif_embed/infrastructure/cohere_bedrock.py +433 -0
  144. atif_embed/infrastructure/corpus_text_rows.py +273 -0
  145. atif_embed/infrastructure/gemma_local.py +355 -0
  146. atif_embed/infrastructure/lake_text_rows.py +203 -0
  147. atif_embed/infrastructure/lance_store.py +440 -0
  148. atif_embed/infrastructure/providers.py +30 -0
  149. atif_embed/infrastructure/settings.py +154 -0
  150. atif_embed/py.typed +0 -0
  151. atif_models/__init__.py +14 -0
  152. atif_models/domain/__init__.py +9 -0
  153. atif_models/domain/ports.py +143 -0
  154. atif_models/domain/registry.py +151 -0
  155. atif_models/domain/schema.py +109 -0
  156. atif_models/infrastructure/__init__.py +8 -0
  157. atif_models/infrastructure/openai_bedrock.py +314 -0
  158. atif_models/infrastructure/settings.py +85 -0
  159. atif_models/py.typed +0 -0
  160. atif_sql-0.2.0.dist-info/METADATA +272 -0
  161. atif_sql-0.2.0.dist-info/RECORD +164 -0
  162. atif_sql-0.2.0.dist-info/WHEEL +4 -0
  163. atif_sql-0.2.0.dist-info/entry_points.txt +2 -0
  164. atif_sql-0.2.0.dist-info/licenses/LICENSE +202 -0
@@ -0,0 +1,192 @@
1
+ # SPDX-License-Identifier: Apache-2.0
2
+
3
+ """edges.jsonl for Codex rollouts — the id-keyed skeleton of a rollout.
4
+
5
+ Pure functions over already-parsed record dicts; no harbor, no I/O. The line
6
+ shape is the SAME nine contract fields the Claude Code emitter produces
7
+ (:mod:`atif_converter.domain.edges`), so atif-duck's ``messages`` view reads
8
+ one projection for every agent:
9
+
10
+ {uuid, parent_uuid, message_id, type, ts, is_sidechain,
11
+ is_compact_summary, source_file, tool_use_ids: [..]}
12
+
13
+ Three fields mean something different for a rollout than for a Claude Code
14
+ transcript, and the difference is the point of this module rather than a
15
+ shortcut:
16
+
17
+ ``uuid``
18
+ A rollout record has no session-wide uuid. ``response_item`` payloads
19
+ usually carry an ``id`` (``msg_*`` / ``call_*`` / ``rs_*``), and newer
20
+ Codex builds add a per-file ``ordinal`` to every record — but a
21
+ ``function_call_output`` carries neither, and rollouts written before the
22
+ ordinal landed carry it nowhere. :func:`record_id` therefore falls back to
23
+ a SYNTHETIC ``<type>:L<line>`` id, deterministic in the file's bytes and
24
+ visibly synthetic so a reader cannot mistake it for something Codex
25
+ assigned.
26
+
27
+ ``parent_uuid``
28
+ Always ``None``. A rollout is a flat append-only log with no parent
29
+ pointer: the ordering IS the structure. Inventing a chain by pairing each
30
+ record with its predecessor would produce a tree that looks authoritative
31
+ and describes nothing, so the column stays null and
32
+ ``turn_id``-derived grouping is left to ``codex_turn_id`` on the steps.
33
+
34
+ ``type``
35
+ The raw ROLE where the record has one (``user`` / ``assistant`` /
36
+ ``developer``), and the item kind otherwise (``function_call``,
37
+ ``reasoning``, ``compacted``, …). This is what preserves the
38
+ developer-vs-system distinction that harbor's step conversion flattens
39
+ (:attr:`~atif_converter.domain.codex_fidelity.CodexFidelityGap.DEVELOPER_ROLE_FLATTENED_TO_SYSTEM`).
40
+ """
41
+
42
+ from __future__ import annotations
43
+
44
+ from typing import Any
45
+
46
+ from atif_converter.domain.edges import EDGE_FIELDS, edge_to_jsonl_line
47
+
48
+ #: Payload types whose ``call_id`` is a tool-call join key back to ATIF
49
+ #: ``tool_calls[].tool_call_id`` / ``observation.results[].source_call_id``.
50
+ _CALL_ID_ITEM_TYPES: frozenset[str] = frozenset(
51
+ {
52
+ "function_call",
53
+ "function_call_output",
54
+ "custom_tool_call",
55
+ "custom_tool_call_output",
56
+ "tool_search_call",
57
+ "tool_search_output",
58
+ }
59
+ )
60
+
61
+
62
+ def _payload(record: dict[str, Any]) -> dict[str, Any]:
63
+ payload = record.get("payload")
64
+ return payload if isinstance(payload, dict) else {}
65
+
66
+
67
+ def record_id(record: dict[str, Any], line_number: int) -> str:
68
+ """A stable id for one rollout record: its own, else a synthetic one.
69
+
70
+ Preference order, and each rung exists because the one above it is absent
71
+ on real rollouts (measured 2026-09-11 over 78 local rollouts): the
72
+ payload's own ``id``, then the record's per-file ``ordinal``, then the
73
+ 1-based ``line_number``. The synthetic forms carry the record type so a
74
+ reader can see at a glance that Codex did not assign the id.
75
+ """
76
+ payload = _payload(record)
77
+ own_id = payload.get("id")
78
+ if isinstance(own_id, str) and own_id:
79
+ return own_id
80
+ record_type = record.get("type")
81
+ kind = record_type if isinstance(record_type, str) and record_type else "record"
82
+ ordinal = record.get("ordinal")
83
+ if isinstance(ordinal, int):
84
+ return f"{kind}:{ordinal}"
85
+ return f"{kind}:L{line_number}"
86
+
87
+
88
+ def _edge_type(record: dict[str, Any]) -> str | None:
89
+ """The raw role for a message, else the item/record kind."""
90
+ record_type = record.get("type")
91
+ payload = _payload(record)
92
+ if record_type == "response_item":
93
+ item_type = payload.get("type")
94
+ if item_type == "message":
95
+ role = payload.get("role")
96
+ return role if isinstance(role, str) and role else "message"
97
+ return item_type if isinstance(item_type, str) else None
98
+ if record_type == "event_msg":
99
+ item_type = payload.get("type")
100
+ return f"event:{item_type}" if isinstance(item_type, str) else "event"
101
+ return record_type if isinstance(record_type, str) else None
102
+
103
+
104
+ def _tool_use_ids(record: dict[str, Any]) -> list[str]:
105
+ """The record's ``call_id``, when it is a tool call or a tool output."""
106
+ payload = _payload(record)
107
+ item_type = payload.get("type")
108
+ # A `type` that is an object or a list is unhashable, and the membership test raised
109
+ # `TypeError` on it (found by the fuzzer). Only a string can name a call item.
110
+ if not isinstance(item_type, str) or item_type not in _CALL_ID_ITEM_TYPES:
111
+ return []
112
+ call_id = payload.get("call_id")
113
+ return [call_id] if isinstance(call_id, str) and call_id else []
114
+
115
+
116
+ def codex_edge_from_record(
117
+ record: dict[str, Any],
118
+ source_file: str,
119
+ line_number: int,
120
+ ) -> dict[str, Any]:
121
+ """Project one rollout record onto the contract's edge shape."""
122
+ payload = _payload(record)
123
+ message_id = payload.get("id") if payload.get("type") == "message" else None
124
+ return {
125
+ "uuid": record_id(record, line_number),
126
+ # Flat log, no parent pointer — see the module docstring.
127
+ "parent_uuid": None,
128
+ "message_id": message_id if isinstance(message_id, str) else None,
129
+ "type": _edge_type(record),
130
+ "ts": record.get("timestamp"),
131
+ # Codex sub-agents get their own rollout file rather than an inlined
132
+ # sidechain, so no record of THIS session is a sidechain record.
133
+ "is_sidechain": False,
134
+ "is_compact_summary": record.get("type") == "compacted",
135
+ "source_file": source_file,
136
+ "tool_use_ids": _tool_use_ids(record),
137
+ }
138
+
139
+
140
+ def build_codex_edges(records: list[tuple[dict[str, Any], str]]) -> list[dict[str, Any]]:
141
+ """Build every edge from ``(record, source_file)`` pairs, in FILE order.
142
+
143
+ File order, not timestamp order: a rollout is append-only, so its line
144
+ order is the causal order, and two records inside one millisecond share a
145
+ timestamp often enough that sorting by it would shuffle a tool call past
146
+ its own output. The line number is also what
147
+ :func:`record_id` falls back to, so the two must be counted the same way —
148
+ per source file, 1-based.
149
+ """
150
+ line_numbers: dict[str, int] = {}
151
+ used_ids: set[str] = set()
152
+ edges: list[dict[str, Any]] = []
153
+ for record, source_file in records:
154
+ line_numbers[source_file] = line_numbers.get(source_file, 0) + 1
155
+ edge = codex_edge_from_record(record, source_file, line_numbers[source_file])
156
+ # UNIQUENESS GUARD. `messages` is a uuid-keyed view, so two edges
157
+ # sharing a uuid would silently double a record. No duplicate payload
158
+ # id appears across the 74 local rollouts measured 2026-09-11, but the
159
+ # ids come from a source this workspace does not control, so a repeat
160
+ # falls back to the synthetic form rather than shadowing the first.
161
+ if str(edge["uuid"]) in used_ids:
162
+ kind = record.get("type")
163
+ prefix = kind if isinstance(kind, str) and kind else "record"
164
+ edge["uuid"] = f"{prefix}:L{line_numbers[source_file]}"
165
+ used_ids.add(str(edge["uuid"]))
166
+ edges.append(edge)
167
+ return edges
168
+
169
+
170
+ def codex_edges_jsonl_lines(records: list[tuple[dict[str, Any], str]]) -> list[str]:
171
+ """``build_codex_edges`` + serialization: the ready-to-write edges.jsonl body."""
172
+ return [edge_to_jsonl_line(edge) for edge in build_codex_edges(records)]
173
+
174
+
175
+ def codex_edge_ids(records: list[tuple[dict[str, Any], str]]) -> list[str]:
176
+ """Every record's edge id, in the same order :func:`build_codex_edges` uses.
177
+
178
+ The enrichment pass attributes steps to these ids, so it must derive them
179
+ from the one function the edges are built with rather than recomputing the
180
+ fallback rules and drifting from them.
181
+ """
182
+ return [str(edge["uuid"]) for edge in build_codex_edges(records)]
183
+
184
+
185
+ __all__ = [
186
+ "EDGE_FIELDS",
187
+ "build_codex_edges",
188
+ "codex_edge_from_record",
189
+ "codex_edge_ids",
190
+ "codex_edges_jsonl_lines",
191
+ "record_id",
192
+ ]
@@ -0,0 +1,419 @@
1
+ # SPDX-License-Identifier: Apache-2.0
2
+
3
+ """Post-conversion enrichment for Codex: restore what harbor 0.22.0 drops.
4
+
5
+ Pure function over a trajectory dict plus the rollout's raw records — no
6
+ harbor import, no I/O. The Claude Code counterpart is
7
+ :mod:`atif_converter.domain.enrichment`; this module repairs the Codex gaps
8
+ that are repairable wrap-locally and, exactly like that one, REFUSES rather
9
+ than guesses when its re-derivation disagrees with harbor.
10
+
11
+ What it repairs, by the gap it answers
12
+ --------------------------------------
13
+ ``ITEM_IDS_NOT_PRESERVED``
14
+ ``step.extra["source_uuids"]`` lists the rollout record ids that
15
+ contributed to the step, using the SAME key the Claude Code pass writes so
16
+ atif-duck's ``steps.source_uuids`` column means one thing for every agent.
17
+
18
+ Cache split
19
+ ``trajectory.extra["cache_creation_total"]`` is surfaced from
20
+ ``final_metrics.extra.total_cache_write_input_tokens``, again matching the
21
+ Claude Code key so one query answers cache-creation for both agents.
22
+
23
+ ``COMPACTION_UNHANDLED``
24
+ ``trajectory.extra["codex_compaction_count"]`` records how many
25
+ ``compacted`` records the rollout carried. The summaries themselves are
26
+ not reconstructable into steps — harbor emitted none — so the count is
27
+ the honest repair: it tells a reader the step list has a seam without
28
+ inventing a step to mark it.
29
+
30
+ ATTRIBUTION STRATEGY: identity first, and position only where the format
31
+ leaves no id
32
+ -----------------------------------------------------------------------------
33
+ Three passes, ordered by how much they can be trusted.
34
+
35
+ 1. **Tool records, by call id.** harbor keeps each Codex ``call_id`` verbatim
36
+ in ``tool_calls[].tool_call_id``, so a step's tool records are an exact
37
+ lookup: no ordering assumption, no off-by-one to make, and a step whose
38
+ call ids match nothing simply gets no tool attribution.
39
+
40
+ 2. **Agent message records, by API call id.** harbor bundles the assistant
41
+ messages of one model request into one step and stamps that request's id
42
+ into ``step.extra.api_call_id``. This pass re-derives the same ids from the
43
+ raw records — the counter advances on a ``token_count`` event that closed a
44
+ request which produced output, which is harbor's own rule — and attributes
45
+ each assistant ``message`` record to the step bearing its request id. That
46
+ is what makes an assistant record with EMPTY text attributable at all: it
47
+ contributes no text to its step (harbor filters empty parts), so no
48
+ text-matching walk could ever place it, while its request id places it
49
+ exactly.
50
+
51
+ 3. **User and system message steps, by position.** These carry no id anywhere:
52
+ harbor emits one step per non-assistant message record, in file order, with
53
+ no bundling and nothing filtered. The walk therefore pairs them 1:1 and
54
+ cross-checks each pairing against harbor's own text extraction. The FIRST
55
+ pair whose text or source disagrees stops this pass and records
56
+ ``trajectory.extra["enrichment_truncated_at_step"]``: once positions
57
+ disagree, every later pairing is suspect.
58
+
59
+ Reasoning records are attributed only to a step that actually carries
60
+ ``reasoning_content``. A reasoning record belongs to the API call that
61
+ consumed it, but recent Codex builds ship the content encrypted and harbor
62
+ reads only the plaintext summary
63
+ (:attr:`~atif_converter.domain.codex_fidelity.CodexFidelityGap.REASONING_ENCRYPTED_DROPPED`),
64
+ so claiming a dropped record "contributed to" a step would overstate what the
65
+ step holds.
66
+
67
+ After all three passes, any step left with no ``source_uuids`` and any
68
+ non-assistant message record never consumed are counted into
69
+ ``trajectory.extra["enrichment_unattributed_steps"]`` and
70
+ ``trajectory.extra["enrichment_leftover_messages"]`` (present only when
71
+ nonzero). ``extra`` is free-form per the contract, so all of this is
72
+ contract-compatible; a truncation marker or a nonzero count means the
73
+ re-derivation disagreed with harbor and the affected attributions were
74
+ REFUSED, not guessed.
75
+ """
76
+
77
+ from __future__ import annotations
78
+
79
+ from dataclasses import dataclass, field
80
+ from typing import Any
81
+
82
+ from loguru import logger
83
+
84
+ from atif_converter.domain.codex_edges import build_codex_edges
85
+
86
+ _SOURCE_UUIDS = "source_uuids"
87
+ _TRUNCATED_AT = "enrichment_truncated_at_step"
88
+ #: Message records that reached no step: the silent end of the positional walk.
89
+ _LEFTOVER_MESSAGES = "enrichment_leftover_messages"
90
+
91
+ #: ATIF ``source`` values keyed by the raw Codex message role. ``developer``
92
+ #: and anything unknown land on ``system`` — the flattening named by
93
+ #: ``DEVELOPER_ROLE_FLATTENED_TO_SYSTEM``.
94
+ _ROLE_TO_SOURCE: dict[str, str] = {"user": "user", "assistant": "agent"}
95
+
96
+ #: Payload types that CONSUME a pending reasoning record, mirroring harbor's
97
+ #: normalization: each one clears ``pending_reasoning`` after reading it.
98
+ _REASONING_CONSUMERS: frozenset[str] = frozenset(
99
+ {
100
+ "message",
101
+ "web_search_call",
102
+ "function_call",
103
+ "custom_tool_call",
104
+ "function_call_output",
105
+ "custom_tool_call_output",
106
+ }
107
+ )
108
+
109
+
110
+ def _payload(record: dict[str, Any]) -> dict[str, Any]:
111
+ payload = record.get("payload")
112
+ return payload if isinstance(payload, dict) else {}
113
+
114
+
115
+ def message_text(payload: dict[str, Any]) -> str:
116
+ """Replicate harbor's ``Codex._extract_message_text``.
117
+
118
+ Every content block's ``text`` field, concatenated with NO separator —
119
+ which is what harbor does, and the reason a two-block developer message
120
+ reproduces as one run-together string rather than two lines.
121
+ """
122
+ content = payload.get("content")
123
+ if not isinstance(content, list):
124
+ return ""
125
+ parts: list[str] = []
126
+ for block in content:
127
+ if isinstance(block, dict):
128
+ text = block.get("text")
129
+ if isinstance(text, str):
130
+ parts.append(text)
131
+ return "".join(parts)
132
+
133
+
134
+ def _source_for_role(role: Any) -> str:
135
+ return _ROLE_TO_SOURCE.get(role, "system") if isinstance(role, str) else "system"
136
+
137
+
138
+ @dataclass(slots=True)
139
+ class _RolloutIndex:
140
+ """Everything the three passes need, derived from the raw records once."""
141
+
142
+ #: call_id -> record ids naming it (the call and its output).
143
+ by_call_id: dict[str, list[str]] = field(default_factory=dict)
144
+ #: api_call_id -> assistant message / web-search record ids in that request.
145
+ by_api_call: dict[str, list[str]] = field(default_factory=dict)
146
+ #: Non-assistant message records in file order: (record id, source, text).
147
+ non_agent_messages: list[tuple[str, str, str]] = field(default_factory=list)
148
+ #: record id or call_id -> the reasoning record it consumed.
149
+ reasoning_for: dict[str, str] = field(default_factory=dict)
150
+ #: ``compacted`` records seen.
151
+ compaction_count: int = 0
152
+
153
+
154
+ def _index_rollout(records: list[tuple[dict[str, Any], str]]) -> _RolloutIndex:
155
+ """Re-derive harbor's grouping over the raw records, ids attached.
156
+
157
+ The ``api_call_id`` counter replicates harbor's exactly: it starts at
158
+ ``api_call_1``, and a ``token_count`` event advances it only when the
159
+ request it closes produced model output. Anything that diverges here
160
+ diverges from the step ids in ``step.extra.api_call_id``, which is what
161
+ pass 2 joins on — so a mismatch shows up as an unattributed step rather
162
+ than as a wrong attribution.
163
+ """
164
+ index = _RolloutIndex()
165
+ edges = build_codex_edges(records)
166
+ api_call_number = 1
167
+ current_api_call = "api_call_1"
168
+ saw_model_output = False
169
+ pending_reasoning: str | None = None
170
+ #: call_id -> the api_call id current when the CALL was seen, because
171
+ #: harbor's normalized event for a tool call is emitted at its OUTPUT and
172
+ #: carries the call's request id, not the one live at output time.
173
+ pending_call_api: dict[str, str] = {}
174
+
175
+ for (record, _source_file), edge in zip(records, edges, strict=True):
176
+ record_id = str(edge["uuid"])
177
+ record_type = record.get("type")
178
+ payload = _payload(record)
179
+
180
+ if record_type == "compacted":
181
+ index.compaction_count += 1
182
+ continue
183
+ if record_type == "event_msg":
184
+ if payload.get("type") == "token_count" and saw_model_output:
185
+ api_call_number += 1
186
+ current_api_call = f"api_call_{api_call_number}"
187
+ saw_model_output = False
188
+ continue
189
+ if record_type != "response_item":
190
+ continue
191
+
192
+ item_type = payload.get("type")
193
+ if item_type == "reasoning":
194
+ pending_reasoning = record_id
195
+ continue
196
+
197
+ if _str(item_type) in _REASONING_CONSUMERS and pending_reasoning is not None:
198
+ key = payload.get("call_id") if "call" in str(item_type) else record_id
199
+ index.reasoning_for[str(key) if isinstance(key, str) and key else record_id] = (
200
+ pending_reasoning
201
+ )
202
+ pending_reasoning = None
203
+
204
+ if item_type == "message":
205
+ role = payload.get("role")
206
+ source = _source_for_role(role)
207
+ if role == "assistant":
208
+ index.by_api_call.setdefault(current_api_call, []).append(record_id)
209
+ saw_model_output = True
210
+ else:
211
+ index.non_agent_messages.append((record_id, source, message_text(payload)))
212
+ continue
213
+
214
+ if item_type == "web_search_call":
215
+ index.by_api_call.setdefault(current_api_call, []).append(record_id)
216
+ saw_model_output = True
217
+ continue
218
+
219
+ call_id = payload.get("call_id")
220
+ if isinstance(call_id, str) and call_id:
221
+ index.by_call_id.setdefault(call_id, []).append(record_id)
222
+ if _str(item_type) in {"function_call", "custom_tool_call"}:
223
+ pending_call_api[call_id] = current_api_call
224
+ saw_model_output = True
225
+
226
+ return index
227
+
228
+
229
+ def _step_call_ids(step: dict[str, Any]) -> set[str]:
230
+ """Every ``tool_call_id`` and ``source_call_id`` this step names."""
231
+ ids: set[str] = set()
232
+ for call in step.get("tool_calls") or []:
233
+ if isinstance(call, dict):
234
+ call_id = call.get("tool_call_id")
235
+ if isinstance(call_id, str) and call_id:
236
+ ids.add(call_id)
237
+ observation = step.get("observation")
238
+ if isinstance(observation, dict):
239
+ for result in observation.get("results") or []:
240
+ if isinstance(result, dict):
241
+ call_id = result.get("source_call_id")
242
+ if isinstance(call_id, str) and call_id:
243
+ ids.add(call_id)
244
+ return ids
245
+
246
+
247
+ def _step_extra(step: dict[str, Any]) -> dict[str, Any]:
248
+ extra = step.get("extra")
249
+ return extra if isinstance(extra, dict) else {}
250
+
251
+
252
+ def _step_id(step: dict[str, Any]) -> int | None:
253
+ step_id = step.get("step_id")
254
+ return step_id if isinstance(step_id, int) else None
255
+
256
+
257
+ def _attach(step: dict[str, Any], ids: list[str]) -> None:
258
+ """Merge ``ids`` into the step's ``extra.source_uuids``, order-preserving."""
259
+ if not ids:
260
+ return
261
+ extra = step.get("extra")
262
+ if not isinstance(extra, dict):
263
+ extra = {}
264
+ step["extra"] = extra
265
+ existing = extra.get(_SOURCE_UUIDS)
266
+ merged: list[str] = [str(value) for value in existing] if isinstance(existing, list) else []
267
+ for value in ids:
268
+ if value not in merged:
269
+ merged.append(value)
270
+ extra[_SOURCE_UUIDS] = merged
271
+
272
+
273
+ def enrich_codex_trajectory(
274
+ trajectory: dict[str, Any],
275
+ records: list[tuple[dict[str, Any], str]],
276
+ *,
277
+ copy_input: bool = True,
278
+ ) -> dict[str, Any]:
279
+ """Return ``trajectory`` with Codex source attribution and cache totals.
280
+
281
+ Parameters
282
+ ----------
283
+ trajectory
284
+ The harbor Codex trajectory as a dict (``model_dump(mode="json")``).
285
+ records
286
+ The rollout's ``(record, source_file)`` pairs in FILE order — the same
287
+ pairs :func:`~atif_converter.domain.codex_edges.build_codex_edges`
288
+ reads, so the ids attributed here are the ids in ``edges.jsonl``.
289
+ copy_input
290
+ Deep-copy before mutating. The production caller passes ``False``
291
+ because the harbor trajectory has no other consumer and a rollout's
292
+ step list is large.
293
+ """
294
+ if copy_input:
295
+ import copy
296
+
297
+ trajectory = copy.deepcopy(trajectory)
298
+
299
+ steps = trajectory.get("steps")
300
+ if not isinstance(steps, list):
301
+ return trajectory
302
+
303
+ index = _index_rollout(records)
304
+
305
+ # Passes 1 and 2 — identity. Neither can misattribute: both are lookups.
306
+ for step in steps:
307
+ if not isinstance(step, dict):
308
+ continue
309
+ attributed: list[str] = []
310
+ call_ids = sorted(_step_call_ids(step))
311
+ for call_id in call_ids:
312
+ attributed.extend(index.by_call_id.get(call_id, ()))
313
+ api_call_id = _step_extra(step).get("api_call_id")
314
+ if isinstance(api_call_id, str):
315
+ attributed.extend(index.by_api_call.get(api_call_id, ()))
316
+ if step.get("reasoning_content"):
317
+ for key in (*call_ids, *attributed):
318
+ reasoning_id = index.reasoning_for.get(key)
319
+ if reasoning_id is not None and reasoning_id not in attributed:
320
+ attributed.append(reasoning_id)
321
+ _attach(step, attributed)
322
+
323
+ # Pass 3 — position, over the non-assistant message steps only.
324
+ truncated_at, leftover_messages = _attribute_non_agent_messages(steps, index.non_agent_messages)
325
+
326
+ extra = trajectory.get("extra")
327
+ if not isinstance(extra, dict):
328
+ extra = {}
329
+ _record_cache_total(trajectory, extra)
330
+ if index.compaction_count:
331
+ extra["codex_compaction_count"] = index.compaction_count
332
+ if truncated_at is not None:
333
+ extra[_TRUNCATED_AT] = truncated_at
334
+ logger.warning(
335
+ "codex enrichment: message attribution refused at step {} — the "
336
+ "positional walk disagreed with harbor's step list; later steps "
337
+ "carry no source_uuids rather than a guessed one",
338
+ truncated_at,
339
+ )
340
+ unattributed = sum(
341
+ 1 for step in steps if isinstance(step, dict) and not _step_extra(step).get(_SOURCE_UUIDS)
342
+ )
343
+ if unattributed:
344
+ extra["enrichment_unattributed_steps"] = unattributed
345
+ if leftover_messages:
346
+ extra[_LEFTOVER_MESSAGES] = leftover_messages
347
+ logger.warning(
348
+ "codex enrichment: {} message record(s) attributed to no step — the "
349
+ "loss report counts them convertible, so this is real unreported loss",
350
+ leftover_messages,
351
+ )
352
+ if extra:
353
+ trajectory["extra"] = extra
354
+ return trajectory
355
+
356
+
357
+ def _attribute_non_agent_messages(
358
+ steps: list[Any],
359
+ messages: list[tuple[str, str, str]],
360
+ ) -> tuple[int | None, int]:
361
+ """Pair user/system steps with user/developer message records, 1:1.
362
+
363
+ Returns the step id the walk refused at (``None`` when every such step
364
+ paired cleanly) and the count of message records left unconsumed. A step is
365
+ refused when the next unconsumed record's ATIF source or harbor-extracted
366
+ text disagrees with it.
367
+
368
+ THE TWO ENDS OF THE WALK FAIL DIFFERENTLY, and both have to be reported.
369
+ Running out of RECORDS refuses the step, which is loud. Running out of
370
+ STEPS is silent: the remaining records attribute to nothing, and the loss
371
+ report still counts them convertible, so it would understate the loss with
372
+ no marker anywhere. The leftover count is that marker.
373
+ """
374
+ cursor = 0
375
+ for step in steps:
376
+ if not isinstance(step, dict):
377
+ continue
378
+ source = step.get("source")
379
+ if _str(source) not in {"user", "system"}:
380
+ continue
381
+ if cursor >= len(messages):
382
+ return _step_id(step), 0
383
+ record_id, record_source, record_text = messages[cursor]
384
+ if record_source != source or record_text != (step.get("message") or ""):
385
+ return _step_id(step), len(messages) - cursor
386
+ _attach(step, [record_id])
387
+ cursor += 1
388
+ return None, len(messages) - cursor
389
+
390
+
391
+ def _record_cache_total(trajectory: dict[str, Any], extra: dict[str, Any]) -> None:
392
+ """Surface the cache-creation total under the Claude Code key.
393
+
394
+ harbor puts the Codex figure in
395
+ ``final_metrics.extra.total_cache_write_input_tokens`` and the Claude Code
396
+ figure in ``final_metrics.extra.cache_creation_input_tokens``. One name
397
+ reaches SQL, so the Codex value is republished under
398
+ ``trajectory.extra["cache_creation_total"]`` — the key the enrichment pass
399
+ for Claude Code already writes and the corpus already carries.
400
+ """
401
+ final_metrics = trajectory.get("final_metrics")
402
+ if not isinstance(final_metrics, dict):
403
+ return
404
+ metrics_extra = final_metrics.get("extra")
405
+ if not isinstance(metrics_extra, dict):
406
+ return
407
+ total = metrics_extra.get("total_cache_write_input_tokens")
408
+ if isinstance(total, int):
409
+ extra["cache_creation_total"] = total
410
+
411
+
412
+ def _str(value: object) -> str | None:
413
+ """``value`` when it is a string, else ``None``: a set-membership key that cannot raise.
414
+
415
+ A transcript field can hold any JSON value, and an object or a list is unhashable, so
416
+ testing one against a set raised ``TypeError`` past the converter (found by the fuzzer).
417
+ A non-string value never matched one of these string sets, so no result changes.
418
+ """
419
+ return value if isinstance(value, str) else None