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 <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#ifdef _WIN32
typedef long long cst_ssize_t;
static cst_ssize_t cst_getline(char **lineptr, size_t *n, FILE *stream) {
  if (!lineptr || !n || !stream)
    return -1;
  if (!*lineptr || *n == 0) {
    *n = 256;
    *lineptr = (char *)malloc(*n);
    if (!*lineptr)
      return -1;
  }
  size_t len = 0;
  int c = 0;
  while ((c = fgetc(stream)) != EOF) {
    if (len + 2 > *n) {
      size_t nn = (*n) * 2;
      char *nb = (char *)realloc(*lineptr, nn);
      if (!nb)
        return -1;
      *lineptr = nb;
      *n = nn;
    }
    (*lineptr)[len++] = (char)c;
    if (c == '\n')
      break;
  }
  if (len == 0 && c == EOF)
    return -1;
  (*lineptr)[len] = 0;
  return (cst_ssize_t)len;
}
#define getline cst_getline
#endif

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);
  m->content = m->tool_calls_json = m->meta_json = NULL;
}

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);
  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 (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;
    session_push_mem(s, &m);
    free(tcj);
    free(meta);
    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->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);
  }
  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 : "");
        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);
  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;
  }
  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);
  return rc;
}