What changed, and why

When a category overspends, the W6 monitor emails an alert, then hands the thread to the W7 budget balancer, which offers coverage moves and reads the owner's plain-English reply. The balancer was posting that offer — and every follow-up — on a brand-new email thread, so a single overspend showed up as two disconnected conversations in the owner's inbox. The reported bug: balancing follow-ups should reply on the alert thread, not start their own.

The separate thread was not an oversight. It was a workaround (commit a60e676) for a real delivery bug: the balancer replied on the alert thread by replying to its latest message, which is the agent's own alert. AgentMail addressed that reply back to the agent's inbox, so the owner got the alert but never the options. The workaround swapped threading for delivery.

This change gets both. AgentMail's reply() takes a to recipient override, so the balancer can reply on the alert thread and still address the owners. The fix is a small recipient parameter in the mail layer, plus reverting the balancer to reply-on-thread now that delivery is handled. The sections below follow the change from the mail seam outward to the workflow and the tests.

The recipient override, threaded through the mail client

send_on_thread is the reply workhorse. It already deduped on a sequence label and replied to the thread's latest message; it now takes an optional to. When set, those recipients override AgentMail's default reply addressing — which is exactly what fails when the latest message is the agent's own alert. The other callers (the per-transaction W2, the autonomy offer) pass nothing and keep the default, because they only ever reply after the owner has spoken.

send_on_thread gains to, forwarded to the backend reply.

src/ynab_agent/mail/client.py · 274 lines
src/ynab_agent/mail/client.py274 lines · Python
⋯ 183 lines hidden (lines 1–183)
1"""AgentMail: the email plumbing for the per-transaction threads (SPEC §5).
2 
3A thin, deterministic client — sending and replying are direct API calls, not
4agentic. One inbox is the agent's address; each transaction is one thread within
5it (started by the first send, continued by replies). The model composes the
6prose elsewhere; this module only puts it on the wire.
7 
8The backend is a small :class:`MailBackend` protocol over *our* two operations
9(send a new thread / reply on one), so tests inject a fake and strict mypy never
10has to reason about the AgentMail SDK's surface. :func:`MailClient.from_env`
11wires the real SDK behind that protocol, reading ``AGENTMAIL_API_KEY``.
12"""
13 
14from __future__ import annotations
15 
16import os
17from typing import TYPE_CHECKING, Protocol
18 
19from pydantic import model_validator
20 
21from ynab_agent.domain.base import Frozen
22 
23if TYPE_CHECKING:
24 from agentmail import AgentMail
25 
26_API_KEY_ENV = "AGENTMAIL_API_KEY"
27# Every agent message carries this label, plus per-transaction / per-action
28# labels that serve as idempotency keys (no separate store — derived from the
29# AgentMail thread, SPEC §5).
30_AGENT_LABEL = "ynab-agent"
31 
32 
33class SentEmail(Frozen):
34 """The identifiers AgentMail returns for a sent message."""
35 
36 message_id: str
37 thread_id: str
38 
39 
40class OutboundEmail(Frozen):
41 """One message to send: a new thread, or a reply continuing one."""
42 
43 inbox_id: str
44 subject: str
45 text: str
46 to: tuple[str, ...] = ()
47 reply_to_message_id: str | None = None
48 
49 @model_validator(mode="after")
50 def _check_addressing(self) -> OutboundEmail:
51 if self.reply_to_message_id is None and not self.to:
52 msg = "a new-thread email needs at least one recipient"
53 raise ValueError(msg)
54 return self
55 
56 
57class MailBackend(Protocol):
58 """The mail operations the client needs (implemented over AgentMail)."""
59 
60 def send_new(
61 self,
62 inbox_id: str,
63 to: list[str],
64 subject: str,
65 text: str,
66 labels: list[str] | None = None,
67 ) -> SentEmail:
68 """Start a new thread."""
69 ...
70 
71 def send_reply(
72 self,
73 inbox_id: str,
74 message_id: str,
75 text: str,
76 labels: list[str] | None = None,
77 to: list[str] | None = None,
78 ) -> SentEmail:
79 """Reply on an existing thread.
80 
81 ``to`` overrides the recipients. AgentMail otherwise addresses a reply
82 to the *sender* of ``message_id``; when that message is the agent's own
83 (replying on a thread the agent last spoke on), the reply loops back to
84 the agent and the owner never sees it. Pass the owners to deliver it.
85 """
86 ...
87 
88 def find_thread(self, inbox_id: str, label: str) -> str | None:
89 """The thread id of the first message carrying ``label``, or None.
90 
91 Labels are the agent's idempotency keys: a per-transaction label lets
92 ``open_thread`` find an already-opened thread instead of duplicating it.
93 """
94 ...
95 
96 def latest_message_id(self, inbox_id: str, thread_id: str) -> str | None:
97 """The newest message id on a thread (the target to reply to)."""
98 ...
99 
100 def archive(self, inbox_id: str, thread_id: str) -> None:
101 """Mark a thread closed (an ``ynab-archived`` label)."""
102 ...
103 
104 
105# The process-wide cached client built by ``from_env`` (see its docstring).
106# ``tests/conftest.py`` resets this between tests.
107_CACHED: MailClient | None = None
108 
109 
110class MailClient:
111 """Sends a transaction's outbound mail through AgentMail (SPEC §5)."""
112 
113 def __init__(self, backend: MailBackend) -> None:
114 """Wrap a backend (the real AgentMail adapter, or a test fake)."""
115 self._backend = backend
116 
117 @classmethod
118 def from_env(cls) -> MailClient:
119 """Build (once) a client backed by the real AgentMail SDK.
120 
121 Cached like the YNAB client: the instrumented httpx client is built once
122 and reused, so a per-call client (never closed) can't leak sockets.
123 Tests reset the cache (see ``tests/conftest.py``).
124 
125 Raises:
126 RuntimeError: If ``AGENTMAIL_API_KEY`` is not set.
127 """
128 global _CACHED
129 if _CACHED is not None:
130 return _CACHED
131 key = os.environ.get(_API_KEY_ENV)
132 if not key:
133 msg = f"{_API_KEY_ENV} is not set"
134 raise RuntimeError(msg)
135 import httpx
136 from agentmail import AgentMail
137 
138 from ynab_agent.telemetry import instrument_httpx
139 
140 # Inject our own httpx client so trace context propagates on mail sends
141 # (the SDK accepts one); instrument it the same way the YNAB client is.
142 http_client = httpx.Client(timeout=60.0)
143 instrument_httpx(http_client)
144 _CACHED = cls(
145 _AgentMailBackend(AgentMail(api_key=key, httpx_client=http_client))
146 )
147 return _CACHED
148 
149 def send(self, email: OutboundEmail) -> SentEmail:
150 """Send a new thread, or a reply if a message to reply to is set."""
151 if email.reply_to_message_id is not None:
152 return self._backend.send_reply(
153 email.inbox_id,
154 email.reply_to_message_id,
155 email.text,
156 to=list(email.to) or None,
157 )
158 return self._backend.send_new(
159 email.inbox_id, list(email.to), email.subject, email.text
160 )
161 
162 def open_thread(
163 self,
164 *,
165 inbox_id: str,
166 to: list[str],
167 subject: str,
168 body: str,
169 txn_label: str,
170 ) -> str:
171 """Open the transaction's thread (idempotent); return its id.
172 
173 If a thread already carries ``txn_label`` (a retry of this open), return
174 it rather than starting a duplicate; otherwise send the opening message.
175 """
176 existing = self._backend.find_thread(inbox_id, txn_label)
177 if existing is not None:
178 return existing
179 sent = self._backend.send_new(
180 inbox_id, to, subject, body, labels=[_AGENT_LABEL, txn_label]
181 )
182 return sent.thread_id
183 
184 def send_on_thread(
185 self,
186 *,
187 inbox_id: str,
188 thread_id: str,
189 body: str,
190 seq_label: str,
191 to: list[str] | None = None,
192 ) -> bool:
193 """Reply on a thread (idempotent on ``seq_label``); True if sent.
194 
195 Skips if a message with ``seq_label`` is already on record, so a retry
196 never double-sends (SPEC §3 outbound dedup). ``to`` overrides the
197 recipients: when the thread's latest message is the agent's own (e.g.
198 the W6 alert the balancer replies on), AgentMail would address the reply
199 back to the agent, so pass the owners to deliver it to them (SPEC §8).
200 """
201 if self._backend.find_thread(inbox_id, seq_label) is not None:
202 return False
203 target = self._backend.latest_message_id(inbox_id, thread_id)
204 if target is None:
205 return False
206 self._backend.send_reply(
207 inbox_id, target, body, labels=[_AGENT_LABEL, seq_label], to=to
208 )
209 return True
⋯ 65 lines hidden (lines 210–274)
210 
211 def close(self, *, inbox_id: str, thread_id: str) -> None:
212 """Mark the transaction's thread closed."""
213 self._backend.archive(inbox_id, thread_id)
214 
215 
216class _AgentMailBackend:
217 """Adapts the AgentMail SDK to the :class:`MailBackend` protocol."""
218 
219 def __init__(self, client: AgentMail) -> None:
220 self._client = client
221 
222 def send_new(
223 self,
224 inbox_id: str,
225 to: list[str],
226 subject: str,
227 text: str,
228 labels: list[str] | None = None,
229 ) -> SentEmail:
230 result = self._client.inboxes.messages.send(
231 inbox_id, to=to, subject=subject, text=text, labels=labels
232 )
233 return SentEmail(
234 message_id=result.message_id, thread_id=result.thread_id
235 )
236 
237 def send_reply(
238 self,
239 inbox_id: str,
240 message_id: str,
241 text: str,
242 labels: list[str] | None = None,
243 to: list[str] | None = None,
244 ) -> SentEmail:
245 # ``to`` overrides the recipients; omit it entirely (not None) when
246 # unset so AgentMail keeps its own reply-addressing default.
247 if to is not None:
248 result = self._client.inboxes.messages.reply(
249 inbox_id, message_id, text=text, labels=labels, to=to
250 )
251 else:
252 result = self._client.inboxes.messages.reply(
253 inbox_id, message_id, text=text, labels=labels
254 )
255 return SentEmail(
256 message_id=result.message_id, thread_id=result.thread_id
257 )
258 
259 def find_thread(self, inbox_id: str, label: str) -> str | None:
260 result = self._client.inboxes.messages.list(
261 inbox_id, labels=[label], limit=1
262 )
263 messages = result.messages
264 return messages[0].thread_id if messages else None
265 
266 def latest_message_id(self, inbox_id: str, thread_id: str) -> str | None:
267 thread = self._client.inboxes.threads.get(inbox_id, thread_id)
268 message_id: str | None = thread.last_message_id
269 return message_id
270 
271 def archive(self, inbox_id: str, thread_id: str) -> None:
272 self._client.inboxes.threads.update(
273 inbox_id, thread_id, add_labels=["ynab-archived"]
274 )

