CST-agent/src/session.c

#ifndef _DEFAULT_SOURCE
#define _DEFAULT_SOURCE
#endif
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include "cJSON.h"
#include "cst_agent.h"
#include "cst_attach.h"
#include "cst_util.h"

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "os_compat.h"

struct cst_session {
  char path[4096];
  cst_message_t *msgs;
  int nmsgs;
  int cap;
  FILE *fp;
  pthread_mutex_t mu;
};

static char *xstrdup(const char *s) {
  if (!s)
    return NULL;
  size_t n = strlen(s);
  char *p = malloc(n + 1);
  if (!p)
    return NULL;
  memcpy(p, s, n + 1);
  return p;
}

static const char *role_str(cst_role_t r) {
  switch (r) {
  case CST_ROLE_SYSTEM:
    return "system";
  case CST_ROLE_USER:
    return "user";
  case CST_ROLE_ASSISTANT:
    return "assistant";
  case CST_ROLE_TOOL:
    return "tool";
  }
  return "user";
}

static cst_role_t role_from(const char *s) {
  if (!s)
    return CST_ROLE_USER;
  if (strcmp(s, "system") == 0)
    return CST_ROLE_SYSTEM;
  if (strcmp(s, "assistant") == 0)
    return CST_ROLE_ASSISTANT;
  if (strcmp(s, "tool") == 0)
    return CST_ROLE_TOOL;
  return CST_ROLE_USER;
}

static void msg_free(cst_message_t *m) {
  free(m->content);
  free(m->tool_calls_json);
  free(m->meta_json);
  free(m->reasoning);
  cst_attachments_free(m->attach, m->nattach);
  m->content = m->tool_calls_json = m->meta_json = m->reasoning = NULL;
  m->attach = NULL;
  m->nattach = 0;
}

static int session_grow(cst_session_t *s) {
  if (s->nmsgs < s->cap)
    return 0;
  int ncap = s->cap ? s->cap * 2 : 32;
  cst_message_t *nm = realloc(s->msgs, (size_t)ncap * sizeof(cst_message_t));
  if (!nm)
    return -1;
  s->msgs = nm;
  s->cap = ncap;
  return 0;
}

static int session_push_mem(cst_session_t *s, const cst_message_t *src) {
  if (session_grow(s) != 0)
    return -1;
  cst_message_t *d = &s->msgs[s->nmsgs];
  memset(d, 0, sizeof(*d));
  d->role = src->role;
  d->content = xstrdup(src->content ? src->content : "");
  if (!d->content)
    return -1;
  snprintf(d->tool_call_id, sizeof(d->tool_call_id), "%s", src->tool_call_id);
  snprintf(d->tool_name, sizeof(d->tool_name), "%s", src->tool_name);
  if (src->tool_calls_json && src->tool_calls_json[0])
    d->tool_calls_json = xstrdup(src->tool_calls_json);
  if (src->meta_json && src->meta_json[0])
    d->meta_json = xstrdup(src->meta_json);
  if (src->reasoning && src->reasoning[0])
    d->reasoning = xstrdup(src->reasoning);
  if (src->attach && src->nattach > 0) {
    d->attach = cst_attachments_dup(src->attach, src->nattach);
    if (!d->attach) {
      msg_free(d);
      return -1;
    }
    d->nattach = src->nattach;
  }
  s->nmsgs++;
  return 0;
}

static const char *jstr(const cJSON *o, const char *key) {
  const cJSON *it = cJSON_GetObjectItemCaseSensitive(o, key);
  return cJSON_IsString(it) ? it->valuestring : NULL;
}

