Skip to main content

objectiveai_cli/db/logs/
write.rs

1//! Per-table flat INSERT and UPDATE helpers — each issuing one CTE
2//! that fires the streaming-content / request-blob write AND the
3//! `objectiveai.messages` / `objectiveai.messages_queue` bookkeeping in a single
4//! postgres round-trip.
5//!
6//! Insert path (streaming-content INSERT or request-blob INSERT):
7//! ```sql
8//! WITH data_ins AS (
9//!     INSERT INTO logs.<table> (…) VALUES (…) RETURNING response_id
10//! )
11//! INSERT INTO objectiveai.messages (response_id, "table", row_index,
12//!                            row_sub_index, "index",
13//!                            agent_instance_hierarchy, "timestamp")
14//! SELECT $resp, $msg_table, $row_idx, $row_sub_idx,
15//!        nextval('objectiveai.messages_index_seq'),
16//!        $hier, $ts
17//! FROM data_ins;
18//! ```
19//!
20//! Update path (streaming-content UPDATE):
21//! ```sql
22//! WITH
23//!     data_upd AS (
24//!         UPDATE logs.<table> SET … WHERE … RETURNING response_id
25//!     ),
26//!     msg AS (
27//!         SELECT "index" AS msg_index FROM objectiveai.messages
28//!         WHERE response_id = $resp AND "table" = $msg_table
29//!           AND row_index IS NOT DISTINCT FROM $row_idx
30//!           AND row_sub_index IS NOT DISTINCT FROM $row_sub_idx
31//!     )
32//! UPDATE objectiveai.messages_queue
33//! SET read_index = msg.msg_index - 1
34//! FROM msg, data_upd
35//! WHERE spawned_agent_instance_hierarchy = $hier
36//!   AND read_index >= msg.msg_index;
37//! ```
38//!
39//! Response-blob writes (the three `_responses` tables) DON'T touch
40//! `objectiveai.messages` — they're not events, just the latest snapshot.
41
42use objectiveai_sdk::agent::completions::message::{File, ImageUrl, InputAudio, VideoUrl};
43use serde::Serialize;
44
45use crate::db::{Error, Pool};
46
47use super::row::{MessageTable, RowValue};
48use super::shadow::WriteOp;
49
50/// Overwrite an AIH's most-recent agent-completion `total_tokens`
51/// (last-write-wins snapshot — not a running sum).
52pub async fn update_agent_token_usage(
53    pool: &Pool,
54    agent_instance_hierarchy: &str,
55    total_tokens: i64,
56) -> Result<(), Error> {
57    sqlx::query(
58        "INSERT INTO objectiveai.agent_token_usage \
59            (agent_instance_hierarchy, total_tokens) \
60         VALUES ($1, $2) \
61         ON CONFLICT (agent_instance_hierarchy) \
62         DO UPDATE SET total_tokens = EXCLUDED.total_tokens",
63    )
64    .bind(agent_instance_hierarchy)
65    .bind(total_tokens)
66    .execute(&**pool)
67    .await?;
68    Ok(())
69}
70
71/// Dispatch SQL for `value` per `op`. `Skip` is a no-op.
72pub async fn write_value<'a>(
73    pool: &Pool,
74    op: WriteOp,
75    value: &RowValue<'a>,
76    timestamp: i64,
77) -> Result<(), Error> {
78    match op {
79        WriteOp::Skip => Ok(()),
80        WriteOp::Insert => insert_value(pool, value, timestamp).await,
81        WriteOp::Update => update_value(pool, value).await,
82    }
83}
84
85async fn insert_value<'a>(
86    pool: &Pool,
87    value: &RowValue<'a>,
88    timestamp: i64,
89) -> Result<(), Error> {
90    // MessageQueueContent: branch early, its helper resolves
91    // both the kind (and thus the objectiveai.message_table enum value)
92    // and the parent message_queue.id from `message_queue_contents`
93    // via SQL CASE/subquery. No call into `value.message_table()`
94    // — that returns `None` for this variant by design.
95    if let RowValue::MessageQueueContent {
96        response_id,
97        agent_instance_hierarchy,
98        message_queue_content_id,
99    } = *value
100    {
101        return insert_message_queue_content_with_msg(
102            pool,
103            response_id,
104            agent_instance_hierarchy,
105            message_queue_content_id,
106            timestamp,
107        )
108        .await;
109    }
110
111    // ToolResponse: the head row is written to `objectiveai.tool_response`
112    // purely as the `tool_call_id` lookup for its content rows (JOINed at
113    // read time by `read_all`). It emits NO `objectiveai.messages` event,
114    // so branch early — `value.message_table()` returns `None` for this
115    // variant and would panic the `.expect()` below.
116    if let RowValue::ToolResponse { response_id, index, tool_call_id, .. } = *value {
117        sqlx::query(
118            "INSERT INTO objectiveai.tool_response (response_id, \"index\", tool_call_id) \
119             VALUES ($1, $2, $3)",
120        )
121        .bind(response_id)
122        .bind(index as i64)
123        .bind(tool_call_id)
124        .execute(&**pool)
125        .await?;
126        return Ok(());
127    }
128
129    let mt = value.message_table();
130    let hier = value.agent_instance_hierarchy();
131    let row_index = value.row_index();
132    let row_sub_index = value.row_sub_index();
133    let response_id = value.response_id();
134
135    match *value {
136        RowValue::MessageQueueContent { .. } => unreachable!(
137            "MessageQueueContent handled by early-return branch above"
138        ),
139        RowValue::ToolResponse { .. } => unreachable!(
140            "ToolResponse handled by early-return branch above"
141        ),
142        RowValue::AssistantResponseRefusal { text, .. } => {
143            sqlx::query(
144                "WITH data_ins AS (\
145                    INSERT INTO objectiveai.assistant_response_refusal (response_id, \"index\", text) \
146                    VALUES ($1, $2, $3) RETURNING response_id\
147                 )\
148                 INSERT INTO objectiveai.messages \
149                    (response_id, \"table\", row_index, row_sub_index, \
150                     agent_instance_hierarchy, \"timestamp\") \
151                 SELECT $1, $4, $5, $6, $7, $8 FROM data_ins",
152            )
153            .bind(response_id)
154            .bind(row_index)
155            .bind(text)
156            .bind(mt)
157            .bind(row_index)
158            .bind(row_sub_index)
159            .bind(hier)
160            .bind(timestamp)
161            .execute(&**pool)
162            .await?;
163        }
164        RowValue::AssistantResponseReasoning { text, .. } => {
165            sqlx::query(
166                "WITH data_ins AS (\
167                    INSERT INTO objectiveai.assistant_response_reasoning (response_id, \"index\", text) \
168                    VALUES ($1, $2, $3) RETURNING response_id\
169                 )\
170                 INSERT INTO objectiveai.messages \
171                    (response_id, \"table\", row_index, row_sub_index, \
172                     agent_instance_hierarchy, \"timestamp\") \
173                 SELECT $1, $4, $5, $6, $7, $8 FROM data_ins",
174            )
175            .bind(response_id)
176            .bind(row_index)
177            .bind(text)
178            .bind(mt)
179            .bind(row_index)
180            .bind(row_sub_index)
181            .bind(hier)
182            .bind(timestamp)
183            .execute(&**pool)
184            .await?;
185        }
186        RowValue::AssistantResponseToolCalls {
187            tool_call_index, tool_call_id, function_name, arguments, ..
188        } => {
189            sqlx::query(
190                "WITH data_ins AS (\
191                    INSERT INTO objectiveai.assistant_response_tool_calls \
192                        (response_id, \"index\", tool_call_index, tool_call_id, function_name, arguments) \
193                    VALUES ($1, $2, $3, $4, $5, $6) RETURNING response_id\
194                 )\
195                 INSERT INTO objectiveai.messages \
196                    (response_id, \"table\", row_index, row_sub_index, \
197                     agent_instance_hierarchy, \"timestamp\") \
198                 SELECT $1, $7, $8, $9, $10, $11 FROM data_ins",
199            )
200            .bind(response_id)
201            .bind(row_index)
202            .bind(tool_call_index as i64)
203            .bind(tool_call_id)
204            .bind(function_name)
205            .bind(arguments)
206            .bind(mt)
207            .bind(row_index)
208            .bind(row_sub_index)
209            .bind(hier)
210            .bind(timestamp)
211            .execute(&**pool)
212            .await?;
213        }
214        RowValue::AssistantResponseContentText { text, .. } => {
215            insert_text_part_with_msg(pool, "objectiveai.assistant_response_content_text", value, text, timestamp).await?;
216        }
217        RowValue::ToolResponseContentText { text, .. } => {
218            insert_text_part_with_msg(pool, "objectiveai.tool_response_content_text", value, text, timestamp).await?;
219        }
220        RowValue::AssistantResponseContentImage { image_url, .. } => {
221            insert_image_part_with_msg(pool, "objectiveai.assistant_response_content_image", value, image_url, timestamp).await?;
222        }
223        RowValue::ToolResponseContentImage { image_url, .. } => {
224            insert_image_part_with_msg(pool, "objectiveai.tool_response_content_image", value, image_url, timestamp).await?;
225        }
226        RowValue::AssistantResponseContentAudio { input_audio, .. } => {
227            insert_audio_part_with_msg(pool, "objectiveai.assistant_response_content_audio", value, input_audio, timestamp).await?;
228        }
229        RowValue::ToolResponseContentAudio { input_audio, .. } => {
230            insert_audio_part_with_msg(pool, "objectiveai.tool_response_content_audio", value, input_audio, timestamp).await?;
231        }
232        RowValue::AssistantResponseContentVideo { video_url, .. } => {
233            insert_video_part_with_msg(pool, "objectiveai.assistant_response_content_video", value, video_url, timestamp).await?;
234        }
235        RowValue::ToolResponseContentVideo { video_url, .. } => {
236            insert_video_part_with_msg(pool, "objectiveai.tool_response_content_video", value, video_url, timestamp).await?;
237        }
238        RowValue::AssistantResponseContentFile { file, .. } => {
239            insert_file_part_with_msg(pool, "objectiveai.assistant_response_content_file", value, file, timestamp).await?;
240        }
241        RowValue::ToolResponseContentFile { file, .. } => {
242            insert_file_part_with_msg(pool, "objectiveai.tool_response_content_file", value, file, timestamp).await?;
243        }
244    }
245    Ok(())
246}
247
248async fn update_value<'a>(pool: &Pool, value: &RowValue<'a>) -> Result<(), Error> {
249    // MessageQueueContent has no updatable body — the shadow's
250    // body_eq returns true for any matching key, so this branch
251    // is unreachable in practice. Short-circuit defensively.
252    if matches!(value, RowValue::MessageQueueContent { .. }) {
253        return Ok(());
254    }
255
256    // ToolResponse: its head row has no `messages` event (it's a
257    // `tool_call_id` lookup only), so there is nothing to downgrade.
258    // `tool_call_id` is immutable in practice (the shadow's `body_eq`
259    // returns true, so an Update is never dispatched), making this
260    // branch effectively unreachable; do a bare UPDATE defensively.
261    // Branch early — `value.message_table()` returns `None` for this
262    // variant and would panic the `.expect()` below.
263    if let RowValue::ToolResponse { response_id, index, tool_call_id, .. } = *value {
264        sqlx::query(
265            "UPDATE objectiveai.tool_response SET tool_call_id = $1 \
266             WHERE response_id = $2 AND \"index\" = $3",
267        )
268        .bind(tool_call_id)
269        .bind(response_id)
270        .bind(index as i64)
271        .execute(&**pool)
272        .await?;
273        return Ok(());
274    }
275
276    let mt = value.message_table();
277    let hier = value.agent_instance_hierarchy();
278    let row_index = value.row_index();
279    let row_sub_index = value.row_sub_index();
280    let response_id = value.response_id();
281
282    match *value {
283        RowValue::MessageQueueContent { .. } => unreachable!(
284            "MessageQueueContent handled by short-circuit above"
285        ),
286        RowValue::ToolResponse { .. } => unreachable!(
287            "ToolResponse handled by short-circuit above"
288        ),
289        RowValue::AssistantResponseRefusal { text, .. } => {
290            run_update_with_downgrade(
291                pool,
292                "UPDATE objectiveai.assistant_response_refusal SET text = $A \
293                 WHERE response_id = $RESP AND \"index\" = $RI",
294                response_id, row_index, row_sub_index, mt, hier,
295                &[("A", BindVal::Str(text))],
296                &[BindIdx::Resp, BindIdx::Ri],
297            ).await?;
298        }
299        RowValue::AssistantResponseReasoning { text, .. } => {
300            run_update_with_downgrade(
301                pool,
302                "UPDATE objectiveai.assistant_response_reasoning SET text = $A \
303                 WHERE response_id = $RESP AND \"index\" = $RI",
304                response_id, row_index, row_sub_index, mt, hier,
305                &[("A", BindVal::Str(text))],
306                &[BindIdx::Resp, BindIdx::Ri],
307            ).await?;
308        }
309        RowValue::AssistantResponseToolCalls { tool_call_index, tool_call_id, function_name, arguments, .. } => {
310            run_update_with_downgrade(
311                pool,
312                "UPDATE objectiveai.assistant_response_tool_calls SET tool_call_id = $A, function_name = $B, arguments = $C \
313                 WHERE response_id = $RESP AND \"index\" = $RI AND tool_call_index = $RSI",
314                response_id, row_index, row_sub_index, mt, hier,
315                &[("A", BindVal::Str(tool_call_id)), ("B", BindVal::Str(function_name)), ("C", BindVal::Str(arguments))],
316                &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
317            ).await?;
318            let _ = tool_call_index;
319        }
320        RowValue::AssistantResponseContentText { text, .. }
321        | RowValue::ToolResponseContentText { text, .. } => {
322            let table = match *value {
323                RowValue::AssistantResponseContentText { .. } => "objectiveai.assistant_response_content_text",
324                _ => "objectiveai.tool_response_content_text",
325            };
326            let sql = format!(
327                "UPDATE {table} SET text = $A \
328                 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
329            );
330            run_update_with_downgrade(
331                pool, &sql,
332                response_id, row_index, row_sub_index, mt, hier,
333                &[("A", BindVal::Str(text))],
334                &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
335            ).await?;
336        }
337        RowValue::AssistantResponseContentImage { image_url, .. }
338        | RowValue::ToolResponseContentImage { image_url, .. } => {
339            let table = match *value {
340                RowValue::AssistantResponseContentImage { .. } => "objectiveai.assistant_response_content_image",
341                _ => "objectiveai.tool_response_content_image",
342            };
343            let detail = image_url.detail.as_ref().and_then(|d| serde_json::to_string(d).ok());
344            let sql = format!(
345                "UPDATE {table} SET url = $A, detail = $B \
346                 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
347            );
348            run_update_with_downgrade(
349                pool, &sql,
350                response_id, row_index, row_sub_index, mt, hier,
351                &[("A", BindVal::Str(image_url.url.as_str())), ("B", BindVal::OptString(detail))],
352                &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
353            ).await?;
354        }
355        RowValue::AssistantResponseContentAudio { input_audio, .. }
356        | RowValue::ToolResponseContentAudio { input_audio, .. } => {
357            let table = match *value {
358                RowValue::AssistantResponseContentAudio { .. } => "objectiveai.assistant_response_content_audio",
359                _ => "objectiveai.tool_response_content_audio",
360            };
361            let sql = format!(
362                "UPDATE {table} SET data = $A, format = $B \
363                 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
364            );
365            run_update_with_downgrade(
366                pool, &sql,
367                response_id, row_index, row_sub_index, mt, hier,
368                &[
369                    ("A", BindVal::Str(input_audio.data.as_str())),
370                    ("B", BindVal::Str(input_audio.format.as_str())),
371                ],
372                &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
373            ).await?;
374        }
375        RowValue::AssistantResponseContentVideo { video_url, .. }
376        | RowValue::ToolResponseContentVideo { video_url, .. } => {
377            let table = match *value {
378                RowValue::AssistantResponseContentVideo { .. } => "objectiveai.assistant_response_content_video",
379                _ => "objectiveai.tool_response_content_video",
380            };
381            let sql = format!(
382                "UPDATE {table} SET url = $A \
383                 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
384            );
385            run_update_with_downgrade(
386                pool, &sql,
387                response_id, row_index, row_sub_index, mt, hier,
388                &[("A", BindVal::Str(video_url.url.as_str()))],
389                &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
390            ).await?;
391        }
392        RowValue::AssistantResponseContentFile { file, .. }
393        | RowValue::ToolResponseContentFile { file, .. } => {
394            let table = match *value {
395                RowValue::AssistantResponseContentFile { .. } => "objectiveai.assistant_response_content_file",
396                _ => "objectiveai.tool_response_content_file",
397            };
398            let sql = format!(
399                "UPDATE {table} SET file_data = $A, file_id = $B, filename = $C, file_url = $D \
400                 WHERE response_id = $RESP AND \"index\" = $RI AND part_index = $RSI"
401            );
402            run_update_with_downgrade(
403                pool, &sql,
404                response_id, row_index, row_sub_index, mt, hier,
405                &[
406                    ("A", BindVal::OptStr(file.file_data.as_deref())),
407                    ("B", BindVal::OptStr(file.file_id.as_deref())),
408                    ("C", BindVal::OptStr(file.filename.as_deref())),
409                    ("D", BindVal::OptStr(file.file_url.as_deref())),
410                ],
411                &[BindIdx::Resp, BindIdx::Ri, BindIdx::Rsi],
412            ).await?;
413        }
414    }
415    Ok(())
416}
417
418// ---- INSERT helpers for content parts (shared CTE shape) -------------
419
420async fn insert_text_part_with_msg<'a>(
421    pool: &Pool,
422    table: &str,
423    value: &RowValue<'a>,
424    text: &str,
425    timestamp: i64,
426) -> Result<(), Error> {
427    let sql = format!(
428        "WITH data_ins AS (\
429            INSERT INTO {table} (response_id, \"index\", part_index, text) \
430            VALUES ($1, $2, $3, $4) RETURNING response_id\
431         )\
432         INSERT INTO objectiveai.messages \
433            (response_id, \"table\", row_index, row_sub_index, \
434             agent_instance_hierarchy, \"timestamp\") \
435         SELECT $1, $5, $6, $7, $8, $9 FROM data_ins"
436    );
437    sqlx::query(&sql)
438        .bind(value.response_id())
439        .bind(value.row_index())
440        .bind(value.row_sub_index())
441        .bind(text)
442        .bind(value.message_table())
443        .bind(value.row_index())
444        .bind(value.row_sub_index())
445        .bind(value.agent_instance_hierarchy())
446        .bind(timestamp)
447        .execute(&**pool)
448        .await?;
449    Ok(())
450}
451
452async fn insert_image_part_with_msg<'a>(
453    pool: &Pool,
454    table: &str,
455    value: &RowValue<'a>,
456    image: &ImageUrl,
457    timestamp: i64,
458) -> Result<(), Error> {
459    let detail = image.detail.as_ref().and_then(|d| serde_json::to_string(d).ok());
460    let sql = format!(
461        "WITH data_ins AS (\
462            INSERT INTO {table} (response_id, \"index\", part_index, url, detail) \
463            VALUES ($1, $2, $3, $4, $5) RETURNING response_id\
464         )\
465         INSERT INTO objectiveai.messages \
466            (response_id, \"table\", row_index, row_sub_index, \
467             agent_instance_hierarchy, \"timestamp\") \
468         SELECT $1, $6, $7, $8, $9, $10 FROM data_ins"
469    );
470    sqlx::query(&sql)
471        .bind(value.response_id())
472        .bind(value.row_index())
473        .bind(value.row_sub_index())
474        .bind(image.url.as_str())
475        .bind(detail)
476        .bind(value.message_table())
477        .bind(value.row_index())
478        .bind(value.row_sub_index())
479        .bind(value.agent_instance_hierarchy())
480        .bind(timestamp)
481        .execute(&**pool)
482        .await?;
483    Ok(())
484}
485
486async fn insert_audio_part_with_msg<'a>(
487    pool: &Pool,
488    table: &str,
489    value: &RowValue<'a>,
490    audio: &InputAudio,
491    timestamp: i64,
492) -> Result<(), Error> {
493    let sql = format!(
494        "WITH data_ins AS (\
495            INSERT INTO {table} (response_id, \"index\", part_index, data, format) \
496            VALUES ($1, $2, $3, $4, $5) RETURNING response_id\
497         )\
498         INSERT INTO objectiveai.messages \
499            (response_id, \"table\", row_index, row_sub_index, \
500             agent_instance_hierarchy, \"timestamp\") \
501         SELECT $1, $6, $7, $8, $9, $10 FROM data_ins"
502    );
503    sqlx::query(&sql)
504        .bind(value.response_id())
505        .bind(value.row_index())
506        .bind(value.row_sub_index())
507        .bind(audio.data.as_str())
508        .bind(audio.format.as_str())
509        .bind(value.message_table())
510        .bind(value.row_index())
511        .bind(value.row_sub_index())
512        .bind(value.agent_instance_hierarchy())
513        .bind(timestamp)
514        .execute(&**pool)
515        .await?;
516    Ok(())
517}
518
519/// Consumption-flip + log emit for a single
520/// `message_queue_contents.id`. One SQL statement:
521///
522/// 1. `content` CTE looks up the content row to get its `kind`
523///    and parent `message_queue_id`.
524/// 2. `flip` CTE flips `message_queue.active = FALSE` for the
525///    parent (no-op if already false via the `AND active = TRUE`
526///    guard, so repeat content_ids sharing one parent fire the
527///    flip exactly once).
528/// 3. INSERT a `objectiveai.messages` row with `"table"` chosen by SQL
529///    CASE off the content's kind (`message_queue_text` / `_image`
530///    / `_audio` / `_video` / `_file`), `row_index = content_id`,
531///    no sub-index.
532async fn insert_message_queue_content_with_msg(
533    pool: &Pool,
534    response_id: &str,
535    agent_instance_hierarchy: &str,
536    message_queue_content_id: i64,
537    timestamp: i64,
538) -> Result<(), Error> {
539    sqlx::query(
540        "WITH content AS (\
541             SELECT id, kind, message_queue_id \
542             FROM objectiveai.message_queue_contents \
543             WHERE id = $1 \
544         ), \
545         flip AS (\
546             UPDATE objectiveai.message_queue \
547             SET active = FALSE \
548             WHERE id = (SELECT message_queue_id FROM content) \
549               AND active = TRUE \
550             RETURNING id \
551         ) \
552         INSERT INTO objectiveai.messages \
553             (response_id, \"table\", row_index, row_sub_index, \
554              agent_instance_hierarchy, \"timestamp\") \
555         SELECT $2, \
556                CASE (SELECT kind FROM content) \
557                    WHEN 'text'  THEN 'message_queue_text'::objectiveai.message_table \
558                    WHEN 'image' THEN 'message_queue_image'::objectiveai.message_table \
559                    WHEN 'audio' THEN 'message_queue_audio'::objectiveai.message_table \
560                    WHEN 'video' THEN 'message_queue_video'::objectiveai.message_table \
561                    WHEN 'file'  THEN 'message_queue_file'::objectiveai.message_table \
562                END, \
563                $1, NULL, $3, $4 \
564         FROM content",
565    )
566    .bind(message_queue_content_id)
567    .bind(response_id)
568    .bind(agent_instance_hierarchy)
569    .bind(timestamp)
570    .execute(&**pool)
571    .await?;
572    Ok(())
573}
574
575async fn insert_video_part_with_msg<'a>(
576    pool: &Pool,
577    table: &str,
578    value: &RowValue<'a>,
579    video: &VideoUrl,
580    timestamp: i64,
581) -> Result<(), Error> {
582    let sql = format!(
583        "WITH data_ins AS (\
584            INSERT INTO {table} (response_id, \"index\", part_index, url) \
585            VALUES ($1, $2, $3, $4) RETURNING response_id\
586         )\
587         INSERT INTO objectiveai.messages \
588            (response_id, \"table\", row_index, row_sub_index, \
589             agent_instance_hierarchy, \"timestamp\") \
590         SELECT $1, $5, $6, $7, $8, $9 FROM data_ins"
591    );
592    sqlx::query(&sql)
593        .bind(value.response_id())
594        .bind(value.row_index())
595        .bind(value.row_sub_index())
596        .bind(video.url.as_str())
597        .bind(value.message_table())
598        .bind(value.row_index())
599        .bind(value.row_sub_index())
600        .bind(value.agent_instance_hierarchy())
601        .bind(timestamp)
602        .execute(&**pool)
603        .await?;
604    Ok(())
605}
606
607async fn insert_file_part_with_msg<'a>(
608    pool: &Pool,
609    table: &str,
610    value: &RowValue<'a>,
611    file: &File,
612    timestamp: i64,
613) -> Result<(), Error> {
614    let sql = format!(
615        "WITH data_ins AS (\
616            INSERT INTO {table} (response_id, \"index\", part_index, file_data, file_id, filename, file_url) \
617            VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING response_id\
618         )\
619         INSERT INTO objectiveai.messages \
620            (response_id, \"table\", row_index, row_sub_index, \
621             agent_instance_hierarchy, \"timestamp\") \
622         SELECT $1, $8, $9, $10, $11, $12 FROM data_ins"
623    );
624    sqlx::query(&sql)
625        .bind(value.response_id())
626        .bind(value.row_index())
627        .bind(value.row_sub_index())
628        .bind(file.file_data.as_deref())
629        .bind(file.file_id.as_deref())
630        .bind(file.filename.as_deref())
631        .bind(file.file_url.as_deref())
632        .bind(value.message_table())
633        .bind(value.row_index())
634        .bind(value.row_sub_index())
635        .bind(value.agent_instance_hierarchy())
636        .bind(timestamp)
637        .execute(&**pool)
638        .await?;
639    Ok(())
640}
641
642// ---- UPDATE helper: streaming row + messages_queue downgrade -----------
643//
644// The update is parameterized on a placeholder SQL template that uses
645// named tokens for the response_id / row_index / row_sub_index binds
646// plus arbitrary per-table value binds (A, B, C, D). The helper
647// rewrites the tokens into positional $N placeholders and appends the
648// messages-queue downgrade CTE.
649
650#[derive(Clone, Copy)]
651enum BindIdx {
652    Resp,
653    Ri,
654    Rsi,
655}
656
657enum BindVal<'a> {
658    Str(&'a str),
659    OptStr(Option<&'a str>),
660    OptString(Option<String>),
661    Bool(bool),
662}
663
664#[allow(clippy::too_many_arguments)]
665async fn run_update_with_downgrade<'a>(
666    pool: &Pool,
667    update_sql_template: &str,
668    response_id: &str,
669    row_index: i64,
670    row_sub_index: Option<i64>,
671    message_table: MessageTable,
672    agent_instance_hierarchy: &str,
673    extra_binds: &[(&str, BindVal<'a>)],
674    update_where_binds: &[BindIdx],
675) -> Result<(), Error> {
676    // Assign positional indices. Order in the final SQL:
677    //   $1..$N  = update_where_binds in order, then extra_binds in
678    //             declaration order. (We rewrite the template's
679    //             named tokens accordingly.)
680    //   $(N+1)  = message_table
681    //   $(N+2)  = row_index
682    //   $(N+3)  = row_sub_index
683    //   $(N+4)  = agent_instance_hierarchy
684    let mut sql = update_sql_template.to_string();
685    let mut pos = 1usize;
686
687    // Replace each WHERE bind token with its positional index. The
688    // resp/ri/rsi positions inside `sql` are written back into the
689    // template via `sql.replace(...)`; we only need to remember
690    // resp_pos for the downgrade CTE below.
691    let mut resp_pos: Option<usize> = None;
692    for slot in update_where_binds {
693        let idx = pos;
694        pos += 1;
695        match slot {
696            BindIdx::Resp => {
697                sql = sql.replace("$RESP", &format!("${idx}"));
698                resp_pos = Some(idx);
699            }
700            BindIdx::Ri => {
701                sql = sql.replace("$RI", &format!("${idx}"));
702            }
703            BindIdx::Rsi => {
704                sql = sql.replace("$RSI", &format!("${idx}"));
705            }
706        }
707    }
708    let resp_pos = resp_pos.expect("Resp bind required");
709
710    // Replace extra-bind tokens ($A, $B, $C, $D) with positional.
711    for (token, _val) in extra_binds {
712        let idx = pos;
713        pos += 1;
714        sql = sql.replace(&format!("${token}"), &format!("${idx}"));
715    }
716
717    let mt_pos = pos; pos += 1;
718    let ri_for_msg_pos = pos; pos += 1;
719    let rsi_for_msg_pos = pos; pos += 1;
720    let hier_pos = pos;
721
722    let final_sql = format!(
723        "WITH \
724            data_upd AS ({sql} RETURNING response_id),\
725            msg AS (\
726                SELECT \"index\" AS msg_index FROM objectiveai.messages \
727                WHERE response_id = ${resp_pos} \
728                  AND \"table\" = ${mt_pos} \
729                  AND row_index IS NOT DISTINCT FROM ${ri_for_msg_pos} \
730                  AND row_sub_index IS NOT DISTINCT FROM ${rsi_for_msg_pos}\
731            )\
732         UPDATE objectiveai.messages_queue \
733         SET read_index = msg.msg_index - 1 \
734         FROM msg, data_upd \
735         WHERE spawned_agent_instance_hierarchy = ${hier_pos} \
736           AND read_index >= msg.msg_index",
737    );
738
739    let mut q = sqlx::query(&final_sql);
740    // Bind WHERE clause values in their declared order.
741    for slot in update_where_binds {
742        q = match slot {
743            BindIdx::Resp => q.bind(response_id),
744            BindIdx::Ri => q.bind(row_index),
745            BindIdx::Rsi => q.bind(row_sub_index),
746        };
747    }
748    // Bind extra values.
749    for (_, val) in extra_binds {
750        q = match val {
751            BindVal::Str(s) => q.bind(*s),
752            BindVal::OptStr(s) => q.bind(*s),
753            BindVal::OptString(s) => q.bind(s.clone()),
754            BindVal::Bool(b) => q.bind(*b),
755        };
756    }
757    // Bind messages-row identification + agent hierarchy.
758    q = q.bind(message_table);
759    q = q.bind(row_index);
760    q = q.bind(row_sub_index);
761    q = q.bind(agent_instance_hierarchy);
762
763    q.execute(&**pool).await?;
764    Ok(())
765}
766
767// =====================================================================
768// Tier blob writes
769// =====================================================================
770
771#[derive(Debug, Clone, Copy, PartialEq, Eq)]
772pub enum Tier {
773    Agent,
774    Vector,
775    Function,
776}
777
778impl Tier {
779    pub fn request_table(self) -> &'static str {
780        match self {
781            Tier::Agent => "objectiveai.agent_completion_requests",
782            Tier::Vector => "objectiveai.vector_completion_requests",
783            Tier::Function => "objectiveai.function_execution_requests",
784        }
785    }
786    pub fn response_table(self) -> &'static str {
787        match self {
788            Tier::Agent => "objectiveai.agent_completion_responses",
789            Tier::Vector => "objectiveai.vector_completion_responses",
790            Tier::Function => "objectiveai.function_execution_responses",
791        }
792    }
793    /// The matching [`MessageTable`] for this tier's request blob.
794    /// Response blobs don't emit messages so there's no equivalent.
795    pub fn request_message_table(self) -> MessageTable {
796        match self {
797            Tier::Agent => MessageTable::AgentCompletionRequest,
798            Tier::Vector => MessageTable::VectorCompletionRequest,
799            Tier::Function => MessageTable::FunctionExecutionRequest,
800        }
801    }
802}
803
804/// INSERT the request blob. Called once per stream, on first chunk
805/// arrival. Request blobs don't carry `agent_instance_hierarchy` —
806/// they're shared across every agent that participates in the stream.
807/// The per-agent "the request was made for me" linkage lives in
808/// `objectiveai.messages` and is written separately by
809/// [`insert_request_messages_row`] the first time each agent appears
810/// in the chunk's row iterator.
811pub async fn insert_request_blob<P: Serialize>(
812    pool: &Pool,
813    tier: Tier,
814    response_id: &str,
815    params: &P,
816    sender_agent_instance_hierarchy: &str,
817    timestamp: i64,
818) -> Result<(), Error> {
819    let body = serde_json::to_value(params)?;
820    let sql = format!(
821        "INSERT INTO {table} \
822            (response_id, body, created_at, sender_agent_instance_hierarchy) \
823         VALUES ($1, $2, $3, $4)",
824        table = tier.request_table()
825    );
826    sqlx::query(&sql)
827        .bind(response_id)
828        .bind(sqlx::types::Json(body))
829        .bind(timestamp)
830        .bind(sender_agent_instance_hierarchy)
831        .execute(&**pool)
832        .await?;
833    Ok(())
834}
835
836/// INSERT a `objectiveai.messages` row that registers this stream's request
837/// blob in the agent's history. Called once per (stream, agent) pair
838/// — the writer tracks which agents it has already seen and only
839/// emits this row the first time it encounters a new one in the row
840/// iterator. By postgres's BIGSERIAL `"index"` assignment, this row
841/// is guaranteed to land earlier in the agent's history than any
842/// subsequent streaming-content row that the same writer call
843/// sequences after it.
844pub async fn insert_request_messages_row(
845    pool: &Pool,
846    tier: Tier,
847    response_id: &str,
848    agent_instance_hierarchy: &str,
849    timestamp: i64,
850) -> Result<(), Error> {
851    sqlx::query(
852        "INSERT INTO objectiveai.messages \
853            (response_id, \"table\", row_index, row_sub_index, \
854             agent_instance_hierarchy, \"timestamp\") \
855         VALUES ($1, $2, NULL, NULL, $3, $4)",
856    )
857    .bind(response_id)
858    .bind(tier.request_message_table())
859    .bind(agent_instance_hierarchy)
860    .bind(timestamp)
861    .execute(&**pool)
862    .await?;
863    Ok(())
864}
865
866/// INSERT the response tier blob (first tick only). Response blobs
867/// don't emit messages — they're the latest snapshot, not events.
868/// Tier-symmetric: every tier's response table now has the same
869/// `(response_id, body, created_at, inserted_at)` shape.
870pub async fn insert_response_blob<C: Serialize>(
871    pool: &Pool,
872    tier: Tier,
873    response_id: &str,
874    chunk: &C,
875    created_at: i64,
876) -> Result<(), Error> {
877    let body = serde_json::to_value(chunk)?;
878    let sql = format!(
879        "INSERT INTO {table} (response_id, body, created_at) VALUES ($1, $2, $3)",
880        table = tier.response_table()
881    );
882    sqlx::query(&sql)
883        .bind(response_id)
884        .bind(sqlx::types::Json(body))
885        .bind(created_at)
886        .execute(&**pool)
887        .await?;
888    Ok(())
889}
890
891/// UPDATE the response tier blob (subsequent ticks).
892pub async fn update_response_blob<C: Serialize>(
893    pool: &Pool,
894    tier: Tier,
895    response_id: &str,
896    chunk: &C,
897    created_at: i64,
898) -> Result<(), Error> {
899    let body = serde_json::to_value(chunk)?;
900    let sql = format!(
901        "UPDATE {table} SET body = $2, created_at = $3 WHERE response_id = $1",
902        table = tier.response_table()
903    );
904    sqlx::query(&sql)
905        .bind(response_id)
906        .bind(sqlx::types::Json(body))
907        .bind(created_at)
908        .execute(&**pool)
909        .await?;
910    Ok(())
911}