main
py 622 lines 20.6 KB
Raw
1 """
2 Email handler — orchestrates poll, dispatch, and reply.
3
4 Requires agent context.
5 """
6
7 import asyncio
8 import base64
9 import json
10 import os
11 import uuid
12
13 from agent import Agent, AgentContext, AgentContextType, UserMessage
14 from helpers import guids, plugins, files, runtime
15 from helpers import message_queue as mq
16 from helpers import integration_commands
17 from helpers.persist_chat import save_tmp_chat
18 from helpers.print_style import PrintStyle
19 from helpers.errors import format_error
20 from initialize import initialize_agent
21
22 from plugins._email_integration.helpers import dispatcher as disp
23 from plugins._model_config.helpers import model_config
24 from plugins._email_integration.helpers.imap_client import (
25 InboundMessage,
26 connect_imap,
27 disconnect_imap,
28 fetch_new,
29 fetch_unread_since,
30 get_highest_uid,
31 connect_exchange,
32 fetch_unread_exchange,
33 )
34 from plugins._email_integration.helpers.smtp_client import SmtpConfig, send_reply
35
36
37 PLUGIN_NAME = "_email_integration"
38 DOWNLOAD_FOLDER = "usr/email/attachments"
39 STATE_FILE = "usr/email/state.json"
40
41
42 # ------------------------------------------------------------------
43 # UID state persistence
44 # ------------------------------------------------------------------
45
46 _state_lock = asyncio.Lock()
47
48 # Poll task registry — lives here (not in extension module) because
49 # extension modules are re-executed on each job_loop tick (cache disabled),
50 # which would reset module-level state and orphan running tasks.
51 _poll_tasks: dict[str, asyncio.Task] = {} # type: ignore[type-arg]
52
53 def _load_state() -> dict:
54 path = files.get_abs_path(STATE_FILE)
55 if os.path.isfile(path):
56 try:
57 return json.loads(files.read_file(path))
58 except Exception:
59 return {}
60 return {}
61
62
63 def _save_state(state: dict):
64 path = files.get_abs_path(STATE_FILE)
65 files.make_dirs(path)
66 files.write_file(path, json.dumps(state))
67
68
69 # ------------------------------------------------------------------
70 # Single handler poll (called from per-handler poll loop)
71 # ------------------------------------------------------------------
72
73 async def _poll_single_handler(handler_cfg: dict, state: dict):
74 name = handler_cfg.get("name", "default")
75 account_type = handler_cfg.get("account_type", "imap")
76 whitelist = handler_cfg.get("sender_whitelist") or []
77 last_uid = state.get(name, {}).get("last_uid", 0)
78 process_unread_days = int(handler_cfg.get("process_unread_days", 0))
79
80 if account_type == "exchange":
81 messages = await _fetch_exchange(handler_cfg, whitelist, process_unread_days)
82 if messages:
83 await _dispatch_all(handler_cfg, messages)
84 return
85
86 client = await connect_imap(
87 server=handler_cfg.get("imap_server", ""),
88 port=int(handler_cfg.get("imap_port", 993)),
89 username=handler_cfg.get("username", ""),
90 password=handler_cfg.get("password", ""),
91 )
92 try:
93 # First run: optionally process unread from last N days
94 if last_uid == 0:
95 if process_unread_days > 0:
96 messages, highest = await fetch_unread_since(
97 client, DOWNLOAD_FOLDER, process_unread_days, whitelist or None,
98 )
99 highest = highest or await get_highest_uid(client)
100 state[name] = {"last_uid": highest}
101 if messages:
102 PrintStyle.info(
103 f"Email ({name}): processing {len(messages)} unread"
104 f" from last {process_unread_days} days"
105 )
106 await _dispatch_all(handler_cfg, messages)
107 else:
108 PrintStyle.info(
109 f"Email ({name}): no unread in last {process_unread_days} days"
110 )
111 else:
112 highest = await get_highest_uid(client)
113 state[name] = {"last_uid": highest}
114 PrintStyle.info(f"Email ({name}): initialized, tracking from UID {highest}")
115 return
116
117 messages, new_uid = await fetch_new(
118 client, DOWNLOAD_FOLDER, last_uid, whitelist or None,
119 )
120
121 if new_uid > last_uid:
122 state[name] = {"last_uid": new_uid}
123
124 if messages:
125 PrintStyle.info(f"Email ({name}): {len(messages)} new messages")
126 await _dispatch_all(handler_cfg, messages)
127
128 finally:
129 await disconnect_imap(client)
130
131
132 async def _fetch_exchange(
133 cfg: dict, whitelist: list[str], since_days: int = 0,
134 ) -> list[InboundMessage]:
135 account = await connect_exchange(
136 server=cfg.get("imap_server", ""),
137 username=cfg.get("username", ""),
138 password=cfg.get("password", ""),
139 )
140 return await fetch_unread_exchange(
141 account, DOWNLOAD_FOLDER, whitelist or None, since_days=since_days,
142 )
143
144
145 async def _dispatch_all(handler_cfg: dict, messages: list[InboundMessage]):
146 own_address = (handler_cfg.get("username") or "").lower()
147
148 # Need an agent for dispatcher AI calls
149 # find existing dispatcher or create new background context
150 ctx = None
151 for c in AgentContext._contexts.values():
152 if isinstance(c, AgentContext) and c.name == "Email Dispatcher":
153 ctx = c
154 break
155
156 if not ctx:
157 agent_config = initialize_agent()
158 ctx = AgentContext(agent_config, name="Email Dispatcher",
159 type=AgentContextType.BACKGROUND)
160 agent = ctx.agent0
161
162 for msg in messages:
163 if own_address and _is_own_email(msg.sender, own_address):
164 PrintStyle.info(f"Email: skipping self-sent from {msg.sender}")
165 continue
166 try:
167 await _dispatch_message(agent, handler_cfg, msg)
168 except Exception as e:
169 PrintStyle.error(f"Email dispatch error: {format_error(e)}")
170
171
172 # ------------------------------------------------------------------
173 # Dispatch a single inbound message
174 # ------------------------------------------------------------------
175
176 async def _dispatch_message(agent: Agent, handler_cfg: dict, msg: InboundMessage):
177 handler_name = handler_cfg.get("name", "default")
178 thread_id = disp.extract_thread_id(msg.subject)
179
180 existing = _find_handler_chats(handler_name, msg.sender)
181
182 if await _handle_control_email(handler_cfg, msg, existing, thread_id):
183 return
184
185 # Fast path: thread ID in subject matches a known chat
186 if thread_id:
187 for chat in existing:
188 if chat["thread_id"] == thread_id:
189 await _route_to_chat(
190 agent, handler_cfg, msg, chat["context_id"],
191 )
192 return
193
194 # Dispatcher AI decides
195 decision = await _call_dispatcher(agent, handler_cfg, msg, existing)
196 reason = decision.reason or ""
197
198 if decision.action == "continue_chat" and decision.context_id:
199 ctx = AgentContext.get(decision.context_id)
200 if ctx:
201 await _route_to_chat(agent, handler_cfg, msg, decision.context_id)
202 return
203 PrintStyle.warning(
204 f"Dispatcher referenced unknown context {decision.context_id}, starting new chat"
205 )
206
207 await _start_new_chat(agent, handler_cfg, msg)
208
209
210 async def _call_model(
211 agent: Agent, handler_cfg: dict, system: str, prompt: str,
212 ):
213 if handler_cfg.get("dispatcher_model", "utility") == "chat":
214 from langchain_core.messages import SystemMessage, HumanMessage
215 messages = [SystemMessage(content=system), HumanMessage(content=prompt)]
216 response, _ = await agent.call_chat_model(messages)
217 return response
218 return await agent.call_utility_model(system=system, message=prompt)
219
220
221 async def _call_dispatcher(
222 agent: Agent,
223 handler_cfg: dict,
224 msg: InboundMessage,
225 existing_chats: list[disp.ChatSummary],
226 ) -> disp.DispatchDecision:
227 body_preview = disp.truncate_body(msg.body)
228 chats_text = disp.format_chats_list(existing_chats)
229
230 prompt = agent.read_prompt(
231 "fw.email.dispatcher_prompt.md",
232 sender=msg.sender,
233 subject=msg.subject,
234 body=body_preview,
235 chats=chats_text,
236 )
237
238 extra = handler_cfg.get("dispatcher_instructions", "")
239 if extra:
240 prompt += agent.read_prompt(
241 "fw.email.dispatcher_extra.md", instructions=extra,
242 )
243
244 system = agent.read_prompt("fw.email.dispatcher_system.md")
245
246 try:
247 response = await _call_model(agent, handler_cfg, system, prompt)
248 return disp.parse_dispatcher_response(str(response))
249
250 except Exception as e:
251 PrintStyle.error(f"Dispatcher error: {format_error(e)}")
252 return disp.DispatchDecision(action="new_chat", reason="dispatcher error")
253
254
255 # ------------------------------------------------------------------
256 # Chat creation and routing
257 # ------------------------------------------------------------------
258
259 async def _start_new_chat(agent: Agent, handler_cfg: dict, msg: InboundMessage):
260 from helpers import projects
261
262 handler_name = handler_cfg.get("name", "default")
263 thread_id = guids.generate_id()
264
265 config = initialize_agent()
266 context = AgentContext(config, name=f"Email: {msg.subject[:50]}")
267
268 context.data[disp.CTX_EMAIL_HANDLER] = handler_name
269 context.data[disp.CTX_EMAIL_SENDER] = msg.sender
270 context.data[disp.CTX_EMAIL_THREAD_ID] = thread_id
271 context.data[disp.CTX_EMAIL_SUBJECT] = msg.subject
272 context.data[disp.CTX_EMAIL_LAST_BODY] = msg.body
273 context.data[disp.CTX_EMAIL_MESSAGE_ID] = msg.message_id
274
275 refs_list = []
276 if msg.references:
277 for r in msg.references.split():
278 if r not in refs_list:
279 refs_list.append(r)
280 if msg.message_id and msg.message_id not in refs_list:
281 refs_list.append(msg.message_id)
282
283 context.data[disp.CTX_EMAIL_REFERENCES] = " ".join(refs_list)
284
285 project = handler_cfg.get("project", "")
286 if project:
287 projects.activate_project(context.id, project)
288
289 _apply_handler_model_preset(context, handler_cfg)
290 save_tmp_chat(context)
291
292 user_msg = _build_user_message(agent, msg, handler_cfg)
293 system_ctx = agent.read_prompt("fw.email.system_context.md")
294
295 msg_id = str(uuid.uuid4())
296 mq.log_user_message(context, user_msg, msg.attachments or [], message_id=msg_id, source=" (email)")
297 context.communicate(UserMessage(
298 message=user_msg,
299 system_message=[system_ctx],
300 attachments=msg.attachments,
301 id=msg_id,
302 ))
303
304 PrintStyle.success(f"Email: new chat {context.id} for '{msg.subject}' from {msg.sender}")
305
306
307 async def _route_to_chat(
308 agent: Agent,
309 handler_cfg: dict,
310 msg: InboundMessage,
311 context_id: str,
312 ):
313 context = AgentContext.get(context_id)
314 if not context:
315 return
316
317 context.data[disp.CTX_EMAIL_MESSAGE_ID] = msg.message_id
318 context.data[disp.CTX_EMAIL_LAST_BODY] = msg.body
319 if not context.get_data("chat_model_override"):
320 _apply_handler_model_preset(context, handler_cfg)
321
322 refs = context.data.get(disp.CTX_EMAIL_REFERENCES, "")
323 refs_list = refs.split() if refs else []
324
325 if msg.references:
326 for r in msg.references.split():
327 if r not in refs_list:
328 refs_list.append(r)
329
330 if msg.message_id and msg.message_id not in refs_list:
331 refs_list.append(msg.message_id)
332
333 context.data[disp.CTX_EMAIL_REFERENCES] = " ".join(refs_list)
334
335 user_msg = _build_user_message(agent, msg, handler_cfg)
336 msg_id = str(uuid.uuid4())
337 mq.log_user_message(context, user_msg, msg.attachments or [], message_id=msg_id, source=" (email)")
338 context.communicate(UserMessage(
339 message=user_msg,
340 attachments=msg.attachments,
341 id=msg_id,
342 ))
343
344 save_tmp_chat(context)
345 PrintStyle.info(f"Email: continuing chat {context_id}")
346
347
348 async def _handle_control_email(
349 handler_cfg: dict,
350 msg: InboundMessage,
351 existing_chats: list[disp.ChatSummary],
352 thread_id: str,
353 ) -> bool:
354 parsed = integration_commands.parse_command(msg.body)
355 if not parsed:
356 return False
357
358 target_context_id = ""
359 if thread_id:
360 for chat in existing_chats:
361 if chat["thread_id"] == thread_id:
362 target_context_id = chat["context_id"]
363 break
364
365 if not target_context_id:
366 if len(existing_chats) == 1:
367 target_context_id = existing_chats[0]["context_id"]
368 elif len(existing_chats) > 1:
369 await _send_control_email_reply(
370 handler_cfg,
371 msg,
372 "Multiple Agent Zero email chats match this sender. Reply inside the thread you want to control.",
373 thread_id=thread_id,
374 )
375 return True
376 else:
377 await _send_control_email_reply(
378 handler_cfg,
379 msg,
380 "No matching Agent Zero email chat was found. Reply inside an existing Agent Zero email thread to use /project, /config, or /send.",
381 thread_id=thread_id,
382 )
383 return True
384
385 context = AgentContext.get(target_context_id)
386 if not context:
387 await _send_control_email_reply(
388 handler_cfg,
389 msg,
390 "The matching Agent Zero email chat is no longer available. Send a normal email to start a fresh thread.",
391 thread_id=thread_id,
392 )
393 return True
394
395 response = integration_commands.try_handle_command(context, msg.body)
396 if response is None:
397 return False
398
399 await _send_control_email_reply(
400 handler_cfg,
401 msg,
402 response,
403 thread_id=context.data.get(disp.CTX_EMAIL_THREAD_ID, "") or thread_id,
404 )
405 return True
406
407
408 def _apply_handler_model_preset(context: AgentContext, handler_cfg: dict) -> None:
409 preset_name = str(handler_cfg.get("chat_model_preset", "") or "").strip()
410 if not preset_name:
411 return
412 if not model_config.is_chat_override_allowed(context.agent0):
413 PrintStyle.warning(
414 f"Email ({handler_cfg.get('name', 'default')}): chat override is disabled,"
415 f" cannot apply preset '{preset_name}'"
416 )
417 return
418 if not model_config.get_preset_by_name(preset_name):
419 PrintStyle.warning(
420 f"Email ({handler_cfg.get('name', 'default')}): preset '{preset_name}' was not found"
421 )
422 return
423 context.set_data("chat_model_override", {"preset_name": preset_name})
424
425
426 async def _send_control_email_reply(
427 handler_cfg: dict,
428 msg: InboundMessage,
429 body: str,
430 *,
431 thread_id: str = "",
432 ) -> str | None:
433 smtp_cfg = SmtpConfig(
434 server=handler_cfg.get("smtp_server", handler_cfg.get("imap_server", "")),
435 port=int(handler_cfg.get("smtp_port", 587)),
436 username=handler_cfg.get("username", ""),
437 password=handler_cfg.get("password", ""),
438 )
439
440 subject = _build_control_reply_subject(msg.subject, thread_id)
441 references = _merge_references(msg.references, msg.message_id)
442
443 return await send_reply(
444 config=smtp_cfg,
445 to=msg.sender,
446 subject=subject,
447 body=body,
448 in_reply_to=msg.message_id,
449 references=references,
450 attachments=None,
451 )
452
453
454 def _build_control_reply_subject(subject: str, thread_id: str) -> str:
455 if thread_id:
456 return disp.build_reply_subject(subject, thread_id)
457 cleaned = subject.strip()
458 if not cleaned.lower().startswith("re:"):
459 cleaned = f"Re: {cleaned}"
460 return cleaned
461
462
463 def _merge_references(existing: str, message_id: str) -> str:
464 refs = []
465 for ref in (existing or "").split():
466 if ref and ref not in refs:
467 refs.append(ref)
468 if message_id and message_id not in refs:
469 refs.append(message_id)
470 return " ".join(refs)
471
472
473 # ------------------------------------------------------------------
474 # Chat discovery
475 # ------------------------------------------------------------------
476
477 HISTORY_PREVIEW_MAX_CHARS: int = 500
478
479
480 def _find_handler_chats(handler_name: str, sender: str) -> list[disp.ChatSummary]:
481 results = []
482 for ctx_id, ctx in AgentContext._contexts.items():
483 if not isinstance(ctx, AgentContext):
484 continue
485 data = ctx.data
486 if data.get(disp.CTX_EMAIL_HANDLER) != handler_name:
487 continue
488 if data.get(disp.CTX_EMAIL_SENDER, "").lower() != sender.lower():
489 continue
490 summary = disp.build_chat_summary(ctx_id, data)
491 summary["history_preview"] = _get_history_preview(ctx)
492 results.append(summary)
493
494 results.sort(key=lambda c: c["context_id"], reverse=True)
495 return results[:20]
496
497
498 def _get_history_preview(ctx: AgentContext) -> str:
499 try:
500 history = ctx.agent0.history
501 text = history.output_text(human_label="user", ai_label="agent")
502 if not text:
503 return "(empty)"
504 if len(text) > HISTORY_PREVIEW_MAX_CHARS:
505 return "..." + text[-HISTORY_PREVIEW_MAX_CHARS:]
506 return text
507 except Exception:
508 return "(unavailable)"
509
510
511 # ------------------------------------------------------------------
512 # Sender helpers
513 # ------------------------------------------------------------------
514
515 def _is_own_email(sender: str, own_address: str) -> bool:
516 sender_lower = sender.lower()
517 if "<" in sender_lower:
518 start = sender_lower.index("<") + 1
519 end = sender_lower.index(">", start)
520 return sender_lower[start:end].strip() == own_address
521 return sender_lower.strip() == own_address
522
523
524 # ------------------------------------------------------------------
525 # Message builders
526 # ------------------------------------------------------------------
527
528 def _build_user_message(agent: Agent, msg: InboundMessage, handler_cfg: dict) -> str:
529 recipient = handler_cfg.get("username", "")
530 return agent.read_prompt(
531 "fw.email.user_message.md",
532 sender=msg.sender,
533 recipient=recipient,
534 subject=msg.subject,
535 body=msg.body,
536 )
537
538
539 # ------------------------------------------------------------------
540 # Reply sending (called from process_chain_end extension)
541 # ------------------------------------------------------------------
542
543 async def send_email_reply(
544 context: AgentContext,
545 response_text: str,
546 attachments: list[str] | None = None,
547 ) -> str | None:
548 handler_name = context.data.get(disp.CTX_EMAIL_HANDLER)
549 if not handler_name:
550 return "No email handler configured"
551
552 cfg = _get_handler_config(handler_name)
553 if not cfg:
554 return f"Handler config not found for '{handler_name}'"
555
556 sender = context.data.get(disp.CTX_EMAIL_SENDER, "")
557 original_subject = context.data.get(disp.CTX_EMAIL_SUBJECT, "")
558 thread_id = context.data.get(disp.CTX_EMAIL_THREAD_ID, "")
559 original_msg_id = context.data.get(disp.CTX_EMAIL_MESSAGE_ID, "")
560 references = context.data.get(disp.CTX_EMAIL_REFERENCES, "")
561
562 subject = disp.build_reply_subject(original_subject, thread_id)
563
564 smtp_cfg = SmtpConfig(
565 server=cfg.get("smtp_server", cfg.get("imap_server", "")),
566 port=int(cfg.get("smtp_port", 587)),
567 username=cfg.get("username", ""),
568 password=cfg.get("password", ""),
569 )
570
571 # Read attachment files via RFC (they live in the execution runtime)
572 attachment_data = await _read_attachments_via_rfc(attachments)
573
574 last_body = context.data.get(disp.CTX_EMAIL_LAST_BODY, "").strip()
575 if last_body:
576 quoted = "\n> " + "\n> ".join(last_body.splitlines())
577 response_text = f"{response_text}\n\nOn previous message:\n{quoted}"
578
579 return await send_reply(
580 config=smtp_cfg,
581 to=sender,
582 subject=subject,
583 body=response_text,
584 in_reply_to=original_msg_id,
585 references=references,
586 attachments=attachment_data or None,
587 )
588
589
590 # ------------------------------------------------------------------
591 # Attachment reading (via RFC into execution runtime)
592 # ------------------------------------------------------------------
593
594 async def _read_attachments_via_rfc(
595 paths: list[str] | None,
596 ) -> list[tuple[str, bytes]]:
597 if not paths:
598 return []
599
600 from plugins._email_integration.helpers.attachment_reader import read_attachment
601
602 results: list[tuple[str, bytes]] = []
603 for path in paths:
604 data = await runtime.call_development_function(read_attachment, path)
605 if data["error"]:
606 PrintStyle.error(f"Email attachment: {data['error']}")
607 continue
608 results.append((data["name"], base64.b64decode(data["content_b64"])))
609 return results
610
611
612 # ------------------------------------------------------------------
613 # Config lookup
614 # ------------------------------------------------------------------
615
616 def _get_handler_config(handler_name: str) -> dict | None:
617 config = plugins.get_plugin_config(PLUGIN_NAME) or {}
618 handlers = config.get("handlers", [])
619 for h in handlers:
620 if h.get("name") == handler_name:
621 return h
622 return None