Refactor (#45)

* progress

* refactor

* backend refactor

* test

* minor stdio changes

* streeeeeaming

* stream statuses

* change required user input

* something

* progress

* new beginnings

* progress

* edge case

* uuhh the claude wrapper def worked before this commit

* progress

* ai slop that works

* tests

* tests

* plan mode handling

* progress

* clean up

* bump

* migrations

* readmes

* no text truncation

* consistent poll interval

* fix time

---------

Co-authored-by: Kartik Sarangmath <kartiksarangmath@Kartiks-MacBook-Air.local>
This commit is contained in:
ksarangmath
2025-08-04 01:44:42 -07:00
committed by GitHub
parent 140be5b512
commit eaca5a0ad0
44 changed files with 5178 additions and 1816 deletions

View File

@@ -0,0 +1,293 @@
"""Add messages table
Revision ID: 40d4252deb5b
Revises: dc285eabea90
Create Date: 2025-07-31 11:25:30.567076
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = "40d4252deb5b"
down_revision: Union[str, None] = "dc285eabea90"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.create_table(
"messages",
sa.Column("id", sa.UUID(), nullable=False),
sa.Column("agent_instance_id", sa.UUID(), nullable=False),
sa.Column(
"sender_type", sa.Enum("AGENT", "USER", name="sendertype"), nullable=False
),
sa.Column("content", sa.Text(), nullable=False),
sa.Column("created_at", sa.DateTime(), nullable=False),
sa.Column("requires_user_input", sa.Boolean(), nullable=False),
sa.Column(
"message_metadata", postgresql.JSONB(astext_type=sa.Text()), nullable=True
),
sa.ForeignKeyConstraint(
["agent_instance_id"],
["agent_instances.id"],
),
sa.PrimaryKeyConstraint("id"),
)
op.create_index(
"idx_messages_instance_created",
"messages",
["agent_instance_id", "created_at"],
unique=False,
)
# Migrate data from old tables to messages table
# This combines AgentStep, AgentQuestion, and AgentUserFeedback into Messages
# ordered by timestamp
op.execute("""
-- Insert agent steps as agent messages
INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
SELECT
id,
agent_instance_id,
'AGENT',
description,
created_at,
false,
jsonb_build_object('source', 'agent_step', 'step_number', step_number)
FROM agent_steps;
-- Insert agent questions as agent messages
INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
SELECT
id,
agent_instance_id,
'AGENT',
question_text,
asked_at,
CASE WHEN answer_text IS NULL THEN true ELSE false END,
jsonb_build_object('source', 'agent_question')
FROM agent_questions;
-- Insert question answers as user messages (only for answered questions)
INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
SELECT
gen_random_uuid(),
agent_instance_id,
'USER',
answer_text,
answered_at,
false,
jsonb_build_object('source', 'question_answer', 'question_id', id::text, 'answered_by_user_id', answered_by_user_id::text)
FROM agent_questions
WHERE answer_text IS NOT NULL AND answered_at IS NOT NULL;
-- Insert user feedback as user messages
INSERT INTO messages (id, agent_instance_id, sender_type, content, created_at, requires_user_input, message_metadata)
SELECT
id,
agent_instance_id,
'USER',
feedback_text,
created_at,
false,
jsonb_build_object('source', 'user_feedback', 'created_by_user_id', created_by_user_id::text)
FROM agent_user_feedback;
-- Mark all agent instances as completed
UPDATE agent_instances SET status = 'COMPLETED' WHERE status != 'COMPLETED';
""")
# Add last_read_message_id to agent_instances
op.add_column(
"agent_instances",
sa.Column("last_read_message_id", postgresql.UUID(as_uuid=True), nullable=True),
)
# Create the foreign key constraint after both tables exist to avoid circular dependency
op.create_foreign_key(
"fk_agent_instances_last_read_message",
"agent_instances",
"messages",
["last_read_message_id"],
["id"],
ondelete="SET NULL",
)
# Create combined function for message notifications
op.execute("""
CREATE OR REPLACE FUNCTION notify_message_change() RETURNS trigger AS $$
DECLARE
channel_name text;
payload text;
event_type text;
BEGIN
-- Create channel name based on instance ID
channel_name := 'message_channel_' || NEW.agent_instance_id::text;
-- Determine event type
IF TG_OP = 'INSERT' THEN
event_type := 'message_insert';
ELSIF TG_OP = 'UPDATE' THEN
event_type := 'message_update';
END IF;
-- Create JSON payload with message data
payload := json_build_object(
'event_type', event_type,
'id', NEW.id,
'agent_instance_id', NEW.agent_instance_id,
'sender_type', NEW.sender_type,
'content', NEW.content,
'created_at', to_char(NEW.created_at AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'),
'requires_user_input', NEW.requires_user_input,
'message_metadata', NEW.message_metadata,
'old_requires_user_input', CASE
WHEN TG_OP = 'UPDATE' THEN OLD.requires_user_input
ELSE NULL
END
)::text;
-- Send notification (quote channel name for UUIDs with hyphens)
EXECUTE format('NOTIFY %I, %L', channel_name, payload);
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
""")
# Create single trigger for both INSERT and UPDATE on messages table
op.execute("""
CREATE TRIGGER message_change_notify
AFTER INSERT OR UPDATE ON messages
FOR EACH ROW
EXECUTE FUNCTION notify_message_change();
""")
# Create function for status change notifications
op.execute("""
CREATE OR REPLACE FUNCTION notify_status_change() RETURNS trigger AS $$
DECLARE
channel_name text;
payload text;
BEGIN
-- Only notify if status actually changed
IF OLD.status IS DISTINCT FROM NEW.status THEN
-- Create channel name based on instance ID
channel_name := 'message_channel_' || NEW.id::text;
-- Create JSON payload with status update data
payload := json_build_object(
'event_type', 'status_update',
'instance_id', NEW.id,
'status', NEW.status,
'timestamp', to_char(NOW() AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"')
)::text;
-- Send notification (quote channel name for UUIDs with hyphens)
EXECUTE format('NOTIFY %I, %L', channel_name, payload);
END IF;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
""")
# Create trigger on agent_instances table for status updates
op.execute("""
CREATE TRIGGER agent_instance_status_notify
AFTER UPDATE OF status ON agent_instances
FOR EACH ROW
EXECUTE FUNCTION notify_status_change();
""")
# Drop the old tables after migration
op.drop_table("agent_user_feedback")
op.drop_table("agent_questions")
op.drop_table("agent_steps")
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
# Recreate the old tables first
op.create_table(
"agent_steps",
sa.Column("id", sa.UUID(), nullable=False),
sa.Column("agent_instance_id", sa.UUID(), nullable=False),
sa.Column("step_number", sa.Integer(), nullable=False),
sa.Column("description", sa.Text(), nullable=False),
sa.Column("created_at", sa.DateTime(), nullable=False),
sa.ForeignKeyConstraint(
["agent_instance_id"],
["agent_instances.id"],
),
sa.PrimaryKeyConstraint("id"),
)
op.create_table(
"agent_questions",
sa.Column("id", sa.UUID(), nullable=False),
sa.Column("agent_instance_id", sa.UUID(), nullable=False),
sa.Column("question_text", sa.Text(), nullable=False),
sa.Column("answer_text", sa.Text(), nullable=True),
sa.Column("answered_by_user_id", sa.UUID(), nullable=True),
sa.Column("asked_at", sa.DateTime(), nullable=False),
sa.Column("answered_at", sa.DateTime(), nullable=True),
sa.Column("is_active", sa.Boolean(), nullable=False),
sa.ForeignKeyConstraint(
["agent_instance_id"],
["agent_instances.id"],
),
sa.ForeignKeyConstraint(
["answered_by_user_id"],
["users.id"],
),
sa.PrimaryKeyConstraint("id"),
)
op.create_table(
"agent_user_feedback",
sa.Column("id", sa.UUID(), nullable=False),
sa.Column("agent_instance_id", sa.UUID(), nullable=False),
sa.Column("created_by_user_id", sa.UUID(), nullable=False),
sa.Column("feedback_text", sa.Text(), nullable=False),
sa.Column("created_at", sa.DateTime(), nullable=False),
sa.Column("retrieved_at", sa.DateTime(), nullable=True),
sa.ForeignKeyConstraint(
["agent_instance_id"],
["agent_instances.id"],
),
sa.ForeignKeyConstraint(
["created_by_user_id"],
["users.id"],
),
sa.PrimaryKeyConstraint("id"),
)
# Note: We cannot fully restore the original data as some information is lost
# (e.g., step_number ordering, which user created feedback vs answered questions)
# This is a best-effort restoration
# Drop status change trigger and function
op.execute(
"DROP TRIGGER IF EXISTS agent_instance_status_notify ON agent_instances;"
)
op.execute("DROP FUNCTION IF EXISTS notify_status_change();")
# Drop message trigger and function
op.execute("DROP TRIGGER IF EXISTS message_change_notify ON messages;")
op.execute("DROP FUNCTION IF EXISTS notify_message_change();")
op.drop_constraint(
"fk_agent_instances_last_read_message", "agent_instances", type_="foreignkey"
)
op.drop_column("agent_instances", "last_read_message_id")
op.drop_index("idx_messages_instance_created", table_name="messages")
op.drop_table("messages")
# ### end Alembic commands ###

View File

@@ -1,11 +1,9 @@
from .enums import AgentStatus
from .enums import AgentStatus, SenderType
from .models import (
AgentInstance,
AgentQuestion,
AgentStep,
AgentUserFeedback,
APIKey,
Base,
Message,
PushToken,
User,
UserAgent,
@@ -20,12 +18,11 @@ __all__ = [
"User",
"UserAgent",
"AgentInstance",
"AgentStep",
"AgentQuestion",
"AgentStatus",
"AgentUserFeedback",
"APIKey",
"Message",
"PushToken",
"SenderType",
"Subscription",
"BillingEvent",
]

View File

@@ -2,11 +2,16 @@ from enum import Enum
class AgentStatus(str, Enum):
ACTIVE = "active"
AWAITING_INPUT = "awaiting_input"
PAUSED = "paused"
STALE = "stale"
COMPLETED = "completed"
FAILED = "failed"
KILLED = "killed"
DISCONNECTED = "disconnected"
ACTIVE = "ACTIVE"
AWAITING_INPUT = "AWAITING_INPUT"
PAUSED = "PAUSED"
STALE = "STALE"
COMPLETED = "COMPLETED"
FAILED = "FAILED"
KILLED = "KILLED"
DISCONNECTED = "DISCONNECTED"
class SenderType(str, Enum):
AGENT = "AGENT"
USER = "USER"

View File

@@ -3,7 +3,7 @@ from uuid import UUID, uuid4
from typing import TYPE_CHECKING
from sqlalchemy import ForeignKey, Index, String, Text, UniqueConstraint
from sqlalchemy.dialects.postgresql import UUID as PostgresUUID
from sqlalchemy.dialects.postgresql import UUID as PostgresUUID, JSONB
from sqlalchemy.orm import (
DeclarativeBase, # type: ignore[attr-defined]
Mapped, # type: ignore[attr-defined]
@@ -12,7 +12,7 @@ from sqlalchemy.orm import (
validates,
)
from .enums import AgentStatus
from .enums import AgentStatus, SenderType
from .utils import is_valid_git_diff
if TYPE_CHECKING:
@@ -57,12 +57,6 @@ class User(Base):
agent_instances: Mapped[list["AgentInstance"]] = relationship(
"AgentInstance", back_populates="user"
)
answered_questions: Mapped[list["AgentQuestion"]] = relationship(
"AgentQuestion", back_populates="answered_by_user"
)
feedback: Mapped[list["AgentUserFeedback"]] = relationship(
"AgentUserFeedback", back_populates="created_by_user"
)
api_keys: Mapped[list["APIKey"]] = relationship("APIKey", back_populates="user")
user_agents: Mapped[list["UserAgent"]] = relationship(
"UserAgent", back_populates="user"
@@ -130,22 +124,33 @@ class AgentInstance(Base):
)
ended_at: Mapped[datetime | None] = mapped_column(default=None)
git_diff: Mapped[str | None] = mapped_column(Text, default=None)
last_read_message_id: Mapped[UUID | None] = mapped_column(
ForeignKey(
"messages.id",
use_alter=True,
name="fk_agent_instances_last_read_message",
ondelete="SET NULL",
),
type_=PostgresUUID(as_uuid=True),
default=None,
)
# Relationships
user_agent: Mapped["UserAgent"] = relationship(
"UserAgent", back_populates="instances"
)
user: Mapped["User"] = relationship("User", back_populates="agent_instances")
steps: Mapped[list["AgentStep"]] = relationship(
"AgentStep", back_populates="instance", order_by="AgentStep.created_at"
)
questions: Mapped[list["AgentQuestion"]] = relationship(
"AgentQuestion", back_populates="instance", order_by="AgentQuestion.asked_at"
)
user_feedback: Mapped[list["AgentUserFeedback"]] = relationship(
"AgentUserFeedback",
messages: Mapped[list["Message"]] = relationship(
"Message",
back_populates="instance",
order_by="AgentUserFeedback.created_at",
order_by="Message.created_at",
foreign_keys="Message.agent_instance_id",
)
last_read_message: Mapped["Message | None"] = relationship(
"Message",
foreign_keys=[last_read_message_id],
post_update=True,
passive_deletes=True,
)
@validates("git_diff")
@@ -154,7 +159,7 @@ class AgentInstance(Base):
Raises ValueError if the git diff is invalid.
"""
if value is None:
if value is None or value == "":
return value
if not is_valid_git_diff(value):
@@ -163,81 +168,6 @@ class AgentInstance(Base):
return value
class AgentStep(Base):
__tablename__ = "agent_steps"
id: Mapped[UUID] = mapped_column(
PostgresUUID(as_uuid=True), primary_key=True, default=uuid4
)
agent_instance_id: Mapped[UUID] = mapped_column(
ForeignKey("agent_instances.id"), type_=PostgresUUID(as_uuid=True)
)
step_number: Mapped[int] = mapped_column()
description: Mapped[str] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
default=lambda: datetime.now(timezone.utc)
)
# Relationships
instance: Mapped["AgentInstance"] = relationship(
"AgentInstance", back_populates="steps"
)
class AgentQuestion(Base):
__tablename__ = "agent_questions"
id: Mapped[UUID] = mapped_column(
PostgresUUID(as_uuid=True), primary_key=True, default=uuid4
)
agent_instance_id: Mapped[UUID] = mapped_column(
ForeignKey("agent_instances.id"), type_=PostgresUUID(as_uuid=True)
)
question_text: Mapped[str] = mapped_column(Text)
answer_text: Mapped[str | None] = mapped_column(Text, default=None)
answered_by_user_id: Mapped[UUID | None] = mapped_column(
ForeignKey("users.id"), type_=PostgresUUID(as_uuid=True), default=None
)
asked_at: Mapped[datetime] = mapped_column(
default=lambda: datetime.now(timezone.utc)
)
answered_at: Mapped[datetime | None] = mapped_column(default=None)
is_active: Mapped[bool] = mapped_column(default=True)
# Relationships
instance: Mapped["AgentInstance"] = relationship(
"AgentInstance", back_populates="questions"
)
answered_by_user: Mapped["User | None"] = relationship(
"User", back_populates="answered_questions"
)
class AgentUserFeedback(Base):
__tablename__ = "agent_user_feedback"
id: Mapped[UUID] = mapped_column(
PostgresUUID(as_uuid=True), primary_key=True, default=uuid4
)
agent_instance_id: Mapped[UUID] = mapped_column(
ForeignKey("agent_instances.id"), type_=PostgresUUID(as_uuid=True)
)
created_by_user_id: Mapped[UUID] = mapped_column(
ForeignKey("users.id"), type_=PostgresUUID(as_uuid=True)
)
feedback_text: Mapped[str] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
default=lambda: datetime.now(timezone.utc)
)
retrieved_at: Mapped[datetime | None] = mapped_column(default=None)
# Relationships
instance: Mapped["AgentInstance"] = relationship(
"AgentInstance", back_populates="user_feedback"
)
created_by_user: Mapped["User"] = relationship("User", back_populates="feedback")
class APIKey(Base):
__tablename__ = "api_keys"
@@ -286,3 +216,31 @@ class PushToken(Base):
# Relationships
user: Mapped["User"] = relationship("User", back_populates="push_tokens")
class Message(Base):
__tablename__ = "messages"
__table_args__ = (
Index("idx_messages_instance_created", "agent_instance_id", "created_at"),
)
id: Mapped[UUID] = mapped_column(
PostgresUUID(as_uuid=True), primary_key=True, default=uuid4
)
agent_instance_id: Mapped[UUID] = mapped_column(
ForeignKey("agent_instances.id"), type_=PostgresUUID(as_uuid=True)
)
sender_type: Mapped[SenderType] = mapped_column()
content: Mapped[str] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
default=lambda: datetime.now(timezone.utc)
)
requires_user_input: Mapped[bool] = mapped_column(default=False)
message_metadata: Mapped[dict | None] = mapped_column(JSONB, default=None)
# Relationships
instance: Mapped["AgentInstance"] = relationship(
"AgentInstance",
back_populates="messages",
foreign_keys=[agent_instance_id],
)