static void load_existing(cst_session_t *s) {
  FILE *rp = fopen(s->path, "rb");
  if (!rp)
    return;
  char *line = NULL;
  size_t cap = 0;
  while (os_getline(&line, &cap, rp) != -1) {
    cJSON *o = cJSON_Parse(line);
    if (!o || !cJSON_IsObject(o)) {
      cJSON_Delete(o);
      continue;
    }
    const char *role = jstr(o, "role");
    if (!role) {
      cJSON_Delete(o);
      continue;
    }
    cst_message_t m = {0};
    m.role = role_from(role);
    m.content = (char *)(jstr(o, "content") ? jstr(o, "content") : "");
    const char *tid = jstr(o, "tool_call_id");
    const char *tname = jstr(o, "tool_name");
    if (tid)
      snprintf(m.tool_call_id, sizeof(m.tool_call_id), "%s", tid);
    if (tname)
      snprintf(m.tool_name, sizeof(m.tool_name), "%s", tname);
    char *tcj = NULL;
    const cJSON *tc = cJSON_GetObjectItemCaseSensitive(o, "tool_calls");
    if (cJSON_IsArray(tc))
      tcj = cJSON_PrintUnformatted(tc);
    else if (jstr(o, "tool_calls_json")) /* legacy string field */
      tcj = xstrdup(jstr(o, "tool_calls_json"));
    m.tool_calls_json = tcj;
    char *meta = NULL;
    const cJSON *mo = cJSON_GetObjectItemCaseSensitive(o, "meta");
    if (cJSON_IsObject(mo))
      meta = cJSON_PrintUnformatted(mo);
    m.meta_json = meta;
    m.reasoning = (char *)jstr(o, "reasoning");
    const cJSON *at = cJSON_GetObjectItemCaseSensitive(o, "attachments");
    if (cJSON_IsArray(at))
      m.attach = cst_attachments_from_json(at, &m.nattach);
    session_push_mem(s, &m);
    free(tcj);
    free(meta);
    cst_attachments_free(m.attach, m.nattach);
    cJSON_Delete(o);
  }
  free(line);
  fclose(rp);
}

static void free_all(cst_session_t *s) {
  for (int i = 0; i < s->nmsgs; i++)
    msg_free(&s->msgs[i]);
  s->nmsgs = 0;
}

cst_session_t *cst_session_open(const char *path) {
  cst_session_t *s = calloc(1, sizeof(*s));
  if (!s)
    return NULL;
  pthread_mutex_init(&s->mu, NULL);
  snprintf(s->path, sizeof(s->path), "%s", path ? path : "session.jsonl");
  load_existing(s);
  s->fp = fopen(s->path, "a+b");
  if (!s->fp || session_grow(s) != 0) {
    if (s->fp)
      fclose(s->fp);
    free_all(s);
    free(s->msgs);
    pthread_mutex_destroy(&s->mu);
    free(s);
    return NULL;
  }
  return s;
}

void cst_session_close(cst_session_t *s) {
  if (!s)
    return;
  free_all(s);
  free(s->msgs);
  if (s->fp)
    fclose(s->fp);
  pthread_mutex_destroy(&s->mu);
  free(s);
}

void cst_session_lock(cst_session_t *s) {
  if (s)
    pthread_mutex_lock(&s->mu);
}

void cst_session_unlock(cst_session_t *s) {
  if (s)
    pthread_mutex_unlock(&s->mu);
}