The AgentMail adapter passes to to the SDK's reply() only when it is set. Omitting the argument (rather than sending None) lets AgentMail keep its own reply-addressing behavior for every caller that does not override it — the empty default stays a true no-op.

_AgentMailBackend.send_reply: include to only when provided.

src/ynab_agent/mail/client.py · 274 lines
src/ynab_agent/mail/client.py274 lines · Python
⋯ 236 lines hidden (lines 1–236)
1"""AgentMail: the email plumbing for the per-transaction threads (SPEC §5).
2 
3A thin, deterministic client — sending and replying are direct API calls, not
4agentic. One inbox is the agent's address; each transaction is one thread within
5it (started by the first send, continued by replies). The model composes the
6prose elsewhere; this module only puts it on the wire.
7 
8The backend is a small :class:`MailBackend` protocol over *our* two operations
9(send a new thread / reply on one), so tests inject a fake and strict mypy never
10has to reason about the AgentMail SDK's surface. :func:`MailClient.from_env`
11wires the real SDK behind that protocol, reading ``AGENTMAIL_API_KEY``.
12"""
13 
14from __future__ import annotations
15 
16import os
17from typing import TYPE_CHECKING, Protocol
18 
19from pydantic import model_validator
20 
21from ynab_agent.domain.base import Frozen
22 
23if TYPE_CHECKING:
24 from agentmail import AgentMail
25 
26_API_KEY_ENV = "AGENTMAIL_API_KEY"
27# Every agent message carries this label, plus per-transaction / per-action
28# labels that serve as idempotency keys (no separate store — derived from the
29# AgentMail thread, SPEC §5).
30_AGENT_LABEL = "ynab-agent"
31 
32 
33class SentEmail(Frozen):
34 """The identifiers AgentMail returns for a sent message."""
35 
36 message_id: str
37 thread_id: str
38 
39 
40class OutboundEmail(Frozen):
41 """One message to send: a new thread, or a reply continuing one."""
42 
43 inbox_id: str
44 subject: str
45 text: str
46 to: tuple[str, ...] = ()
47 reply_to_message_id: str | None = None
48 
49 @model_validator(mode="after")
50 def _check_addressing(self) -> OutboundEmail:
51 if self.reply_to_message_id is None and not self.to:
52 msg = "a new-thread email needs at least one recipient"
53 raise ValueError(msg)
54 return self
55 
56 
57class MailBackend(Protocol):
58 """The mail operations the client needs (implemented over AgentMail)."""
59 
60 def send_new(
61 self,
62 inbox_id: str,
63 to: list[str],
64 subject: str,
65 text: str,
66 labels: list[str] | None = None,
67 ) -> SentEmail:
68 """Start a new thread."""
69 ...
70 
71 def send_reply(
72 self,
73 inbox_id: str,
74 message_id: str,
75 text: str,
76 labels: list[str] | None = None,
77 to: list[str] | None = None,
78 ) -> SentEmail:
79 """Reply on an existing thread.
80 
81 ``to`` overrides the recipients. AgentMail otherwise addresses a reply
82 to the *sender* of ``message_id``; when that message is the agent's own
83 (replying on a thread the agent last spoke on), the reply loops back to
84 the agent and the owner never sees it. Pass the owners to deliver it.
85 """
86 ...
87 
88 def find_thread(self, inbox_id: str, label: str) -> str | None:
89 """The thread id of the first message carrying ``label``, or None.
90 
91 Labels are the agent's idempotency keys: a per-transaction label lets
92 ``open_thread`` find an already-opened thread instead of duplicating it.
93 """
94 ...
95 
96 def latest_message_id(self, inbox_id: str, thread_id: str) -> str | None:
97 """The newest message id on a thread (the target to reply to)."""
98 ...
99 
100 def archive(self, inbox_id: str, thread_id: str) -> None:
101 """Mark a thread closed (an ``ynab-archived`` label)."""
102 ...
103 
104 
105# The process-wide cached client built by ``from_env`` (see its docstring).
106# ``tests/conftest.py`` resets this between tests.
107_CACHED: MailClient | None = None
108 
109 
110class MailClient:
111 """Sends a transaction's outbound mail through AgentMail (SPEC §5)."""
112 
113 def __init__(self, backend: MailBackend) -> None:
114 """Wrap a backend (the real AgentMail adapter, or a test fake)."""
115 self._backend = backend
116 
117 @classmethod
118 def from_env(cls) -> MailClient:
119 """Build (once) a client backed by the real AgentMail SDK.
120 
121 Cached like the YNAB client: the instrumented httpx client is built once
122 and reused, so a per-call client (never closed) can't leak sockets.
123 Tests reset the cache (see ``tests/conftest.py``).
124 
125 Raises:
126 RuntimeError: If ``AGENTMAIL_API_KEY`` is not set.
127 """
128 global _CACHED
129 if _CACHED is not None:
130 return _CACHED
131 key = os.environ.get(_API_KEY_ENV)
132 if not key:
133 msg = f"{_API_KEY_ENV} is not set"
134 raise RuntimeError(msg)
135 import httpx
136 from agentmail import AgentMail
137 
138 from ynab_agent.telemetry import instrument_httpx
139 
140 # Inject our own httpx client so trace context propagates on mail sends
141 # (the SDK accepts one); instrument it the same way the YNAB client is.
142 http_client = httpx.Client(timeout=60.0)
143 instrument_httpx(http_client)
144 _CACHED = cls(
145 _AgentMailBackend(AgentMail(api_key=key, httpx_client=http_client))
146 )
147 return _CACHED
148 
149 def send(self, email: OutboundEmail) -> SentEmail:
150 """Send a new thread, or a reply if a message to reply to is set."""
151 if email.reply_to_message_id is not None:
152 return self._backend.send_reply(
153 email.inbox_id,
154 email.reply_to_message_id,
155 email.text,
156 to=list(email.to) or None,
157 )
158 return self._backend.send_new(
159 email.inbox_id, list(email.to), email.subject, email.text
160 )
161 
162 def open_thread(
163 self,
164 *,
165 inbox_id: str,
166 to: list[str],
167 subject: str,
168 body: str,
169 txn_label: str,
170 ) -> str:
171 """Open the transaction's thread (idempotent); return its id.
172 
173 If a thread already carries ``txn_label`` (a retry of this open), return
174 it rather than starting a duplicate; otherwise send the opening message.
175 """
176 existing = self._backend.find_thread(inbox_id, txn_label)
177 if existing is not None:
178 return existing
179 sent = self._backend.send_new(
180 inbox_id, to, subject, body, labels=[_AGENT_LABEL, txn_label]
181 )
182 return sent.thread_id
183 
184 def send_on_thread(
185 self,
186 *,
187 inbox_id: str,
188 thread_id: str,
189 body: str,
190 seq_label: str,
191 to: list[str] | None = None,
192 ) -> bool:
193 """Reply on a thread (idempotent on ``seq_label``); True if sent.
194 
195 Skips if a message with ``seq_label`` is already on record, so a retry
196 never double-sends (SPEC §3 outbound dedup). ``to`` overrides the
197 recipients: when the thread's latest message is the agent's own (e.g.
198 the W6 alert the balancer replies on), AgentMail would address the reply
199 back to the agent, so pass the owners to deliver it to them (SPEC §8).
200 """
201 if self._backend.find_thread(inbox_id, seq_label) is not None:
202 return False
203 target = self._backend.latest_message_id(inbox_id, thread_id)
204 if target is None:
205 return False
206 self._backend.send_reply(
207 inbox_id, target, body, labels=[_AGENT_LABEL, seq_label], to=to
208 )
209 return True
210 
211 def close(self, *, inbox_id: str, thread_id: str) -> None:
212 """Mark the transaction's thread closed."""
213 self._backend.archive(inbox_id, thread_id)
214 
215 
216class _AgentMailBackend:
217 """Adapts the AgentMail SDK to the :class:`MailBackend` protocol."""
218 
219 def __init__(self, client: AgentMail) -> None:
220 self._client = client
221 
222 def send_new(
223 self,
224 inbox_id: str,
225 to: list[str],
226 subject: str,
227 text: str,
228 labels: list[str] | None = None,
229 ) -> SentEmail:
230 result = self._client.inboxes.messages.send(
231 inbox_id, to=to, subject=subject, text=text, labels=labels
232 )
233 return SentEmail(
234 message_id=result.message_id, thread_id=result.thread_id
235 )
236 
237 def send_reply(
238 self,
239 inbox_id: str,
240 message_id: str,
241 text: str,
242 labels: list[str] | None = None,
243 to: list[str] | None = None,
244 ) -> SentEmail:
245 # ``to`` overrides the recipients; omit it entirely (not None) when
246 # unset so AgentMail keeps its own reply-addressing default.
247 if to is not None:
248 result = self._client.inboxes.messages.reply(
249 inbox_id, message_id, text=text, labels=labels, to=to
250 )
251 else:
252 result = self._client.inboxes.messages.reply(
253 inbox_id, message_id, text=text, labels=labels
254 )
255 return SentEmail(
256 message_id=result.message_id, thread_id=result.thread_id
257 )
⋯ 17 lines hidden (lines 258–274)
258 
259 def find_thread(self, inbox_id: str, label: str) -> str | None:
260 result = self._client.inboxes.messages.list(
261 inbox_id, labels=[label], limit=1
262 )
263 messages = result.messages
264 return messages[0].thread_id if messages else None
265 
266 def latest_message_id(self, inbox_id: str, thread_id: str) -> str | None:
267 thread = self._client.inboxes.threads.get(inbox_id, thread_id)
268 message_id: str | None = thread.last_message_id
269 return message_id
270 
271 def archive(self, inbox_id: str, thread_id: str) -> None:
272 self._client.inboxes.threads.update(
273 inbox_id, thread_id, add_labels=["ynab-archived"]
274 )