static int write_line(FILE *fp, const cst_message_t *msg) {
  cJSON *o = cJSON_CreateObject();
  if (!o)
    return -1;
  cJSON_AddStringToObject(o, "role", role_str(msg->role));
  cJSON_AddStringToObject(o, "content", msg->content ? msg->content : "");
  if (msg->tool_call_id[0])
    cJSON_AddStringToObject(o, "tool_call_id", msg->tool_call_id);
  if (msg->tool_name[0])
    cJSON_AddStringToObject(o, "tool_name", msg->tool_name);
  if (msg->tool_calls_json && msg->tool_calls_json[0]) {
    cJSON *tc = cJSON_Parse(msg->tool_calls_json);
    if (cJSON_IsArray(tc))
      cJSON_AddItemToObject(o, "tool_calls", tc);
    else {
      cJSON_Delete(tc);
      cJSON_AddStringToObject(o, "tool_calls_json", msg->tool_calls_json);
    }
  }
  if (msg->reasoning && msg->reasoning[0])
    cJSON_AddStringToObject(o, "reasoning", msg->reasoning);
  if (msg->meta_json && msg->meta_json[0]) {
    cJSON *mo = cJSON_Parse(msg->meta_json);
    if (cJSON_IsObject(mo))
      cJSON_AddItemToObject(o, "meta", mo);
    else
      cJSON_Delete(mo);
  }
  if (msg->attach && msg->nattach > 0) {
    /* One writer for every attachment (cst_attachments_to_json): a field it
     * carries must not have to be remembered here too -- forgetting
     * `for_user` would put a delivered file back in front of the model. */
    char *json = cst_attachments_to_json(msg->attach, msg->nattach);
    if (json) {
      cJSON *arr = cJSON_Parse(json);
      free(json);
      if (cJSON_IsArray(arr))
        cJSON_AddItemToObject(o, "attachments", arr);
      else
        cJSON_Delete(arr);
    }
  }
  char *txt = cJSON_PrintUnformatted(o);
  cJSON_Delete(o);
  if (!txt)
    return -1;
  fputs(txt, fp);
  fputc('\n', fp);
  free(txt);
  fflush(fp);
  return 0;
}

int cst_session_append(cst_session_t *s, const cst_message_t *msg) {
  if (!s || !msg)
    return -1;
  pthread_mutex_lock(&s->mu);
  int rc = session_push_mem(s, msg);
  if (rc == 0 && s->fp)
    rc = write_line(s->fp, msg);
  pthread_mutex_unlock(&s->mu);
  return rc;
}

int cst_session_count(const cst_session_t *s) { return s ? s->nmsgs : 0; }

const cst_message_t *cst_session_get(const cst_session_t *s, int i) {
  if (!s || i < 0 || i >= s->nmsgs)
    return NULL;
  return &s->msgs[i];
}

const char *cst_session_path(const cst_session_t *s) {
  return s ? s->path : NULL;
}

static int reopen_truncate(cst_session_t *s) {
  if (s->fp) {
    fclose(s->fp);
    s->fp = NULL;
  }
  s->fp = fopen(s->path, "w+b");
  return s->fp ? 0 : -1;
}

int cst_session_clear(cst_session_t *s, int keep_system) {
  if (!s)
    return -1;
  pthread_mutex_lock(&s->mu);
  cst_message_t sys = {0};
  int have_sys = 0;
  if (keep_system) {
    for (int i = 0; i < s->nmsgs; i++) {
      if (s->msgs[i].role == CST_ROLE_SYSTEM) {
        sys.role = CST_ROLE_SYSTEM;
        sys.content = xstrdup(s->msgs[i].content ? s->msgs[i].content : "");
        sys.meta_json = xstrdup(s->msgs[i].meta_json);
        have_sys = sys.content != NULL;
        break;
      }
    }
  }
  free_all(s);
  int rc = reopen_truncate(s);
  if (rc == 0 && have_sys) {
    session_push_mem(s, &sys);
    write_line(s->fp, &sys);
  }
  free(sys.content);
  free(sys.meta_json);
  pthread_mutex_unlock(&s->mu);
  return rc;
}

int cst_session_rewrite(cst_session_t *s, const cst_message_t *msgs, int n) {
  if (!s || n < 0)
    return -1;
  /* msgs may point into this session's own storage: copy before freeing. */
  cst_message_t *copy = calloc((size_t)(n ? n : 1), sizeof(*copy));
  if (!copy)
    return -1;
  for (int i = 0; i < n; i++) {
    copy[i] = msgs[i];
    copy[i].content = strdup(msgs[i].content ? msgs[i].content : "");
    copy[i].tool_calls_json =
        msgs[i].tool_calls_json ? strdup(msgs[i].tool_calls_json) : NULL;
    copy[i].meta_json = msgs[i].meta_json ? strdup(msgs[i].meta_json) : NULL;
    copy[i].reasoning = msgs[i].reasoning ? strdup(msgs[i].reasoning) : NULL;
    copy[i].attach = msgs[i].attach && msgs[i].nattach > 0
                         ? cst_attachments_dup(msgs[i].attach, msgs[i].nattach)
                         : NULL;
    copy[i].nattach = copy[i].attach ? msgs[i].nattach : 0;
  }
  pthread_mutex_lock(&s->mu);
  free_all(s);
  int rc = reopen_truncate(s);
  for (int i = 0; rc == 0 && i < n; i++) {
    if (session_push_mem(s, &copy[i]) != 0 || write_line(s->fp, &copy[i]) != 0)
      rc = -1;
  }
  pthread_mutex_unlock(&s->mu);
  for (int i = 0; i < n; i++) {
    free(copy[i].content);
    free(copy[i].tool_calls_json);
    free(copy[i].meta_json);
    free(copy[i].reasoning);
    cst_attachments_free(copy[i].attach, copy[i].nattach);
  }
  free(copy);
  return rc;
}

int cst_message_is_tool_result(const cst_message_t *m) {
  return m && (m->role == CST_ROLE_TOOL ||
               (m->role == CST_ROLE_USER && m->tool_call_id[0]));
}

static void md_fence(cst_buf_t *b, const char *s, size_t max) {
  size_t n = strlen(s);
  if (n > max)
    n = cst_utf8_floor(s, max);
  cst_buf_puts(b, "````\n");
  cst_buf_append(b, s, n);
  cst_buf_puts(b, n < strlen(s) ? "\n…(truncated)\n````\n\n" : "\n````\n\n");
}

char *cst_session_export_md(cst_session_t *s, const char *title,
                            const char *sid, const char *model) {
  if (!s)
    return NULL;
  cst_buf_t b = {0};
  cst_buf_printf(&b, "# %s\n\nsession `%s` · model `%s`\n\n",
                 title ? title : "", sid ? sid : "", model ? model : "");
  cst_session_lock(s);
  for (int i = 0; i < cst_session_count(s); i++) {
    const cst_message_t *m = cst_session_get(s, i);
    const char *c = m->content ? m->content : "";
    if (m->role == CST_ROLE_SYSTEM)
      continue;
    if (cst_message_is_tool_result(m)) {
      cst_buf_printf(&b, "**Tool result** `%s`\n\n", m->tool_name);
      md_fence(&b, c, 4000);
    } else if (cst_message_is_note(m)) {
      cst_buf_puts(&b, "> ");
      cst_buf_append(&b, c, strcspn(c, "\n"));
      cst_buf_puts(&b, "\n\n");
    } else if (m->role == CST_ROLE_USER) {
      cst_buf_puts(&b, "## User\n\n");
      cst_buf_puts(&b, c);
      cst_buf_puts(&b, "\n\n");
      for (int k = 0; k < m->nattach; k++) {
        char d[128];
        cst_attach_describe(&m->attach[k], d, sizeof(d));
        cst_buf_printf(&b, "- %s `%s` (%s)\n",
                       cst_attach_is_image(&m->attach[k]) ? "image"
                                                         : "document",
                       m->attach[k].name, d);
      }
      if (m->nattach)
        cst_buf_puts(&b, "\n");
    } else {
      cst_buf_puts(&b, "## Assistant\n\n");
      cst_buf_puts(&b, c);
      cst_buf_puts(&b, "\n\n");
      if (m->tool_calls_json && m->tool_calls_json[0]) {
        cst_buf_puts(&b, "**Tool calls**\n\n");
        md_fence(&b, m->tool_calls_json, 4000);
      }
    }
  }
  cst_session_unlock(s);
  return cst_buf_steal(&b);
}