The balancer replies on the alert thread

With delivery handled in the mail layer, the balancer's send activity goes back to a reply. send_balance_email now replies on the alert thread and passes to=list(settings.owners), because that thread's latest message is often the agent's own — the opening alert, or a back-to-back agent message — so the recipients must be set explicitly. The open_balance_thread activity that opened a separate thread is gone, along with its entry in the activity registry.

send_balance_email: reply on the thread, addressed to the owners.

src/ynab_agent/workflow/balance_activities.py · 295 lines
src/ynab_agent/workflow/balance_activities.py295 lines · Python
⋯ 270 lines hidden (lines 1–270)
1"""The I/O ports of the W7 budget balancer, as Temporal activities (SPEC §8).
2 
3Its own module so the balance workflow's sandbox import graph stays minimal (see
4``offer_activities``). Heavy clients (Temporal, YNAB, AgentMail, the model) are
5imported lazily inside the bodies so they never enter a workflow sandbox.
6 
7Division of labor: these activities read YNAB, run the model, and write one
8category's ``budgeted`` (idempotent absolute set + read-back verify); the
9*workflow* computes absolute targets from a baseline snapshot and orchestrates
10the writes, so a write retry re-sets the same value and never double-applies.
11"""
12 
13from __future__ import annotations
14 
15from typing import TYPE_CHECKING
16 
17from temporalio import activity
18 
19from ynab_agent.budget.balance import (
20 BalanceOption,
21 BalanceOutcome,
22 BudgetMove,
24from ynab_agent.budget.overspend import OverspendAssessment
25from ynab_agent.domain.money import Money
26from ynab_agent.workflow.balance_types import BalanceParams, BudgetState
27 
28if TYPE_CHECKING:
29 # ``SourceFunds`` lives in ``agentic.balance``, which pulls in pydantic-ai —
30 # never import it at module scope here (this module is passed through the
31 # workflow sandbox). The runtime import is lazy, inside the bodies.
32 from ynab_agent.agentic.balance import SourceFunds
33 from ynab_agent.budget.balance import Source
34 
35 
36def _dollars(amount: Money) -> float:
37 """A Money amount as a float dollar figure for the model's context."""
38 return float(amount.currency_amount)
39 
40 
41def _overspend_note(assessment: OverspendAssessment) -> str:
42 """A one-line human summary of how the category is tracking."""
43 return (
44 f"{assessment.name}: {assessment.spent} spent of "
45 f"{assessment.budgeted} budgeted, projected ~{assessment.projected} "
46 "by month-end."
47 )
48 
49 
50def _to_source_funds(
51 sources: tuple[Source, ...], names: dict[str, str]
52) -> tuple[SourceFunds, ...]:
53 """Domain sources as the model's ``SourceFunds`` (dollars + names)."""
54 from ynab_agent.agentic.balance import SourceFunds
55 from ynab_agent.budget.balance import READY_TO_ASSIGN_SOURCE
56 
57 funds: list[SourceFunds] = []
58 for source in sources:
59 is_rta = source.category == READY_TO_ASSIGN_SOURCE
60 funds.append(
61 SourceFunds(
62 id=str(source.category),
63 name="Ready to Assign"
64 if is_rta
65 else names.get(str(source.category), str(source.category)),
66 available=_dollars(source.available),
67 kind="ready-to-assign" if is_rta else "category",
68 )
69 )
70 return tuple(funds)
71 
72 
73@activity.defn
74async def start_balance_offer(
75 assessment: OverspendAssessment, thread_id: str
76) -> None:
77 """Start the balance offer for an alerted category (SPEC §8, the W6→W7 tie).
78 
79 Started ``REJECT_DUPLICATE`` on the (category, period) id, so a worsening
80 re-alert in the same month is a no-op, not a second offer; at most one
81 coverage offer per category per period, matching the monitor's own dedupe.
82 """
83 import datetime
84 
85 from temporalio.common import WorkflowIDReusePolicy
86 from temporalio.exceptions import WorkflowAlreadyStartedError
87 
88 from ynab_agent.budget.balance import need_from_assessment
89 from ynab_agent.workflow.balance_types import (
90 BalanceParams,
91 balance_workflow_id,
92 )
93 from ynab_agent.workflow.temporal_client import client, task_queue
94 
95 if need_from_assessment(assessment).shortfall.is_zero:
96 return # nothing to cover
97 period = datetime.datetime.now(datetime.UTC).strftime("%Y-%m")
98 temporal = await client()
99 try:
100 await temporal.start_workflow(
101 "BudgetBalanceWorkflow",
102 BalanceParams(
103 assessment=assessment, thread_id=thread_id, period=period
104 ),
105 id=balance_workflow_id(str(assessment.category), period),
106 task_queue=task_queue(),
107 id_reuse_policy=WorkflowIDReusePolicy.REJECT_DUPLICATE,
108 )
109 except WorkflowAlreadyStartedError:
110 return
111 
112 
113@activity.defn
114async def propose_balance_options(
115 params: BalanceParams,
116) -> list[BalanceOption]:
117 """Read the budget, ask the model for options, keep the feasible ones (§8).
118 
119 The model proposes; the deterministic guard (:func:`feasible_options`) drops
120 anything that doesn't add up or pulls money a source doesn't have. When the
121 model yields nothing usable, fall back to the greedy plan; an empty list
122 means even that can't cover it from current funds.
123 """
124 import asyncio
125 
126 from ynab_agent.agentic.balance import (
127 BalanceContext,
128 propose_balance,
129 to_options,
130 )
131 from ynab_agent.budget.balance import (
132 fallback_option,
133 feasible_options,
134 need_from_assessment,
135 sources_from_spends,
136 )
137 from ynab_agent.ynab.client import YnabClient
138 
139 assessment = params.assessment
140 need = need_from_assessment(assessment)
141 if need.shortfall.is_zero:
142 return []
143 client = YnabClient.from_env()
144 spends = await asyncio.to_thread(client.category_spends)
145 rta = await asyncio.to_thread(client.ready_to_assign)
146 sources = sources_from_spends(spends, rta, exclude=assessment.category)
147 if not sources:
148 return []
149 names = {str(spend.category): spend.name for spend in spends}
150 context = BalanceContext(
151 needy_category_id=str(assessment.category),
152 needy_category_name=assessment.name,
153 shortfall=_dollars(need.shortfall),
154 overspend_note=_overspend_note(assessment),
155 sources=_to_source_funds(sources, names),
156 )
157 proposal = await propose_balance(context)
158 options = to_options(proposal, destination=assessment.category)
159 feasible = feasible_options(options, need, sources)
160 if feasible:
161 return list(feasible)
162 fallback = fallback_option(need, sources)
163 return [fallback] if fallback is not None else []
164 
165 
166@activity.defn
167async def interpret_balance_reply(
168 params: BalanceParams, reply_text: str, options: list[BalanceOption]
169) -> BalanceOutcome:
170 """Read the owner's free-text reply into a concrete outcome (SPEC §8).
171 
172 Re-reads the sources (funds may have shifted since the offer) so the model
173 resolves "from dining instead" against current reality; the ``to_*`` seam
174 maps the reading onto a domain outcome the workflow branches on.
175 """
176 import asyncio
177 
178 from ynab_agent.agentic.balance import (
179 BalanceReplyRequest,
180 MoveSpec,
181 OfferedOption,
182 to_balance_outcome,
183 )
184 from ynab_agent.agentic.balance import (
185 interpret_balance_reply as run_reply,
186 )
187 from ynab_agent.budget.balance import (
188 need_from_assessment,
189 sources_from_spends,
190 )
191 from ynab_agent.ynab.client import YnabClient
192 
193 assessment = params.assessment
194 client = YnabClient.from_env()
195 spends = await asyncio.to_thread(client.category_spends)
196 rta = await asyncio.to_thread(client.ready_to_assign)
197 sources = sources_from_spends(spends, rta, exclude=assessment.category)
198 names = {str(spend.category): spend.name for spend in spends}
199 offered = tuple(
200 OfferedOption(
201 label=option.label,
202 moves=tuple(
203 MoveSpec(
204 source_category_id=str(move.source),
205 amount=_dollars(move.amount),
206 )
207 for move in option.moves
208 ),
209 rationale=option.rationale,
210 )
211 for option in options
212 )
213 request = BalanceReplyRequest(
214 reply_text=reply_text,
215 needy_category_name=assessment.name,
216 shortfall=_dollars(need_from_assessment(assessment).shortfall),
217 options=offered,
218 sources=_to_source_funds(sources, names),
219 )
220 reading = await run_reply(request)
221 return to_balance_outcome(reading, destination=assessment.category)
222 
223 
224@activity.defn
225async def read_budget_state() -> BudgetState:
226 """Snapshot available funds and current budgets for the apply (SPEC §8)."""
227 import asyncio
228 
229 from ynab_agent.budget.balance import READY_TO_ASSIGN_SOURCE
230 from ynab_agent.ynab.client import YnabClient
231 
232 client = YnabClient.from_env()
233 spends = await asyncio.to_thread(client.category_spends)
234 rta = await asyncio.to_thread(client.ready_to_assign)
235 available = {spend.category: spend.balance for spend in spends}
236 available[READY_TO_ASSIGN_SOURCE] = rta
237 budgeted = {spend.category: spend.budgeted for spend in spends}
238 return BudgetState(available=available, budgeted=budgeted)
239 
240 
241@activity.defn
242async def set_category_budgeted(category_id: str, target: Money) -> bool:
243 """Set a category's month budget to ``target`` and verify it (SPEC §8).
244 
245 An absolute write (idempotent on retry), then an independent read-back —
246 never trusting its echo (SPEC §0.6). ``True`` when the read confirms.
247 """
248 import asyncio
249 
250 from ynab_agent.ynab.client import YnabClient
251 
252 client = YnabClient.from_env()
253 await asyncio.to_thread(client.set_budgeted, category_id, target)
254 actual = await asyncio.to_thread(client.read_budgeted, category_id)
255 return actual == target
256 
257 
258@activity.defn
259async def log_budget_moves(moves: list[BudgetMove], period: str) -> None:
260 """Record each applied reallocation to the audit trail (SPEC §8, §9)."""
261 from ynab_agent.audit.record import record_budget_move
262 
263 for move in moves:
264 event = record_budget_move(move, period)
265 activity.logger.info(
266 "budget move applied", extra={"audit": event.model_dump()}
267 )
268 
269 
270@activity.defn
271async def send_balance_email(thread_id: str, body: str, seq_label: str) -> None:
272 """Reply on the overspend-alert thread, addressed to the owners (SPEC §8).
273 
274 The whole balance conversation lives on the W6 alert thread (the W6→W7 tie):
275 options, clarifications, and the apply/decline confirmation all reply there,
276 so the owner sees one thread per overspend. ``to`` is set explicitly because
277 that thread's latest message is often the agent's own (the alert, or a
278 back-to-back agent reply); without it AgentMail addresses the reply back to
279 the agent and the owner never receives it. Idempotent on ``seq_label``.
280 """
281 import asyncio
282 
283 from ynab_agent.mail.client import MailClient
284 from ynab_agent.settings import Settings
285 
286 settings = Settings()
287 mail = MailClient.from_env()
288 await asyncio.to_thread(
289 mail.send_on_thread,
290 inbox_id=settings.inbox,
291 thread_id=thread_id,
292 body=body,
293 seq_label=seq_label,
294 to=list(settings.owners),
295 )

Indexing the workflow by the alert thread

The workflow stamps its BalanceThreadId search attribute with params.thread_id — the alert thread — so the owner's reply on that thread routes back here through W3. Every post (the options offer, the could-not-cover note, clarifications, the apply/decline confirmation) goes through one _send helper that replies on that thread; the old _open / self._thread_id indirection is removed.

Stamp the alert thread, then post the options as a reply on it.

src/ynab_agent/workflow/balance_workflow.py · 253 lines
src/ynab_agent/workflow/balance_workflow.py253 lines · Python
⋯ 121 lines hidden (lines 1–121)
1"""W7 · the budget-balance workflow — propose, confirm by NL, apply (SPEC §8).
2 
3One short-lived workflow per (overspent category, period), started by the W6
4monitor (id ``balance-offer-{category}-{period}``, ``REJECT_DUPLICATE`` so it is
5one offer per period). It asks the model for coverage options, posts them as a
6reply on the *same* overspend-alert thread, stamps a ``BalanceThreadId`` search
7attribute so W3 routes the reply back here, and acts on it: apply
8the chosen plan, decline, or answer a clarifying question and keep waiting.
9 
10The apply is the careful part. The workflow reads a baseline snapshot, validates
11the moves against real funds and the hard floor, then computes each category's
12*absolute* target ``budgeted`` here, in durable workflow state — so re-driving a
13write activity re-sets the same value and never double-applies. Every move is
14read-back verified before the workflow calls it done.
15 
16The seam mirrors the offer/W2 drivers: ``workflow.*`` for all clock reads, side
17effects behind activities, pure domain types for the branch — so the model and
18mail stacks never enter this sandbox.
19"""
20 
21from __future__ import annotations
22 
23from collections import deque
24from datetime import timedelta
25 
26from temporalio import workflow
27from temporalio.common import SearchAttributeKey
28from temporalio.exceptions import ActivityError
29 
30from ynab_agent.workflow.balance_types import (
31 BALANCE_PATIENCE,
32 BALANCE_THREAD_ID,
34 
35_BALANCE_THREAD_ID = SearchAttributeKey.for_keyword(BALANCE_THREAD_ID)
36 
37with workflow.unsafe.imports_passed_through():
38 from ynab_agent.agentic.compose import (
39 render_balance_applied,
40 render_balance_could_not_cover,
41 render_balance_declined,
42 render_balance_failed,
43 render_balance_options,
44 )
45 from ynab_agent.budget.balance import (
46 ApplyMoves,
47 BudgetMove,
48 DeclineBalance,
49 OptionRejection,
50 check_moves,
51 move_targets,
52 )
53 from ynab_agent.dispatch.classify import InboundMessage
54 from ynab_agent.domain.money import Money
55 from ynab_agent.workflow import alert_activities, balance_activities
56 from ynab_agent.workflow.alerting import build_failure_alert
57 from ynab_agent.workflow.balance_types import (
58 BalanceParams,
59 BalanceResult,
60 )
61 from ynab_agent.workflow.constants import (
62 ACTIVITY_RETRY,
63 ACTIVITY_TIMEOUT,
64 ALERT_BUDGET,
65 ALERT_RETRY,
66 ALERT_TIMEOUT,
67 )
68 
69 
70def _reason(rejection: OptionRejection) -> str:
71 """The owner-facing reason an approved plan can't be applied."""
72 if rejection is OptionRejection.OVER_CEILING:
73 return "one of the moves is larger than I can make automatically"
74 if rejection is OptionRejection.INSUFFICIENT_SOURCE:
75 return "a source no longer has enough to cover it"
76 return "the plan didn't add up"
77 
78 
79@workflow.defn
80class BudgetBalanceWorkflow:
81 """One overspent category's coverage offer for a budget period."""
82 
83 def __init__(self) -> None:
84 """Start with no reply buffered; ``run`` posts the offer and waits."""
85 self._responses: deque[InboundMessage] = deque()
86 self._clarifications = 0
87 
88 @workflow.signal
89 def submit_response(self, message: InboundMessage) -> None:
90 """The owner replied on the alert thread (delivered by W3)."""
91 self._responses.append(message)
92 
93 @workflow.run
94 async def run(self, params: BalanceParams) -> BalanceResult:
95 """Run the offer, paging the owner on a terminal failure."""
96 try:
97 return await self._run(params)
98 except ActivityError as exc:
99 assessment = params.assessment
100 await workflow.execute_activity(
101 alert_activities.alert_failure,
102 build_failure_alert(
103 key=f"balance-{assessment.category}-{params.period}",
104 context=f"budget balance for {assessment.name}",
105 exc=exc,
106 ),
107 start_to_close_timeout=ALERT_TIMEOUT,
108 schedule_to_close_timeout=ALERT_BUDGET,
109 retry_policy=ALERT_RETRY,
110 )
111 raise
112 
113 async def _run(self, params: BalanceParams) -> BalanceResult:
114 assessment = params.assessment
115 label = f"{assessment.category}-{params.period}"
116 options = await workflow.execute_activity(
117 balance_activities.propose_balance_options,
118 params,
119 start_to_close_timeout=ACTIVITY_TIMEOUT,
120 retry_policy=ACTIVITY_RETRY,
121 )
122 if not options:
123 await self._send(
124 params,
125 render_balance_could_not_cover(assessment.name),
126 f"yb-nocover-{label}",
127 )
128 return BalanceResult(outcome="could-not-cover")
129 
130 # Index this workflow by the overspend-alert thread so the owner's reply
131 # routes back here (W3 → BalanceThreadId), then post the options as a
132 # reply on that same thread — the W6→W7 tie, one conversation (SPEC §8).
133 workflow.upsert_search_attributes(
134 [_BALANCE_THREAD_ID.value_set(params.thread_id)]
135 )
136 await self._send(
137 params,
138 render_balance_options(assessment.name, tuple(options)),
139 f"yb-cover-{label}",
140 )
⋯ 113 lines hidden (lines 141–253)
141 
142 deadline = workflow.now() + BALANCE_PATIENCE
143 while True:
144 timeout = deadline - workflow.now()
145 if timeout < timedelta(0):
146 timeout = timedelta(0)
147 try:
148 await workflow.wait_condition(
149 lambda: len(self._responses) > 0, timeout=timeout
150 )
151 except TimeoutError:
152 return BalanceResult(outcome="timed-out")
153 message = self._responses.popleft()
154 outcome = await workflow.execute_activity(
155 balance_activities.interpret_balance_reply,
156 args=[params, message.body, options],
157 start_to_close_timeout=ACTIVITY_TIMEOUT,
158 retry_policy=ACTIVITY_RETRY,
159 )
160 if isinstance(outcome, DeclineBalance):
161 await self._send(
162 params,
163 render_balance_declined(assessment.name),
164 f"ybalance-declined-{label}",
165 )
166 return BalanceResult(outcome="declined")
167 if isinstance(outcome, ApplyMoves):
168 return await self._apply(params, outcome.moves, label)
169 # ClarifyBalance: answer, then keep waiting for a clearer reply.
170 self._clarifications += 1
171 await self._send(
172 params,
173 outcome.question,
174 f"ybalance-clarify-{label}-{self._clarifications}",
175 )
176 
177 async def _apply(
178 self,
179 params: BalanceParams,
180 moves: tuple[BudgetMove, ...],
181 label: str,
182 ) -> BalanceResult:
183 """Validate, write targets, verify, audit, confirm (SPEC §8)."""
184 assessment = params.assessment
185 state = await workflow.execute_activity(
186 balance_activities.read_budget_state,
187 start_to_close_timeout=ACTIVITY_TIMEOUT,
188 retry_policy=ACTIVITY_RETRY,
189 )
190 rejection = check_moves(moves, state.available)
191 if rejection is not None:
192 await self._send(
193 params,
194 render_balance_failed(assessment.name, _reason(rejection)),
195 f"ybalance-failed-{label}",
196 )
197 return BalanceResult(outcome="rejected", detail=rejection.value)
198 
199 # Targets computed here, in durable state, from the baseline snapshot —
200 # so a write retry re-sets the same absolute value (no double-apply).
201 targets = move_targets(moves, state.budgeted)
202 verified = True
203 for category, target in targets.items():
204 ok = await workflow.execute_activity(
205 balance_activities.set_category_budgeted,
206 args=[str(category), target],
207 start_to_close_timeout=ACTIVITY_TIMEOUT,
208 retry_policy=ACTIVITY_RETRY,
209 )
210 verified = verified and ok
211 if not verified:
212 await self._send(
213 params,
214 render_balance_failed(
215 assessment.name, "the change didn't take effect cleanly"
216 ),
217 f"ybalance-failed-{label}-verify",
218 )
219 return BalanceResult(outcome="verify-failed")
220 
221 await workflow.execute_activity(
222 balance_activities.log_budget_moves,
223 args=[list(moves), params.period],
224 start_to_close_timeout=ACTIVITY_TIMEOUT,
225 retry_policy=ACTIVITY_RETRY,
226 )
227 total = Money.zero()
228 for move in moves:
229 total = total + move.amount
230 await self._send(
231 params,
232 render_balance_applied(assessment.name, total),
233 f"ybalance-applied-{label}",
234 )
235 return BalanceResult(outcome="applied")
236 
237 async def _send(
238 self, params: BalanceParams, body: str, seq_label: str
239 ) -> None:
240 """Reply on the overspend-alert thread, addressed to the owners (§8).
241 
242 Every balancer message — options, clarifications, the apply/decline
243 confirmation — replies on the W6 alert thread (``params.thread_id``) so
244 the owner sees one conversation. The activity sets the recipients
245 explicitly, so a reply on a thread the agent last spoke on still reaches
246 the owner. Idempotent on ``seq_label``.
247 """
248 await workflow.execute_activity(
249 balance_activities.send_balance_email,
250 args=[params.thread_id, body, seq_label],
251 start_to_close_timeout=ACTIVITY_TIMEOUT,
252 retry_policy=ACTIVITY_RETRY,
253 )

One reply helper, always on params.thread_id.

src/ynab_agent/workflow/balance_workflow.py · 253 lines
src/ynab_agent/workflow/balance_workflow.py253 lines · Python
⋯ 236 lines hidden (lines 1–236)
1"""W7 · the budget-balance workflow — propose, confirm by NL, apply (SPEC §8).
2 
3One short-lived workflow per (overspent category, period), started by the W6
4monitor (id ``balance-offer-{category}-{period}``, ``REJECT_DUPLICATE`` so it is
5one offer per period). It asks the model for coverage options, posts them as a
6reply on the *same* overspend-alert thread, stamps a ``BalanceThreadId`` search
7attribute so W3 routes the reply back here, and acts on it: apply
8the chosen plan, decline, or answer a clarifying question and keep waiting.
9 
10The apply is the careful part. The workflow reads a baseline snapshot, validates
11the moves against real funds and the hard floor, then computes each category's
12*absolute* target ``budgeted`` here, in durable workflow state — so re-driving a
13write activity re-sets the same value and never double-applies. Every move is
14read-back verified before the workflow calls it done.
15 
16The seam mirrors the offer/W2 drivers: ``workflow.*`` for all clock reads, side
17effects behind activities, pure domain types for the branch — so the model and
18mail stacks never enter this sandbox.
19"""
20 
21from __future__ import annotations
22 
23from collections import deque
24from datetime import timedelta
25 
26from temporalio import workflow
27from temporalio.common import SearchAttributeKey
28from temporalio.exceptions import ActivityError
29 
30from ynab_agent.workflow.balance_types import (
31 BALANCE_PATIENCE,
32 BALANCE_THREAD_ID,
34 
35_BALANCE_THREAD_ID = SearchAttributeKey.for_keyword(BALANCE_THREAD_ID)
36 
37with workflow.unsafe.imports_passed_through():
38 from ynab_agent.agentic.compose import (
39 render_balance_applied,
40 render_balance_could_not_cover,
41 render_balance_declined,
42 render_balance_failed,
43 render_balance_options,
44 )
45 from ynab_agent.budget.balance import (
46 ApplyMoves,
47 BudgetMove,
48 DeclineBalance,
49 OptionRejection,
50 check_moves,
51 move_targets,
52 )
53 from ynab_agent.dispatch.classify import InboundMessage
54 from ynab_agent.domain.money import Money
55 from ynab_agent.workflow import alert_activities, balance_activities
56 from ynab_agent.workflow.alerting import build_failure_alert
57 from ynab_agent.workflow.balance_types import (
58 BalanceParams,
59 BalanceResult,
60 )
61 from ynab_agent.workflow.constants import (
62 ACTIVITY_RETRY,
63 ACTIVITY_TIMEOUT,
64 ALERT_BUDGET,
65 ALERT_RETRY,
66 ALERT_TIMEOUT,
67 )
68 
69 
70def _reason(rejection: OptionRejection) -> str:
71 """The owner-facing reason an approved plan can't be applied."""
72 if rejection is OptionRejection.OVER_CEILING:
73 return "one of the moves is larger than I can make automatically"
74 if rejection is OptionRejection.INSUFFICIENT_SOURCE:
75 return "a source no longer has enough to cover it"
76 return "the plan didn't add up"
77 
78 
79@workflow.defn
80class BudgetBalanceWorkflow:
81 """One overspent category's coverage offer for a budget period."""
82 
83 def __init__(self) -> None:
84 """Start with no reply buffered; ``run`` posts the offer and waits."""
85 self._responses: deque[InboundMessage] = deque()
86 self._clarifications = 0
87 
88 @workflow.signal
89 def submit_response(self, message: InboundMessage) -> None:
90 """The owner replied on the alert thread (delivered by W3)."""
91 self._responses.append(message)
92 
93 @workflow.run
94 async def run(self, params: BalanceParams) -> BalanceResult:
95 """Run the offer, paging the owner on a terminal failure."""
96 try:
97 return await self._run(params)
98 except ActivityError as exc:
99 assessment = params.assessment
100 await workflow.execute_activity(
101 alert_activities.alert_failure,
102 build_failure_alert(
103 key=f"balance-{assessment.category}-{params.period}",
104 context=f"budget balance for {assessment.name}",
105 exc=exc,
106 ),
107 start_to_close_timeout=ALERT_TIMEOUT,
108 schedule_to_close_timeout=ALERT_BUDGET,
109 retry_policy=ALERT_RETRY,
110 )
111 raise
112 
113 async def _run(self, params: BalanceParams) -> BalanceResult:
114 assessment = params.assessment
115 label = f"{assessment.category}-{params.period}"
116 options = await workflow.execute_activity(
117 balance_activities.propose_balance_options,
118 params,
119 start_to_close_timeout=ACTIVITY_TIMEOUT,
120 retry_policy=ACTIVITY_RETRY,
121 )
122 if not options:
123 await self._send(
124 params,
125 render_balance_could_not_cover(assessment.name),
126 f"yb-nocover-{label}",
127 )
128 return BalanceResult(outcome="could-not-cover")
129 
130 # Index this workflow by the overspend-alert thread so the owner's reply
131 # routes back here (W3 → BalanceThreadId), then post the options as a
132 # reply on that same thread — the W6→W7 tie, one conversation (SPEC §8).
133 workflow.upsert_search_attributes(
134 [_BALANCE_THREAD_ID.value_set(params.thread_id)]
135 )
136 await self._send(
137 params,
138 render_balance_options(assessment.name, tuple(options)),
139 f"yb-cover-{label}",
140 )
141 
142 deadline = workflow.now() + BALANCE_PATIENCE
143 while True:
144 timeout = deadline - workflow.now()
145 if timeout < timedelta(0):
146 timeout = timedelta(0)
147 try:
148 await workflow.wait_condition(
149 lambda: len(self._responses) > 0, timeout=timeout
150 )
151 except TimeoutError:
152 return BalanceResult(outcome="timed-out")
153 message = self._responses.popleft()
154 outcome = await workflow.execute_activity(
155 balance_activities.interpret_balance_reply,
156 args=[params, message.body, options],
157 start_to_close_timeout=ACTIVITY_TIMEOUT,
158 retry_policy=ACTIVITY_RETRY,
159 )
160 if isinstance(outcome, DeclineBalance):
161 await self._send(
162 params,
163 render_balance_declined(assessment.name),
164 f"ybalance-declined-{label}",
165 )
166 return BalanceResult(outcome="declined")
167 if isinstance(outcome, ApplyMoves):
168 return await self._apply(params, outcome.moves, label)
169 # ClarifyBalance: answer, then keep waiting for a clearer reply.
170 self._clarifications += 1
171 await self._send(
172 params,
173 outcome.question,
174 f"ybalance-clarify-{label}-{self._clarifications}",
175 )
176 
177 async def _apply(
178 self,
179 params: BalanceParams,
180 moves: tuple[BudgetMove, ...],
181 label: str,
182 ) -> BalanceResult:
183 """Validate, write targets, verify, audit, confirm (SPEC §8)."""
184 assessment = params.assessment
185 state = await workflow.execute_activity(
186 balance_activities.read_budget_state,
187 start_to_close_timeout=ACTIVITY_TIMEOUT,
188 retry_policy=ACTIVITY_RETRY,
189 )
190 rejection = check_moves(moves, state.available)
191 if rejection is not None:
192 await self._send(
193 params,
194 render_balance_failed(assessment.name, _reason(rejection)),
195 f"ybalance-failed-{label}",
196 )
197 return BalanceResult(outcome="rejected", detail=rejection.value)
198 
199 # Targets computed here, in durable state, from the baseline snapshot —
200 # so a write retry re-sets the same absolute value (no double-apply).
201 targets = move_targets(moves, state.budgeted)
202 verified = True
203 for category, target in targets.items():
204 ok = await workflow.execute_activity(
205 balance_activities.set_category_budgeted,
206 args=[str(category), target],
207 start_to_close_timeout=ACTIVITY_TIMEOUT,
208 retry_policy=ACTIVITY_RETRY,
209 )
210 verified = verified and ok
211 if not verified:
212 await self._send(
213 params,
214 render_balance_failed(
215 assessment.name, "the change didn't take effect cleanly"
216 ),
217 f"ybalance-failed-{label}-verify",
218 )
219 return BalanceResult(outcome="verify-failed")
220 
221 await workflow.execute_activity(
222 balance_activities.log_budget_moves,
223 args=[list(moves), params.period],
224 start_to_close_timeout=ACTIVITY_TIMEOUT,
225 retry_policy=ACTIVITY_RETRY,
226 )
227 total = Money.zero()
228 for move in moves:
229 total = total + move.amount
230 await self._send(
231 params,
232 render_balance_applied(assessment.name, total),
233 f"ybalance-applied-{label}",
234 )
235 return BalanceResult(outcome="applied")
236 
237 async def _send(
238 self, params: BalanceParams, body: str, seq_label: str
239 ) -> None:
240 """Reply on the overspend-alert thread, addressed to the owners (§8).
241 
242 Every balancer message — options, clarifications, the apply/decline
243 confirmation — replies on the W6 alert thread (``params.thread_id``) so
244 the owner sees one conversation. The activity sets the recipients
245 explicitly, so a reply on a thread the agent last spoke on still reaches
246 the owner. Idempotent on ``seq_label``.
247 """
248 await workflow.execute_activity(
249 balance_activities.send_balance_email,
250 args=[params.thread_id, body, seq_label],
251 start_to_close_timeout=ACTIVITY_TIMEOUT,
252 retry_policy=ACTIVITY_RETRY,
⋯ 1 line hidden (lines 253–253)
253 )

Tests: a regression guard and a live proof

The workflow test is the regression guard. It runs the balancer through an offer-then-decline, then asserts two things the old code would fail: every balancer email replied on the alert thread ("thr-overspend"), and BalanceThreadId is stamped with that same thread. On the old code these were a freshly opened thread id instead.

Offer + decline: posts and the search attribute land on the alert thread.

tests/workflow/test_balance_workflow.py · 343 lines
tests/workflow/test_balance_workflow.py343 lines · Python
⋯ 306 lines hidden (lines 1–306)
1"""End-to-end tests for the W7 budget-balance workflow (SPEC §8).
2 
3Exercised on the time-skipping server with mock activities. The workflow stamps
4a ``BalanceThreadId`` search attribute on offer, which the test server needs
5registered (as the real cluster does via manage/search-attributes.yaml). The
6apply path runs the *real* pure guard (``check_moves`` / ``move_targets``), so
7the floor-refusal and target-math tests are genuine.
8"""
9 
10from __future__ import annotations
11 
12from typing import TYPE_CHECKING
13 
14from temporalio import activity
15from temporalio.api.enums.v1 import IndexedValueType
16from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest
17from temporalio.common import SearchAttributeKey
18from temporalio.testing import WorkflowEnvironment
19from temporalio.worker import Worker
20 
21from ynab_agent.budget.balance import (
22 ApplyMoves,
23 BalanceOption,
24 BalanceOutcome,
25 BudgetMove,
26 ClarifyBalance,
27 DeclineBalance,
29from ynab_agent.budget.overspend import OverspendAssessment, OverspendVerdict
30from ynab_agent.dispatch.classify import InboundMessage
31from ynab_agent.domain.ids import CategoryId, MessageId, ThreadId
32from ynab_agent.domain.money import Money
33from ynab_agent.workflow.balance_types import (
34 BalanceParams,
35 BalanceResult,
36 BudgetState,
37 balance_workflow_id,
39from ynab_agent.workflow.balance_workflow import BudgetBalanceWorkflow
40from ynab_agent.workflow.runtime import DATA_CONVERTER
41 
42if TYPE_CHECKING:
43 from collections.abc import Callable
44 
45_TASK_QUEUE = "balance-wf-test"
46_PERIOD = "2026-06"
47 
48 
49def _assessment() -> OverspendAssessment:
50 return OverspendAssessment(
51 category=CategoryId("dining"),
52 name="Dining Out",
53 verdict=OverspendVerdict.TRENDING_OVER,
54 budgeted=Money.from_currency("400"),
55 spent=Money.from_currency("250"),
56 projected=Money.from_currency("520"),
57 )
58 
59 
60def _params() -> BalanceParams:
61 return BalanceParams(
62 assessment=_assessment(), thread_id="thr-overspend", period=_PERIOD
63 )
64 
65 
66def _option(amount: str = "120") -> BalanceOption:
67 return BalanceOption(
68 label="From Buffer",
69 moves=(
70 BudgetMove(
71 source=CategoryId("buffer"),
72 destination=CategoryId("dining"),
73 amount=Money.from_currency(amount),
74 ),
75 ),
76 rationale="Buffer has plenty.",
77 )
78 
79 
80def _apply_moves(amount: str) -> ApplyMoves:
81 return ApplyMoves(
82 moves=(
83 BudgetMove(
84 source=CategoryId("buffer"),
85 destination=CategoryId("dining"),
86 amount=Money.from_currency(amount),
87 ),
88 )
89 )
90 
91 
92def _state() -> BudgetState:
93 return BudgetState(
94 available={
95 CategoryId("buffer"): Money.from_currency("500"),
96 CategoryId("dining"): Money.from_currency("-20"),
97 },
98 budgeted={
99 CategoryId("buffer"): Money.from_currency("500"),
100 CategoryId("dining"): Money.from_currency("400"),
101 },
102 )
103 
104 
105def _reply(text: str = "do it") -> InboundMessage:
106 return InboundMessage(
107 message_id=MessageId("m1"),
108 from_address="matthew@example.com",
109 subject="re: Dining Out: trending over budget",
110 body=text,
111 thread_id=ThreadId("thr-overspend"),
112 signature_verified=True,
113 )
114 
115 
116class _Recorder:
117 def __init__(self) -> None:
118 self.sends: list[str] = []
119 self.threads: list[str] = [] # thread each balance email replied on
120 self.sets: list[tuple[str, str]] = [] # (category, "$amount")
121 self.logged = 0
122 
123 
124def _activities(
125 *,
126 options: list[BalanceOption],
127 outcomes: list[BalanceOutcome],
128 rec: _Recorder,
129 set_ok: bool = True,
130) -> list[Callable[..., object]]:
131 pending = list(outcomes)
132 
133 @activity.defn(name="propose_balance_options")
134 async def propose_balance_options(
135 params: BalanceParams,
136 ) -> list[BalanceOption]:
137 return options
138 
139 @activity.defn(name="interpret_balance_reply")
140 async def interpret_balance_reply(
141 params: BalanceParams, reply_text: str, opts: list[BalanceOption]
142 ) -> BalanceOutcome:
143 return pending.pop(0)
144 
145 @activity.defn(name="read_budget_state")
146 async def read_budget_state() -> BudgetState:
147 return _state()
148 
149 @activity.defn(name="set_category_budgeted")
150 async def set_category_budgeted(category_id: str, target: Money) -> bool:
151 rec.sets.append((category_id, str(target)))
152 return set_ok
153 
154 @activity.defn(name="log_budget_moves")
155 async def log_budget_moves(moves: list[BudgetMove], period: str) -> None:
156 rec.logged += 1
157 
158 @activity.defn(name="send_balance_email")
159 async def send_balance_email(
160 thread_id: str, body: str, seq_label: str
161 ) -> None:
162 rec.sends.append(seq_label)
163 rec.threads.append(thread_id)
164 
165 @activity.defn(name="alert_failure")
166 async def alert_failure(alert: object) -> None:
167 return None
168 
169 return [
170 propose_balance_options,
171 interpret_balance_reply,
172 read_budget_state,
173 set_category_budgeted,
174 log_budget_moves,
175 send_balance_email,
176 alert_failure,
177 ]
178 
179 
180async def _start_env() -> WorkflowEnvironment:
181 env = await WorkflowEnvironment.start_time_skipping(
182 data_converter=DATA_CONVERTER
183 )
184 await env.client.operator_service.add_search_attributes(
185 AddSearchAttributesRequest(
186 namespace="default",
187 search_attributes={
188 "BalanceThreadId": IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD
189 },
190 )
191 )
192 return env
193 
194 
195async def _run(
196 *,
197 options: list[BalanceOption],
198 outcomes: list[BalanceOutcome],
199 replies: int,
200 set_ok: bool = True,
201) -> tuple[BalanceResult, _Recorder]:
202 rec = _Recorder()
203 async with (
204 await _start_env() as env,
205 Worker(
206 env.client,
207 task_queue=_TASK_QUEUE,
208 workflows=[BudgetBalanceWorkflow],
209 activities=_activities(
210 options=options, outcomes=outcomes, rec=rec, set_ok=set_ok
211 ),
212 ),
213 ):
214 handle = await env.client.start_workflow(
215 BudgetBalanceWorkflow.run,
216 _params(),
217 id=balance_workflow_id("dining", _PERIOD),
218 task_queue=_TASK_QUEUE,
219 )
220 for _ in range(replies):
221 await handle.signal(BudgetBalanceWorkflow.submit_response, _reply())
222 result = await handle.result()
223 return result, rec
224 
225 
226async def test_apply_writes_targets_verifies_and_confirms() -> None:
227 result, rec = await _run(
228 options=[_option()], outcomes=[_apply_moves("120")], replies=1
229 )
230 assert result.outcome == "applied"
231 # Destination raised to $520, source lowered to $380 (absolute targets).
232 assert ("dining", "$520.00") in rec.sets
233 assert ("buffer", "$380.00") in rec.sets
234 assert rec.logged == 1
235 assert any(s.startswith("ybalance-applied") for s in rec.sends)
236 
237 
238async def test_natural_language_modified_plan_is_applied() -> None:
239 # The owner asked to cover only $50; the workflow applies the modified plan.
240 result, rec = await _run(
241 options=[_option()], outcomes=[_apply_moves("50")], replies=1
242 )
243 assert result.outcome == "applied"
244 assert ("dining", "$450.00") in rec.sets # 400 + 50
245 assert ("buffer", "$450.00") in rec.sets # 500 - 50
246 
247 
248async def test_over_ceiling_move_is_refused_by_the_floor() -> None:
249 # A $600 move exceeds the $500 per-move ceiling: the real check_moves guard
250 # refuses it, nothing is written, even though the owner "approved" it.
251 result, rec = await _run(
252 options=[_option()], outcomes=[_apply_moves("600")], replies=1
253 )
254 assert result.outcome == "rejected"
255 assert result.detail == "over_ceiling"
256 assert rec.sets == [] # no writes
257 assert any(s.startswith("ybalance-failed") for s in rec.sends)
258 
259 
260async def test_decline_sends_a_note_and_writes_nothing() -> None:
261 result, rec = await _run(
262 options=[_option()], outcomes=[DeclineBalance()], replies=1
263 )
264 assert result.outcome == "declined"
265 assert rec.sets == []
266 assert any(s.startswith("ybalance-declined") for s in rec.sends)
267 
268 
269async def test_clarify_then_apply() -> None:
270 result, rec = await _run(
271 options=[_option()],
272 outcomes=[ClarifyBalance(question="From which?"), _apply_moves("120")],
273 replies=2,
274 )
275 assert result.outcome == "applied"
276 assert any(s.startswith("ybalance-clarify") for s in rec.sends)
277 assert ("dining", "$520.00") in rec.sets
278 
279 
280async def test_no_feasible_options_sends_could_not_cover() -> None:
281 result, rec = await _run(options=[], outcomes=[], replies=0)
282 assert result.outcome == "could-not-cover"
283 assert any(s.startswith("yb-nocover") for s in rec.sends)
284 assert rec.sets == []
285 
286 
287async def test_verify_failure_reports_and_does_not_confirm() -> None:
288 result, rec = await _run(
289 options=[_option()],
290 outcomes=[_apply_moves("120")],
291 replies=1,
292 set_ok=False,
293 )
294 assert result.outcome == "verify-failed"
295 assert any("failed" in s for s in rec.sends)
296 assert not any(s.startswith("ybalance-applied") for s in rec.sends)
297 
298 
299async def test_no_reply_times_out() -> None:
300 result, rec = await _run(options=[_option()], outcomes=[], replies=0)
301 # Wait — replies=0 with options posts the offer then the patience window
302 # elapses (time-skipped) with no answer.
303 assert result.outcome == "timed-out"
304 assert any(s.startswith("yb-cover") for s in rec.sends)
305 
306 
307async def test_offer_posts_and_routes_on_the_alert_thread() -> None:
308 """The balancer posts on the W6 alert thread and indexes itself by it.
309 
310 Every balancer email replies on ``thread_id`` (the overspend-alert thread,
311 ``"thr-overspend"``), and ``BalanceThreadId`` is stamped with that same
312 thread — so the owner sees one conversation and a reply on it routes back
313 here (the W6→W7 tie, SPEC §8). This is the regression guard against posting
314 on a freshly opened thread instead.
315 """
316 rec = _Recorder()
317 async with (
318 await _start_env() as env,
319 Worker(
320 env.client,
321 task_queue=_TASK_QUEUE,
322 workflows=[BudgetBalanceWorkflow],
323 activities=_activities(
324 options=[_option()], outcomes=[DeclineBalance()], rec=rec
325 ),
326 ),
327 ):
328 handle = await env.client.start_workflow(
329 BudgetBalanceWorkflow.run,
330 _params(),
331 id=balance_workflow_id("dining", _PERIOD),
332 task_queue=_TASK_QUEUE,
333 )
334 await handle.signal(BudgetBalanceWorkflow.submit_response, _reply())
335 await handle.result()
336 desc = await handle.describe()
337 stamped = desc.typed_search_attributes.get(
338 SearchAttributeKey.for_keyword("BalanceThreadId")
339 )
340 assert stamped == "thr-overspend"
341 # The options offer and the decline confirmation both replied on the alert
342 # thread — never on a separately opened one.
343 assert rec.threads == ["thr-overspend", "thr-overspend"]

Offline tests run against a fake backend, so they prove the code calls reply with the right recipients but not how AgentMail behaves. A gated live test closes that gap: it reproduces the W6→W7 shape (open an alert, then reply on it), then fetches the sent reply from the real server.

Live: the reply threads, and its To is the owner — not the agent.

tests/mail/test_mail_client.py · 294 lines
tests/mail/test_mail_client.py294 lines · Python
⋯ 252 lines hidden (lines 1–252)
1"""Tests for the AgentMail client (SPEC §5).
2 
3Offline tests inject a fake backend (no network). One opt-in live test sends a
4real email through AgentMail and is skipped unless ``YNAB_AGENT_LIVE_EMAIL`` is
5set, so the gate stays offline.
6 
7The fake backend models the one property the client relies on for store-free
8idempotency: labels are sticky per thread, so ``find_thread(label)`` re-finds an
9already-opened thread (open dedup) or an already-sent action (send dedup).
10"""
11 
12from __future__ import annotations
13 
14import os
15 
16import pytest
17 
18from ynab_agent.mail.client import MailClient, OutboundEmail, SentEmail
19 
20 
21class _FakeBackend:
22 """An in-memory AgentMail: threads, sticky labels, last-message tracking."""
23 
24 def __init__(self) -> None:
25 self.calls: list[tuple[str, ...]] = []
26 self.reply_tos: list[list[str] | None] = [] # `to` per send_reply call
27 self._labels: dict[str, str] = {} # label -> thread_id (first to carry)
28 self._last: dict[str, str] = {} # thread_id -> last message_id
29 self._msg_thread: dict[str, str] = {} # message_id -> thread_id
30 self._archived: set[str] = set()
31 self._seq = 0
32 
33 def _next(self, prefix: str) -> str:
34 self._seq += 1
35 return f"{prefix}{self._seq}"
36 
37 def _register(self, thread_id: str, labels: list[str] | None) -> None:
38 for label in labels or []:
39 self._labels.setdefault(label, thread_id)
40 
41 def send_new(
42 self,
43 inbox_id: str,
44 to: list[str],
45 subject: str,
46 text: str,
47 labels: list[str] | None = None,
48 ) -> SentEmail:
49 self.calls.append(("new", inbox_id, ",".join(to), subject, text))
50 thread_id = self._next("t")
51 message_id = self._next("m")
52 self._last[thread_id] = message_id
53 self._msg_thread[message_id] = thread_id
54 self._register(thread_id, labels)
55 return SentEmail(message_id=message_id, thread_id=thread_id)
56 
57 def send_reply(
58 self,
59 inbox_id: str,
60 message_id: str,
61 text: str,
62 labels: list[str] | None = None,
63 to: list[str] | None = None,
64 ) -> SentEmail:
65 self.calls.append(("reply", inbox_id, message_id, text))
66 self.reply_tos.append(to)
67 thread_id = self._msg_thread.get(message_id, "t1")
68 new_id = self._next("m")
69 self._last[thread_id] = new_id
70 self._msg_thread[new_id] = thread_id
71 self._register(thread_id, labels)
72 return SentEmail(message_id=new_id, thread_id=thread_id)
73 
74 def find_thread(self, inbox_id: str, label: str) -> str | None:
75 return self._labels.get(label)
76 
77 def latest_message_id(self, inbox_id: str, thread_id: str) -> str | None:
78 return self._last.get(thread_id)
79 
80 def archive(self, inbox_id: str, thread_id: str) -> None:
81 self.calls.append(("archive", inbox_id, thread_id))
82 self._archived.add(thread_id)
83 
84 
85def test_new_thread_routes_to_send_new() -> None:
86 backend = _FakeBackend()
87 out = MailClient(backend).send(
88 OutboundEmail(
89 inbox_id="ib",
90 to=("wife@example.com",),
91 subject="[YNAB] $4.50",
92 text="best guess: Dining",
93 )
94 )
95 assert out.thread_id
96 assert backend.calls[0][0] == "new"
97 assert backend.calls[0][3] == "[YNAB] $4.50"
98 
99 
100def test_reply_routes_to_send_reply() -> None:
101 backend = _FakeBackend()
102 out = MailClient(backend).send(
103 OutboundEmail(
104 inbox_id="ib",
105 subject="[YNAB] $4.50",
106 text="applied, thanks",
107 reply_to_message_id="m0",
108 )
109 )
110 assert out.message_id
111 assert backend.calls[0] == ("reply", "ib", "m0", "applied, thanks")
112 
113 
114def test_new_thread_without_recipients_is_rejected() -> None:
115 with pytest.raises(ValueError, match="recipient"):
116 OutboundEmail(inbox_id="ib", subject="x", text="y")
117 
118 
119def test_open_thread_sends_once_then_dedups() -> None:
120 backend = _FakeBackend()
121 client = MailClient(backend)
122 first = client.open_thread(
123 inbox_id="ib",
124 to=["wife@example.com"],
125 subject="[YNAB] $4.50",
126 body="best guess: Dining",
127 txn_label="yatxn-abc",
128 )
129 second = client.open_thread(
130 inbox_id="ib",
131 to=["wife@example.com"],
132 subject="[YNAB] $4.50",
133 body="best guess: Dining",
134 txn_label="yatxn-abc",
135 )
136 assert first == second
137 # Only the first call actually sent a new thread (idempotent open).
138 assert [c[0] for c in backend.calls] == ["new"]
139 
140 
141def test_send_on_thread_dedups_on_seq_label() -> None:
142 backend = _FakeBackend()
143 client = MailClient(backend)
144 thread_id = client.open_thread(
145 inbox_id="ib",
146 to=["wife@example.com"],
147 subject="[YNAB] $4.50",
148 body="proposal",
149 txn_label="yatxn-abc",
150 )
151 sent_first = client.send_on_thread(
152 inbox_id="ib",
153 thread_id=thread_id,
154 body="a reminder",
155 seq_label="yaseq-abc-2",
156 )
157 sent_again = client.send_on_thread(
158 inbox_id="ib",
159 thread_id=thread_id,
160 body="a reminder",
161 seq_label="yaseq-abc-2",
162 )
163 assert sent_first is True
164 assert sent_again is False
165 assert [c[0] for c in backend.calls].count("reply") == 1
166 
167 
168def test_send_on_thread_false_when_thread_empty() -> None:
169 backend = _FakeBackend()
170 sent = MailClient(backend).send_on_thread(
171 inbox_id="ib",
172 thread_id="t-unknown",
173 body="hi",
174 seq_label="yaseq-x-1",
175 )
176 assert sent is False
177 assert backend.calls == []
178 
179 
180def test_send_on_thread_forwards_recipients_to_the_reply() -> None:
181 # The balancer replies on the W6 alert thread but must address the owners
182 # explicitly, else AgentMail loops the reply back to the agent (SPEC §8).
183 backend = _FakeBackend()
184 client = MailClient(backend)
185 thread_id = client.open_thread(
186 inbox_id="ib",
187 to=["wife@example.com"],
188 subject="Dining: over budget",
189 body="alert",
190 txn_label="yaspend-dining",
191 )
192 sent = client.send_on_thread(
193 inbox_id="ib",
194 thread_id=thread_id,
195 body="cover the overspend?",
196 seq_label="yb-cover-dining",
197 to=["wife@example.com"],
198 )
199 assert sent is True
200 assert backend.reply_tos == [["wife@example.com"]]
201 
202 
203def test_send_on_thread_without_recipients_omits_the_override() -> None:
204 # The other callers (W2, the autonomy offer) reply only after the owner has
205 # spoken, so they leave AgentMail's default reply-addressing in place.
206 backend = _FakeBackend()
207 client = MailClient(backend)
208 thread_id = client.open_thread(
209 inbox_id="ib",
210 to=["wife@example.com"],
211 subject="[YNAB] $4.50",
212 body="proposal",
213 txn_label="yatxn-abc",
214 )
215 client.send_on_thread(
216 inbox_id="ib",
217 thread_id=thread_id,
218 body="a reminder",
219 seq_label="yaseq-abc-2",
220 )
221 assert backend.reply_tos == [None]
222 
223 
224def test_close_archives_the_thread() -> None:
225 backend = _FakeBackend()
226 MailClient(backend).close(inbox_id="ib", thread_id="t1")
227 assert backend.calls == [("archive", "ib", "t1")]
228 
229 
230@pytest.mark.skipif(
231 not os.environ.get("YNAB_AGENT_LIVE_EMAIL"),
232 reason="set YNAB_AGENT_LIVE_EMAIL=1 to send a real AgentMail message",
234def test_live_send_reaches_the_owner() -> None:
235 inbox = os.environ.get("YNAB_AGENT_INBOX", "alivehoney93@agentmail.to")
236 to = os.environ.get("YNAB_AGENT_OWNER", "matthew@mseeks.me")
237 out = MailClient.from_env().send(
238 OutboundEmail(
239 inbox_id=inbox,
240 to=(to,),
241 subject="[YNAB agent] live mail smoke",
242 text="This is the YNAB agent's AgentMail send path, working.",
243 )
244 )
245 assert out.message_id
246 assert out.thread_id
247 
248 
249@pytest.mark.skipif(
250 not os.environ.get("YNAB_AGENT_LIVE_EMAIL"),
251 reason="set YNAB_AGENT_LIVE_EMAIL=1 to send real AgentMail messages",
253def test_live_reply_threads_and_addresses_the_owner() -> None:
254 """The W7 fix's load-bearing claim, against the real server (SPEC §8).
255 
256 Reproduce the W6→W7 shape: open an alert thread (the agent's own outbound),
257 then reply on it addressed to the owner — the exact case commit a60e676
258 found broke (the reply's ``To`` looped back to the agent, so the owner never
259 saw it). Assert the reply lands on the *same* thread (AgentMail threaded it)
260 and its ``To`` is the owner (delivered, not looped back). These two server
261 behaviors are what the offline fake cannot prove.
262 """
263 from agentmail import AgentMail
264 
265 inbox = os.environ.get("YNAB_AGENT_INBOX", "alivehoney93@agentmail.to")
266 owner = os.environ.get("YNAB_AGENT_OWNER", "matthew@mseeks.me")
267 nonce = os.urandom(4).hex()
268 
269 mail = MailClient.from_env()
270 thread = mail.open_thread(
271 inbox_id=inbox,
272 to=[owner],
273 subject="[YNAB agent] live reply-threading smoke",
274 body="alert: Dining is over budget.",
275 txn_label=f"live-alert-{nonce}",
276 )
277 sent = mail.send_on_thread(
278 inbox_id=inbox,
279 thread_id=thread,
280 body="Cover the overspend? (live reply-threading test.)",
281 seq_label=f"live-cover-{nonce}",
282 to=[owner],
283 )
284 assert sent is True
285 
286 raw = AgentMail(api_key=os.environ["AGENTMAIL_API_KEY"])
287 last_id = raw.inboxes.threads.get(inbox, thread).last_message_id
288 reply = raw.inboxes.messages.get(inbox, last_id)
289 # Threaded on the server, and addressed to the owner (not looped back to the
290 # agent — the a60e676 symptom). ``in_reply_to`` confirms the threading
291 # headers the owner's mail client needs are set.
292 assert reply.thread_id == thread
293 assert owner in str(reply.to)
294 assert reply.in_reply_to is not None